"""
Gradio Chat UI for PydanticAI Multi-Agent System

A simple web interface for interacting with the agents.
Run separately from the main app.

Usage:
    cd demo
    pip install gradio
    python app.py

Or from project root:
    python demo/app.py
"""

import asyncio
import sys
from pathlib import Path

# Load environment variables FIRST, before any imports that need API keys
from dotenv import load_dotenv

env_path = Path(__file__).parent.parent / ".env"
if env_path.exists():
    load_dotenv(env_path)
    print(f"Loaded environment from {env_path}")
else:
    print(f"Warning: No .env file found at {env_path}")
    print("Make sure OPENAI_API_KEY and ANTHROPIC_API_KEY are set.")

# Add parent directory to path so we can import from pydanticai_multiagent
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

import gradio as gr

from pydanticai_multiagent.agents import (
    analyst_agent,
    code_agent,
    research_agent,
    router_agent,
    support_agent,
    writer_agent,
)
from pydanticai_multiagent.dependencies import (
    create_mock_auth_deps,
    create_mock_base_deps,
    create_mock_search_deps,
)


# Agent configurations
AGENTS = {
    "Router (Auto-delegates)": {
        "agent": router_agent,
        "deps_factory": create_mock_search_deps,
        "description": "Analyzes your request and delegates to the best specialist",
    },
    "Research": {
        "agent": research_agent,
        "deps_factory": create_mock_search_deps,
        "description": "Information gathering, fact-finding, and synthesis",
    },
    "Analyst": {
        "agent": analyst_agent,
        "deps_factory": create_mock_base_deps,
        "description": "Data analysis, statistics, and visualizations",
    },
    "Code": {
        "agent": code_agent,
        "deps_factory": create_mock_base_deps,
        "description": "Programming help, debugging, and code review",
    },
    "Writer": {
        "agent": writer_agent,
        "deps_factory": create_mock_base_deps,
        "description": "Content creation - blogs, emails, documentation",
    },
    "Support": {
        "agent": support_agent,
        "deps_factory": create_mock_auth_deps,
        "description": "Customer service and account inquiries",
    },
}


def format_response(agent_name: str, result) -> str:
    """Format agent response for display."""
    output = result.output

    # Handle different output types
    if hasattr(output, "format_with_sources"):
        # ResearchResult
        return output.format_with_sources()
    elif hasattr(output, "explanation"):
        # CodeResponse
        response = f"**Explanation:**\n{output.explanation}"
        if output.code:
            lang = output.language or ""
            response += f"\n\n**Code:**\n```{lang}\n{output.code}\n```"
        if output.suggestions:
            response += "\n\n**Suggestions:**\n" + "\n".join(
                f"- {s}" for s in output.suggestions
            )
        return response
    elif hasattr(output, "answer") and hasattr(output, "category"):
        # SupportResponse
        response = output.answer
        if output.escalate:
            response += f"\n\n*Escalated: {output.escalation_reason}*"
        if output.follow_up_actions:
            response += "\n\n**Follow-up actions:**\n" + "\n".join(
                f"- {a}" for a in output.follow_up_actions
            )
        return response
    elif hasattr(output, "summary"):
        # AnalysisResult
        response = f"**Summary:** {output.summary}\n\n"
        response += f"**Confidence:** {output.confidence:.0%}\n"
        if output.recommendations:
            response += "\n**Recommendations:**\n" + "\n".join(
                f"- {r}" for r in output.recommendations
            )
        if output.methodology:
            response += f"\n\n*Methodology: {output.methodology}*"
        return response
    else:
        # String or other
        return str(output)


async def chat_async(message: str, agent_name: str) -> str:
    """Process a chat message asynchronously."""
    if not message.strip():
        return "Please enter a message."

    agent_config = AGENTS[agent_name]
    agent = agent_config["agent"]
    deps = agent_config["deps_factory"](user_id="demo-user")

    try:
        result = await agent.run(message, deps=deps)
        return format_response(agent_name, result)
    except Exception as e:
        return f"**Error:** {str(e)}\n\nMake sure you have valid API keys in your `.env` file."


def chat(message: str, agent_name: str) -> str:
    """Sync wrapper for async chat."""
    return asyncio.run(chat_async(message, agent_name))


def create_demo():
    """Create the Gradio interface."""

    with gr.Blocks() as demo:
        gr.Markdown(
            """
            # PydanticAI Multi-Agent Demo

            Chat with specialized AI agents. Select an agent below or use the Router
            to automatically delegate to the best specialist.
            """
        )

        with gr.Row():
            agent_dropdown = gr.Dropdown(
                choices=list(AGENTS.keys()),
                value="Router (Auto-delegates)",
                label="Select Agent",
                interactive=True,
            )
            agent_description = gr.Markdown(
                value=f"*{AGENTS['Router (Auto-delegates)']['description']}*"
            )

        def update_description(agent_name):
            return f"*{AGENTS[agent_name]['description']}*"

        agent_dropdown.change(
            update_description, inputs=[agent_dropdown], outputs=[agent_description]
        )

        chatbot = gr.Chatbot(
            label="Conversation",
            height=500,
        )

        with gr.Row():
            msg = gr.Textbox(
                label="Your message",
                placeholder="Type your message here...",
                lines=2,
                scale=4,
            )
            submit_btn = gr.Button("Send", variant="primary", scale=1)

        clear_btn = gr.Button("Clear Conversation")

        # Example prompts
        gr.Markdown("### Example prompts")
        with gr.Row():
            gr.Examples(
                examples=[
                    ["What are the key features of Python 3.12?"],
                    ["Write a haiku about programming"],
                    ["Explain how async/await works in Python"],
                    ["Help me write a professional email declining a meeting"],
                    ["What's the difference between SQL and NoSQL databases?"],
                ],
                inputs=msg,
                label="Click an example to try it",
            )

        def respond(message, chat_history, agent_name):
            if not message.strip():
                return "", chat_history

            # Add user message to history (Gradio 6.x format)
            chat_history = chat_history + [{"role": "user", "content": message}]

            # Get bot response
            bot_message = chat(message, agent_name)
            chat_history = chat_history + [{"role": "assistant", "content": bot_message}]

            return "", chat_history

        # Wire up the interface
        msg.submit(
            respond,
            inputs=[msg, chatbot, agent_dropdown],
            outputs=[msg, chatbot],
        )
        submit_btn.click(
            respond,
            inputs=[msg, chatbot, agent_dropdown],
            outputs=[msg, chatbot],
        )
        clear_btn.click(lambda: [], None, chatbot, queue=False)

        gr.Markdown(
            """
            ---
            **Note:** This demo uses the agents directly. Make sure you have valid
            `OPENAI_API_KEY` and `ANTHROPIC_API_KEY` in your `.env` file.

            Some tools (web search, database, etc.) use mock data in this demo.
            """
        )

    return demo


if __name__ == "__main__":
    demo = create_demo()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,  # Set to True to get a public URL
    )
