import gradio as gr
from llm_agent import AgriculturalAgent
import os

agent = None


def initialize_agent(api_key: str) -> str:
    """
    Initialize the agent with the provided API key.
    
    Args:
        api_key: Groq API key
        
    Returns:
        Status message
    """
    global agent
    if not api_key or not api_key.strip():
        return "Please enter a valid API key."
    
    try:
        agent = AgriculturalAgent(api_key=api_key.strip())
        return "Agent initialized successfully! Dataset loaded and ready. You can now start asking questions."
    except Exception as e:
        return f"Error initializing agent: {str(e)}"


def chat_with_agent(message: str, history: list, api_key: str) -> tuple:
    """
    Handle chat messages with the agent.
    
    Args:
        message: User's message
        history: Chat history (Gradio format: list of tuples [(user_msg, assistant_msg), ...])
        api_key: Groq API key
        
    Returns:
        Tuple of (empty string, updated history in Gradio format)
    """
    global agent
    
    print(f"\n[APP] Received message: {message[:100]}...") 
    
    if agent is None:
        if not api_key or not api_key.strip():
            history.append((message, "Please initialize the agent with an API key first."))
            return "", history
        print("[APP] Initializing agent...")
        init_msg = initialize_agent(api_key)
        if "Error" in init_msg:
            history.append((message, init_msg))
            return "", history
    
    if agent is None:
        history.append((message, "Please initialize the agent with an API key first."))
        return "", history
    
    print("[APP] Sending message to agent...")
    
    response = agent.chat(message)
    print(f"[APP] Received response from agent: {response[:100]}...")
    
    history.append((message, response))
    
    return "", history


def reset_chat():
    """Reset the chat conversation."""
    global agent
    if agent:
        agent.reset_conversation()
    return []  


with gr.Blocks(title="Agricultural Research AI Agent") as demo:
    gr.Markdown("""
    # 🌾 Agricultural Research AI Agent
    
    This AI agent helps you explore and find information from a comprehensive collection of 
    **45,232 agricultural research publications** from CGIAR (Consultative Group on International Agricultural Research).
    
    ## Features
    - 🔍 Search for research documents on specific agricultural topics
    - 📚 Browse documents by topic (crop management, pest control, climate adaptation, etc.)
    - 📄 Get detailed information about specific research papers
    - 💡 Get insights and answers based on the latest agricultural research
    
    ## How to Use
    1. Enter your Groq API key in the field below (get a free key at https://console.groq.com)
    2. Click "Initialize Agent" to start
    3. Ask questions about agricultural topics, search for documents, or browse research papers
    
    **Example questions:**
    - "What research is available on rice cultivation?"
    - "Find documents about pest control in agriculture"
    - "Tell me about climate adaptation strategies for small-scale farmers"
    - "What does the dataset contain?"
    
    ---
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            api_key_input = gr.Textbox(
                label="Groq API Key",
                placeholder="Enter your Groq API key here...",
                type="password",
                info="Your API key is not stored and is only used for this session. Get a free API key at https://console.groq.com"
            )
            init_btn = gr.Button("Initialize Agent", variant="primary")
            init_status = gr.Textbox(label="Status", interactive=False)
        
        with gr.Column(scale=1):
            reset_btn = gr.Button("Reset Chat", variant="secondary")
    
    chatbot = gr.Chatbot(
        label="Chat with Agricultural Research Agent",
        height=500
    )
    
    with gr.Row():
        msg = gr.Textbox(
            label="Your Question",
            placeholder="Ask about agricultural research, search for documents, or browse topics...",
            scale=4,
            lines=2
        )
        submit_btn = gr.Button("Send", variant="primary", scale=1)
    
    init_btn.click(
        fn=initialize_agent,
        inputs=[api_key_input],
        outputs=[init_status]
    )
    
    submit_btn.click(
        fn=chat_with_agent,
        inputs=[msg, chatbot, api_key_input],
        outputs=[msg, chatbot]
    )
    
    msg.submit(
        fn=chat_with_agent,
        inputs=[msg, chatbot, api_key_input],
        outputs=[msg, chatbot]
    )
    
    reset_btn.click(
        fn=reset_chat,
        outputs=[chatbot]
    )
    
    gr.Markdown("""
    ---
    **Dataset Information:**
    - Source: [CGIAR/gardian-ai-ready-docs](https://huggingface.co/datasets/CGIAR/gardian-ai-ready-docs) on HuggingFace
    - Contains 45,232 structured agricultural research publications
    - Topics: Crop management, pest control, climate adaptation, farming systems, and more
    
    **Note:** This application uses Groq's free API for the LLM agent. Get your free API key at [console.groq.com](https://console.groq.com)
    """)

if __name__ == "__main__":
    demo.launch(share=False, server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())