from langchain_community.llms import LlamaCpp
from langchain_core.prompts import ChatPromptTemplate
from crewai import Agent, Task, Crew
import os

# 1. Charger le modèle avec ROCm (GPU) + CPU
llm = LlamaCpp(
    model_path="/models/model.gguf",
    n_ctx=4096,  # Contexte étendu pour le code
    n_gpu_layers=32,  # Utilise le GPU pour 32 couches (ajuste selon ta RAM)
    n_threads=8,  # Utilise 8 threads CPU pour le reste
    temperature=0.1,
    verbose=True,
)

# 2. Définir l'agent GSD
code_agent = Agent(
    role="Senior Full-Stack Developer",
    goal="Développer des fonctionnalités complètes (backend + frontend + DB) en Rust, TypeScript, Kotlin, SQL",
    backstory="Tu es un expert en Rust (Actix), TypeScript (Angular), Kotlin (Spring Boot), et SQL (PostgreSQL). "
              "Tu utilises la méthodologie GSD : tu décomposes les tâches, génères du code propre, et valides chaque étape.",
    tools=[],  # On ajoutera des outils plus tard
    llm=llm,
    verbose=True,
)

# 3. Prompt personnalisé pour le mode GSD
gsd_prompt = """
Tu es en mode **Get Shit Done (GSD)**. Voici comment procéder :
1. **Analyse** : Décompose la demande en sous-tâches claires (ex: backend, frontend, DB).
2. **Génération** : Pour chaque sous-tâche, génère le code **complet et fonctionnel** dans le langage demandé.
3. **Validation** : Vérifie que le code est compilable/lintable. Si non, corrige-le.
4. **Retour** : Fournis le code final + des instructions pour l'intégrer.

**Exemple de demande** :
"J'aimerai rendre configurable la durée de validité des tokens JWT dans mon backend Rust + frontend Angular."

**Ta réponse doit inclure** :
- [ ] Code Rust (Actix) pour modifier la configuration JWT.
- [ ] Schéma SQL pour stocker la configuration en DB.
- [ ] Endpoint HTTP pour modifier la configuration.
- [ ] Code TypeScript (Angular) pour le formulaire.
- [ ] Instructions pour tester chaque partie.

**Contraintes** :
- Utilise **Actix-Web** pour Rust.
- Utilise **Angular Material** pour le frontend.
- Pour Kotlin, utilise **Spring Boot 3 + JPA**.
- Pour SQL, utilise **PostgreSQL**.

**Demande actuelle** : {user_request}
"""

# 4. Exécuter une tâche GSD
def run_gsd_task(user_request: str):
    task = Task(
        description=gsd_prompt.format(user_request=user_request),
        agent=code_agent,
        expected_output="Code complet et fonctionnel pour chaque sous-tâche, avec instructions d'intégration.",
    )
    crew = Crew(agents=[code_agent], tasks=[task], verbose=2)
    result = crew.kickoff()
    return result

if __name__ == "__main__":
    print("Agent GSD démarré. Tape 'quit' pour arrêter.")
    while True:
        user_request = input("\n> ")
        if user_request.lower() == "quit":
            break
        result = run_gsd_task(user_request)
        print("\n--- Résultat ---\n")
        print(result)