# base/base.py
import os
import time
from langgraph.prebuilt import create_react_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, trim_messages
from langchain_core.callbacks import BaseCallbackHandler
from langchain_openai import ChatOpenAI
from langchain_oci import ChatOCIGenAI
#from langfuse import get_client
#from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
from functools import lru_cache
import queue, threading

from dotenv import load_dotenv
import os
load_dotenv()

class StreamCaptureHandler(BaseCallbackHandler): 
    def __init__(self): 
        self.tokens = [] 
        self.buffer = "" 
        self.final_result = None 
        self.steps = [] 

    def on_llm_new_token(self, token, **kwargs): 
        self.tokens.append(token) 

    def on_chain_end(self, outputs, **kwargs): 
        self.final_result = outputs 

    def on_tool_end(self, output, **kwargs): 
        self.steps.append(output)

@lru_cache(maxsize=32)
def load_langfuse_prompt(
    prompt_name: str,
    ttl_seconds: int = 300,
) -> str:
    """
    Carrega prompt do Langfuse com cache por processo
    e revalidação automática a cada 5 minutos.
    """
    prompt = langfuse.get_prompt(
        name=prompt_name,
        cache_ttl_seconds=ttl_seconds,
    )
    return prompt.prompt

#langfuse = get_client()
#lf_handler = LangfuseCallbackHandler()

class BaseAgent:
    def __init__(self, 
                 tools: list, 
                 streaming: bool = False,
                 prompt_vars: dict | None = None,
                 agent_name: str = None,
                 dynamic_prompt: bool = False):

        self.streaming = streaming
        self.prompt_vars = prompt_vars or {}
        self.agent_name = agent_name
        self.dynamic_prompt = dynamic_prompt

        self.llm = ChatOCIGenAI(
            model_id=os.getenv("OCI_ENDPOINT_ID", ""),
            service_endpoint=os.getenv("OCI_ENDPOINT", ""),
            compartment_id=os.getenv("OCI_COMPARTMENT_ID", ""),
            auth_file_location=os.getenv("OCI_AUTH_FILE_LOCATION", "./config"),
            model_kwargs={"temperature": 0.0,  
                        "top_p": 0.1,
                        "reasoning_effort":"MINIMAL"}
        )
        """self.llm = ChatOpenAI(
            model="openai/gpt-oss-20b",
            api_key="fake-key",
            base_url="http://10.153.34.154/gpt-oss-20b/v1",
            temperature=0.0,
            top_p=0.1,
            reasoning_effort="low",
            streaming=self.streaming,
        )"""

        self.stream_handler = StreamCaptureHandler()

        raw_prompt = self._load_prompt()
        #raw_prompt = load_langfuse_prompt(self.agent_name)

        # Aplica variáveis parciais no template do prompt
        if self.prompt_vars:
            for key, value in self.prompt_vars.items():
                raw_prompt = raw_prompt.replace("{" + key + "}", str(value))

        self.system_prompt = raw_prompt

        # Histórico de mensagens gerenciado manualmente (substitui ConversationBufferWindowMemory)
        self.messages: list = []
        self.k = 30  # window size — últimas k*2 mensagens (k pares human/ai)

        # Cria o agente usando langgraph-prebuilt create_react_agent
        # Se dynamic_prompt=True, usa lambda que lê self.system_prompt a cada invocação
        # permitindo que o prompt seja atualizado entre chamadas (ex: UnifiedAgent)
        if self.dynamic_prompt:
            self.agent_exec = create_react_agent(
                model=self.llm,
                tools=tools,
                prompt=lambda state: [SystemMessage(content=self.system_prompt)] + state["messages"],
            )
        else:
            self.agent_exec = create_react_agent(
                model=self.llm,
                tools=tools,
                prompt=self.system_prompt,
            )

    def _load_prompt(self) -> str:
        raise NotImplementedError

    def _get_trimmed_messages(self) -> list:
        """Retorna as últimas k*2 mensagens (janela deslizante)."""
        max_messages = self.k * 2
        if len(self.messages) > max_messages:
            return self.messages[-max_messages:]
        return list(self.messages)

    def inject_user_message(self, text: str):
        self.messages.append(HumanMessage(content=text))

    def inject_ai_message(self, text: str):
        self.messages.append(AIMessage(content=text))

    # ================= RUN NORMAL ======================
    def run(self, user_input: str):
        # Adiciona a mensagem do usuário ao histórico
        self.messages.append(HumanMessage(content=user_input))

        # Prepara mensagens com window trimming
        input_messages = self._get_trimmed_messages()

        result = self.agent_exec.invoke(
            {"messages": input_messages},
            #config={"callbacks": [self.stream_handler, lf_handler]}
        )

        #print("Texto:",result["messages"][-1].content)
        #print("Metadata:",result["messages"][-1].usage_metadata)
        #print("*"*100)

        # Extrai a última mensagem do AI do resultado
        output_messages = result.get("messages", [])
        ai_response = ""
        tool_steps = []

        for msg in output_messages:
            if isinstance(msg, AIMessage) and msg.content and not getattr(msg, 'tool_calls', None):
                ai_response = msg.content

        # Coleta intermediate tool steps
        for msg in output_messages:
            if hasattr(msg, 'tool_calls') and msg.tool_calls:
                tool_steps.extend(msg.tool_calls)

        # Adiciona a resposta do AI ao histórico
        if ai_response:
            ai_response = ai_response.split('###commentary')[0].strip() # remove commentary
            self.messages.append(AIMessage(content=ai_response))

        #print(self.messages)

        return {
            "output": ai_response,
            "tools": tool_steps
        }

    # ================= STREAMING ======================
    def stream_run(self, user_input: str): 
        self.stream_handler.tokens = []

        # Adiciona a mensagem do usuário ao histórico
        self.messages.append(HumanMessage(content=user_input))

        # Prepara mensagens com window trimming
        input_messages = self._get_trimmed_messages()

        # Usa streaming do agente
        ai_response = ""
        for chunk in self.agent_exec.stream(
            {"messages": input_messages},
            config={"callbacks": [self.stream_handler]}#, lf_handler
        ):
            # Processa chunks de streaming
            for node_name, node_output in chunk.items():
                if node_name == "agent" and "messages" in node_output:
                    for msg in node_output["messages"]:
                        if isinstance(msg, AIMessage) and msg.content:
                            ai_response = msg.content
                            for token in msg.content:
                                yield token

        # Adiciona a resposta ao histórico
        if ai_response:
            self.messages.append(AIMessage(content=ai_response))

        yield ""
