import os
import json
from smolagents import MultiStepAgent
from smolagents.agents import ActionOutput
from models import ModelManager
from tools import search_web, scrape_website, read_file, PythonInterpreterTool


class MonAgent(MultiStepAgent):
    """
    Un agent multi-étapes avec l'implémentation COMPLÈTE de TOUTES les méthodes
    requises et optionnelles pour garantir un fonctionnement sans erreur
    d'implémentation manquante.
    """

    # --- MÉTHODES ABSTRAITES OBLIGATOIRES ---

    def initialize_system_prompt(self) -> str:
        """
        [PIÈCE 1/4 - OBLIGATOIRE] Définit la "personnalité" et les instructions
        de base de l'agent. C'est la première chose que l'on doit implémenter.
        """
        # Ce prompt est optimisé pour forcer une sortie JSON claire.
        return """You are a world-class autonomous agent. Your goal is to fully answer the user's question by creating a plan and using tools.

You MUST format your response as a JSON object containing a "plan" key, which is a list of tool calls.
Each tool call is a dictionary with "tool" and "args".

Example of a valid response with a tool call:
{
  "plan": [
    {
      "tool": "search_web",
      "args": {
        "query": "Who is the current president of France?"
      }
    }
  ]
}

If you have the final answer, respond with an empty plan and the answer in the 'final_answer' key:
{
  "plan": [],
  "final_answer": "The final answer is..."
}
"""

    def _step_stream(self, memory_step):
        """
        [PIÈCE 2/4 - OBLIGATOIRE] C'est le cœur de l'agent. Cette méthode est
        appelée à chaque étape pour décider de la prochaine action.
        Elle doit générer un appel d'outil ou une réponse finale.
        """
        # 1. Préparer les messages pour le modèle
        memory_messages = self.write_memory_to_messages()
        memory_step.model_input_messages = memory_messages

        # 2. Appeler le modèle pour obtenir sa décision
        try:
            chat_message = self.model.generate(
                memory_messages,
                tools_to_call_from=list(self.tools.values()),
            )
            memory_step.model_output_message = chat_message
            memory_step.token_usage = chat_message.token_usage
        except Exception as e:
            raise Exception(f"Erreur lors de la génération du modèle : {e}")

        # 3. Analyser la réponse du modèle pour trouver des appels d'outils
        tool_calls = chat_message.tool_calls
        if not tool_calls:
            # Si le modèle ne retourne pas d'appel d'outil structuré, on essaie de le parser depuis le texte brut
            try:
                chat_message = self.model.parse_tool_calls(chat_message)
                tool_calls = chat_message.tool_calls
            except Exception:
                 # Si le parsing échoue, on considère que c'est une réponse finale
                yield ActionOutput(output=chat_message.content, is_final_answer=True)
                return

        # 4. Traiter les appels d'outils
        final_answer = None
        is_final = False
        for tool_call in tool_calls:
            # On stocke l'appel d'outil dans la mémoire
            memory_step.tool_calls = tool_calls
            yield tool_call

            # On exécute l'outil
            tool_output_value = self.execute_tool_call(tool_call.name, tool_call.arguments)
            
            # On vérifie si c'est la réponse finale
            if tool_call.name == "final_answer":
                final_answer = tool_output_value
                is_final = True

            # On génère une observation textuelle pour la prochaine étape de l'IA
            observation = self.render_tool_result(tool_output_value)
            memory_step.observations = observation
            
            yield {"tool_call_id": tool_call.id, "output": observation}

        # 5. Renvoyer le résultat final de l'étape
        yield ActionOutput(output=final_answer, is_final_answer=is_final)


    # --- MÉTHODES ADDITIONNELLES POUR LA ROBUSTESSE ---
    # Bien que non explicitement listées comme abstraites, elles sont
    # appelées par la logique interne et doivent être présentes.

    def parse_plan(self, response: str) -> list[dict]:
        """
        [PIÈCE 3/4 - REQUISE POUR LE PLANNING] Transforme le texte brut du modèle
        en une liste d'actions structurées.
        """
        cleaned_response = response.strip().removeprefix("```json").removesuffix("```").strip()
        try:
            parsed_json = json.loads(cleaned_response)
            return parsed_json.get("plan", [])
        except json.JSONDecodeError:
            print(f"⚠️ Erreur de parsing JSON dans `parse_plan`. Réponse reçue:\n{response}")
            return []

    def render_tool_result(self, tool_output: any) -> str:
        """
        [PIÈCE 4/4 - REQUISE POUR LA BOUCLE] Transforme le résultat d'un outil
        (qui peut être n'importe quel objet Python) en un texte simple que l'IA
        peut comprendre pour la suite. C'était probablement la pièce manquante principale.
        """
        print(f"⚙️  Formatage du résultat de l'outil: {str(tool_output)[:300]}...")
        if isinstance(tool_output, str):
            return tool_output
        if isinstance(tool_output, (list, dict)):
            try:
                return json.dumps(tool_output, indent=2, ensure_ascii=False)
            except TypeError:
                return str(tool_output)
        return str(tool_output)

    def render_final_answer(self, final_context: dict, final_response: str) -> str:
        """
        [BONUS] Est appelée à la toute fin pour formater la réponse finale.
        """
        # Essaye de parser une réponse finale structurée en JSON
        cleaned_response = final_response.strip().removeprefix("```json").removesuffix("```").strip()
        try:
            parsed_json = json.loads(cleaned_response)
            return parsed_json.get("final_answer", final_response)
        except json.JSONDecodeError:
            return final_response


# --- Le reste du fichier reste identique, il utilise maintenant notre classe MonAgent "blindée" ---

class BasicAgent:
    """
    Classe de compatibilité qui utilise notre nouvel agent complet et robuste.
    """
    def __init__(self):
        print("🚀 Initialisation du BasicAgent (version blindée)...")
        try:
            if not os.getenv("HF_TOKEN"):
                print("⚠️  Attention: Le token Hugging Face (HF_TOKEN) n'est pas défini.")

            self.agent = MonAgent(
                model=ModelManager().get_orchestrator(),
                tools=[search_web, scrape_website, read_file, PythonInterpreterTool()]
            )
            print("✅ BasicAgent initialisé avec succès!")

        except Exception as e:
            print(f"❌ Erreur critique lors de l'initialisation: {e}")
            self.agent = None

    def __call__(self, question: str) -> str:
        if self.agent is None:
            return "Erreur: L'agent n'a pas pu être initialisé. Vérifiez les logs et la configuration (HF_TOKEN)."
        
        print(f"\n{'='*40}\n🤖 NOUVELLE QUESTION: {question}\n{'='*40}")

        try:
            # agent.run() va maintenant utiliser notre implémentation complète de MonAgent
            return self.agent.run(question)
        except Exception as e:
            import traceback
            print(f"❌ Erreur irrécupérable lors du traitement par MonAgent: {e}\n{traceback.format_exc()}")
            return f"Une erreur irrécupérable s'est produite: {e}"