import os
import torch
from huggingface_hub import InferenceClient

# Importera LangChain-komponenter
from langchain_community.llms import HuggingFaceHub # För att använda HF Inference API som LLM
from langchain.agents import AgentExecutor, create_react_agent # Agentens exekverare och ReAct-agent konstruktorn
from langchain.tools import Tool # Verktygsklassen i LangChain
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage

# Importera dina befintliga, anpassade verktygsfunktioner
from tavily_search import search_tavily
from asr_tool import transcribe_audio
from excel_tool import analyze_excel
from math_tool import calculate_math

class GaiaAgent:
    def __init__(self, model_id: str = "google/gemma-2b-it"):
        """
        Initialiserar GaiaAgent, nu med LangChain.
        """
        print(f"Initialiserar GaiaAgent med modell: {model_id}")

        hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
        if not hf_token:
            raise ValueError(
                "Hugging Face token (HF_TOKEN eller HUGGING_FACE_HUB_TOKEN) är inte konfigurerad i miljövariabler."
                "Vänligen lägg till din token som en 'Repository secret' i dina Space-inställningar."
            )

        # 1. Initialisera LLM med LangChain's HuggingFaceHub
        try:
            # HuggingFaceHub ansluter till en fjärrmodell via HF Inference API
            self.llm = HuggingFaceHub(
                repo_id=model_id,
                huggingfacehub_api_token=hf_token,
                task="text-generation", # Specifiera task
                # model_kwargs={"temperature": 0.1, "max_new_tokens": 512} # Exempel på modell-kwargs
            )
            print("LangChain HuggingFaceHub LLM laddad framgångsrikt.")
        except Exception as e:
            raise RuntimeError(
                f"Misslyckades att initialisera HuggingFaceHub LLM: {e}."
                "Vänligen kontrollera din HF token och att modellen är tillgänglig/laddningsbar."
            )

        # 2. Definiera dina anpassade verktyg som LangChain Tool-objekt
        tools_list = [
            Tool.from_function(
                func=search_tavily,
                name="search_tavily",
                description="Användbart för att söka information online med Tavily Search. Returnerar en sammanfattning av de mest relevanta resultaten från webben. Kräver en fråga som input.",
            ),
            Tool.from_function(
                func=transcribe_audio,
                name="transcribe_audio",
                description="Transkriberar ljudfil till text. Användbart för att omvandla tal till text från en angiven ljudfilsväg. Kräver en filsökväg till ljudfilen som input.",
            ),
            Tool.from_function(
                func=analyze_excel,
                name="analyze_excel",
                description="Analysera Excel-filer och returnera detaljerad information om rader, kolumner, datatyper och statistik (summa, medelvärde, max, min för numeriska kolumner). Kan ta både en lokal filväg eller en URL till Excel-filen som input.",
            ),
            Tool.from_function(
                func=calculate_math,
                name="calculate_math",
                description="Beräkna matematiska uttryck. Användbart för att utföra aritmetiska operationer som addition, subtraktion, multiplikation, division och potenser. Tar ett matematiskt uttryck som en sträng som input.",
            )
        ]
        print(f"Laddade {len(tools_list)} anpassade verktyg för LangChain.")

        # 3. Skapa en prompt för ReAct-agenten
        # Detta prompt-format är viktigt för hur LLM:en förstår att använda verktyg.
        # MessagesPlaceholder används för att injicera verktyg och meddelandehistorik dynamiskt.
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "Du är en hjälpsam AI-assistent. Använd tillgängliga verktyg för att svara på frågor."),
                MessagesPlaceholder("chat_history", optional=True),
                ("human", "{input}"),
                MessagesPlaceholder("agent_scratchpad"), # Detta är där agentens tankar och verktygskall kommer att finnas
            ]
        )

        # 4. Initialisera LangChain ReAct-agenten
        # create_react_agent är en konstruktorfunktion för en ReAct-baserad agent
        agent = create_react_agent(self.llm, tools_list, prompt)

        # 5. Skapa AgentExecutor för att köra agenten
        # AgentExecutor är den körbara delen som hanterar agentens "tankeloop" och verktygskall
        self.agent_executor = AgentExecutor(
            agent=agent,
            tools=tools_list,
            verbose=True, # Sätt till True för att se agentens tankeprocess i loggarna
            handle_parsing_errors=True # Hantera parsningsfel graciöst
        )
        print("LangChain AgentExecutor initialiserad.")

    def process_task(self, task_prompt: str) -> str:
        """
        Bearbetar en uppgift med den interna LangChain AgentExecutor.
        """
        print(f"\nBearbetar uppgift med LangChain AgentExecutor: '{task_prompt}'")
        try:
            # Anropa agenten med invoke. Den returnerar ett dictionary.
            # "input" är användarens prompt.
            # "chat_history" kan skickas in om du har kontext från tidigare konversationer.
            result = self.agent_executor.invoke({"input": task_prompt})
            
            # Det slutgiltiga svaret finns vanligtvis under nyckeln "output"
            final_answer = result.get("output", "Agenten kunde inte generera ett slutgiltigt svar.")
            
            print(f"\nLangChain AgentExecutor avslutad. Slutgiltigt svar: {final_answer}")
            return final_answer
        except Exception as e:
            error_message = f"Ett fel uppstod under agentens bearbetning: {e}"
            print(error_message)
            return f"Agenten kunde inte slutföra uppgiften på grund av ett fel: {error_message}"