import os
import chromadb
import gradio as gr
from dotenv import load_dotenv
from llama_index.core import VectorStoreIndex
from llama_index.core.agent import ReActAgent
from llama_index.core.tools import QueryEngineTool
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.gemini import Gemini
from llama_index.core.workflow import Context
from llama_index.vector_stores.chroma import ChromaVectorStore

load_dotenv()

agent = None
conversation_context = None

async def initialize_agent():
    """Initialize the agent once"""
    global agent, conversation_context
    if agent is not None:
        return agent, conversation_context

    llm = Gemini(
        model="models/gemini-flash-latest",
        api_key=os.getenv("GEMINI_API"),
        temperature=0.3,
    )

    embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

    db = chromadb.PersistentClient(path="./product_db")
    chroma_collection = db.get_collection(name="product_catalog")
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

    index = VectorStoreIndex.from_vector_store(
        vector_store=vector_store,
        embed_model=embed_model
    )

    query_engine = index.as_query_engine(llm=llm)

    query_tool = QueryEngineTool.from_defaults(
        query_engine=query_engine,
        name="ProductInfoTool",
        description="A tool to retrieve information about camping products, including their stock availability.",
    )

    agent = ReActAgent(
        llm=llm,
        tools=[query_tool],
        verbose=False,
        system_prompt="""You are a friendly and knowledgeable camping gear expert.
Your goal is to find the perfect product for the user and tell them about it in a helpful, conversational way.
Use the `ProductInfoTool` to find the best match for the user's query.
In your final response to the user, you MUST include the following three pieces of information:
1. The full product name.
2. A brief, one-sentence reason why it's a good choice for them.
3. The exact stock status (e.g., '15 available' or 'out of stock').
If the tool cannot find a suitable product, just say: 'I'm sorry, I couldn't find a product that matches your request.'
Remember conversation context and refer back to previous messages when appropriate."""
    )

    conversation_context = Context(agent)

    return agent, conversation_context


async def chat_with_agent(message, history):
    """Handle chat messages with the agent"""
    global conversation_context, agent

    try:
        agent, ctx = await initialize_agent()
        response = await agent.run(message, ctx=conversation_context)
        return str(response.response)
    except Exception as e:
        if "index out of range" in str(e):
            conversation_context = Context(agent)
            response = await agent.run(message, ctx=conversation_context)
            return str(response.response)
        return f"Sorry, I encountered an error: {str(e)}"


def main():
    """Launch the simple Gradio interface"""
    demo = gr.ChatInterface(
        fn=chat_with_agent,
        title="🏕️ Tory - The Camping Gear Assistant",
        description="Ask me about camping products and I'll help you find the perfect gear!",
        examples=[
            "I need a lightweight tent for 2 people",
            "What sleeping bags do you have?",
            "Show me available camping stoves"
        ],
    )
    demo.launch()

if __name__ == "__main__":
    main()