import os
import gradio as gr
from langchain.agents import create_agent
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langgraph.checkpoint.memory import InMemorySaver
from dotenv import load_dotenv
from agents import BookRetriverAgent
from tools import get_retrieve_book_context_rag_agent_tool


class GradioAgent:

    def __init__(self):
        self.agent = self.__create_agent()

    def inicialize(self):
        chatbot = gr.ChatInterface(
            self._respond,
            type="messages"
        )

        with gr.Blocks(fill_height=True) as demo:
            chatbot.render()

        demo.launch()

    def __create_agent(self):
        hf_model = HuggingFaceEndpoint(
            repo_id="Qwen/Qwen3-30B-A3B-Instruct-2507",
            task="text-generation",
            provider="auto",
            huggingfacehub_api_token=os.getenv("HF_TOKEN"),
            do_sample=False
        )
        llm = ChatHuggingFace(llm=hf_model)

        rag_agent = BookRetriverAgent().agent

        tools = [
            get_retrieve_book_context_rag_agent_tool(rag_agent)
        ]

        return create_agent(
            tools=tools,
            model=llm,
            checkpointer=InMemorySaver(),
            system_prompt="You are a helpful and usefull assistant."
        )

    def _respond(
        self,
        message,
        history
    ):
        result = self.agent.invoke(
            {"messages": [{"role": "user", "content": message}]},
            {"configurable": {"thread_id": "1"}},
        )
        output = result['messages'][-1].content

        yield output


if __name__ == "__main__":
    load_dotenv()
    gradio = GradioAgent()
    gradio.inicialize()