import gradio as gr
from agno.agent import Agent
from agno.embedder.google import GeminiEmbedder
from agno.knowledge.url import UrlKnowledge
from agno.knowledge.text import TextKnowledgeBase
from agno.knowledge.markdown import MarkdownKnowledgeBase
from agno.models.google import Gemini
from agno.storage.sqlite import SqliteStorage
from agno.vectordb.lancedb import LanceDb, SearchType

# Load Agno documentation in a knowledge base
# You can also use `https://docs.agno.com/llms-full.txt` for the full documentation
# knowledge = UrlKnowledge(
#     urls=["https://www.mea.or.th/about-mea/background/history"],
#     vector_db=LanceDb(
#         uri="tmp/lancedb",
#         table_name="mea-hist",
#         search_type=SearchType.hybrid,
#         # embedder=OllamaEmbedder(id="bge-m3"),
#         embedder=GeminiEmbedder(),
#     ),
# )

# Alternative knowledge base using company policies
knowledge = MarkdownKnowledgeBase(
    path="./meahist.md",
    vector_db=LanceDb(
        uri="tmp/lancedb",
        table_name="mea-docs",
        search_type=SearchType.hybrid,
        # Use OpenAI for embeddings
        embedder=GeminiEmbedder(dimensions=768),
    ),
)

# Store agent sessions in a SQLite database
storage = SqliteStorage(table_name="agent_sessions", db_file="tmp/agent.db")

# model_id = "gemini-2.5-flash"
model_id = "gemini-2.5-flash-lite"
# model_id = "gemma-3-12b-it"

agent = Agent(
    name="Agno Assist",
    model=Gemini(id=model_id),
    instructions=[
        "Search your knowledge before answering the question.",
    ],
    knowledge=knowledge,
    storage=storage,
    search_knowledge=True,
    # add_datetime_to_instructions=True,
    # Add the chat history to the messages
    # add_history_to_messages=True,
    read_chat_history=True,
    show_tool_calls=False,
    # Number of history runs
    num_history_runs=3,
    markdown=True,
)


def chat(message, history):
    response = agent.run(message, stream=True)
    content = ""
    for event in response:
        if event.event == "RunCompleted":
            break
        if event.event == "RunResponseContent":
            content += event.content
            yield content
    yield content


# Load the knowledge base, comment out after first run
# Set recreate to True to recreate the knowledge base if needed
agent.knowledge.load(recreate=False)

# Create the Gradio interface
demo = gr.ChatInterface(
    chat,
    title="MEA History Chat Bot",
    description="Ask questions about MEA (Metropolitan Electricity Authority) history and get AI-powered answers. Data source: https://www.mea.or.th/about-mea/background/history",
    examples=[
        "ประวัติความเป็นมาของ MEA คืออะไร",
        "MEA ก่อตั้งเมื่อไหร่",
        "เหตุการณ์สำคัญในการพัฒนาของ MEA คืออะไร",
    ],
)

if __name__ == "__main__":
    demo.launch()