import gradio as gr
import asyncio
import json
import time
import io
from typing import List
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from core.agents.SkeletonGraphAgent import SkeletonGraphAgent
from core.config.prompt import get_default_system_prompt
from core.config.metadata import create_metadata

# ----- State -----
conversation_history = []

# ----- Agent Initialization -----
async def initialize_agent_async(metadata):
    try:
        new_agent = await SkeletonGraphAgent.create(metadata)
        return new_agent, None
    except Exception as e:
        return None, str(e)

def initialize_agent_handler(model_name, temperature, max_tokens, system_prompt, api_key):
    metadata = create_metadata(model_name, temperature, max_tokens, system_prompt, api_key)
    
    # Use the existing event loop if available, otherwise create a new one
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
    agent, error = loop.run_until_complete(initialize_agent_async(metadata))
    if agent:
        return agent, "✅ Agent initialized successfully!", "🟢 Agent Ready"
    else:
        return None, f"❌ Failed to initialize agent: {error}", "🔴 Agent Failed"

# ----- Chat Handler -----
async def get_agent_response_async(agent, user_input):
    global conversation_history
    human_msg = HumanMessage(content=user_input)
    conversation_history.append(human_msg)

    input_message = {"messages": conversation_history.copy()}

    try:
        result = await agent.workflow.ainvoke(input=input_message)
        if result and 'messages' in result:
            for msg in reversed(result['messages']):
                if isinstance(msg, AIMessage):
                    conversation_history.append(msg)
                    return msg.content, None
        fallback = AIMessage(content="No response generated")
        conversation_history.append(fallback)
        return fallback.content, None
    except Exception as e:
        conversation_history.pop()
        return None, str(e)

def chat_with_agent(agent, user_input, chat_history):
    if not agent:
        chat_history.append([user_input, "❌ Please initialize your agent first."])
        return chat_history, ""

    if not user_input.strip():
        return chat_history, ""

    # Use the existing event loop if available, otherwise create a new one
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    
    response, error = loop.run_until_complete(get_agent_response_async(agent, user_input))

    if response:
        chat_history.append([user_input, response])
    else:
        chat_history.append([user_input, f"❌ Error: {error}"])

    return chat_history, ""

def reset_conversation():
    global conversation_history
    conversation_history = []
    return [], "🔄 Conversation reset!", "🟢 Agent Ready"

# ----- Gradio Interface -----
def create_interface():
    with gr.Blocks(
        title="TrackMate AI",
        theme=gr.themes.Soft(),
        css="""
        /* Dark theme colors */
        :root {
            --bg-primary: #121212;
            --bg-secondary: #1e1e1e;
            --bg-tertiary: #252525;
            --text-primary: #e0e0e0;
            --text-secondary: #a0a0a0;
            --accent-primary: #3a506b;
            --accent-secondary: #1c2541;
            --accent-tertiary: #0b132b;
            --highlight: #5bc0be;
            --error: #ff6b6b;
            --success: #6bd425;
        }
        
        .documentation-container {
            background-color: var(--bg-secondary);
            color: var(--text-primary);
            padding: 2rem;
            border-radius: 15px;
            margin-bottom: 2rem;
            box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3);
        }
        
        .documentation-header {
            text-align: center;
            margin-bottom: 1.5rem;
        }
        
        .documentation-section {
            background-color: var(--bg-tertiary);
            padding: 1.5rem;
            border-radius: 10px;
            margin-bottom: 1.5rem;
            border-left: 4px solid var(--highlight);
        }
        
        .config-panel {
            background-color: var(--bg-secondary);
            padding: 1.5rem;
            border-radius: 15px;
            box-shadow: 0 4px 20px rgba(0, 0, 0, 0.3);
            height: 100%;
            display: flex;
            flex-direction: column;
        }
        
        .chat-panel {
            background-color: var(--bg-secondary);
            padding: 1.5rem;
            border-radius: 15px;
            box-shadow: 0 4px 20px rgba(0, 0, 0, 0.3);
            height: 100%;
            display: flex;
            flex-direction: column;
        }
        
        .panel-header {
            text-align: center;
            padding: 1rem;
            background-color: var(--accent-primary);
            color: var(--text-primary);
            border-radius: 10px;
            margin-bottom: 1rem;
        }
        
        .tools-table {
            width: 100%;
            border-collapse: collapse;
            margin-top: 1rem;
        }
        
        .tools-table th, .tools-table td {
            padding: 0.75rem;
            text-align: left;
            border-bottom: 1px solid var(--accent-secondary);
        }
        
        .tools-table th {
            font-weight: bold;
            background-color: var(--accent-tertiary);
        }
        
        .feature-list {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
            gap: 0.5rem;
            margin-top: 1rem;
        }
        
        .feature-item {
            background-color: var(--accent-secondary);
            padding: 0.75rem;
            border-radius: 10px;
        }
        
        a {
            color: var(--highlight);
            text-decoration: none;
            font-weight: bold;
        }
        
        a:hover {
            text-decoration: underline;
        }
        
        /* Chat avatar styling */
        .chatbot .avatar {
            display: flex !important;
            align-items: center;
            justify-content: center;
            width: 40px !important;
            height: 40px !important;
            border-radius: 50%;
            background-color: var(--accent-secondary);
            color: var(--text-primary);
            font-size: 20px;
            margin-right: 10px;
        }
        
        /* Ensure chat messages have proper contrast */
        .chatbot .user-message {
            background-color: var(--accent-primary) !important;
            color: var(--text-primary) !important;
            border-radius: 10px !important;
        }
        
        .chatbot .bot-message {
            background-color: var(--accent-secondary) !important;
            color: var(--text-primary) !important;
            border-radius: 10px !important;
        }
        
        /* Override any theme styles that might hide avatars */
        .chatbot .message-container {
            display: flex !important;
            align-items: flex-start !important;
        }
        """
    ) as app:
        agent_state = gr.State(value=None)
        
        # Consolidated Documentation Section
        gr.HTML("""
        <div class="documentation-container">
            <div class="documentation-header">
                <h1 style="margin: 0; font-size: 3rem; font-weight: 700;">🤖 TrackMate AI</h1>
                <p style="margin: 0.5rem 0 0 0; font-size: 1.2rem; opacity: 0.9;">Your smart companion for tracking and managing orders.</p>
            </div>
            
            <div class="documentation-section">
                <h2 style="margin: 0 0 1rem 0; font-size: 1.8rem;">📋 About TrackMate AI</h2>
                <p style="margin: 0; font-size: 1.1rem; line-height: 1.6;">
                    <strong>TrackMate AI</strong> is a Gradio-powered autonomous agent designed to interface with an ERP system via MCP tools. 
                    It connects to a custom MCP Server and lets users chat with an LLM to place, modify, or track orders — all through natural language. Designed for seamless interaction with ERP infrastructure, TrackMate AI empowers users to manage operational workflows via a conversational interface — no dashboards, no dropdowns, just intelligent dialogue. Whether you're tracking an order, generating invoices, or accessing global risk intelligence, TrackMate AI brings an intuitive, natural language layer to your business operations.
                </p>
                
                <h3 style="margin: 1.5rem 0 0.5rem 0; font-size: 1.4rem;">📺 Demo Video</h3>
                <p style="margin: 0; font-size: 1.1rem;">
                    <a href="https://youtu.be/r4l0KA9oy6Q" target="_blank">
                        📹 Watch Demo Video - TrackMate AI Agent
                    </a>
                </p>
                
                <h3 style="margin: 1.5rem 0 0.5rem 0; font-size: 1.4rem;">✨ What It Can Do</h3>
                <div class="feature-list">
                    <div class="feature-item">🔍 Query live ERP data (orders, invoices, customers)</div>
                    <div class="feature-item">�� Place or cancel orders</div>
                    <div class="feature-item">🌐 Check for geopolitical disruptions</div>
                    <div class="feature-item">🧠 Leverage multiple tools via MCP + Gradio agent framework</div>
                    <div class="feature-item">🧮 Use a calculator for quick computations</div>
                    <div class="feature-item">☀️🌧️ Get live weather updates for any location</div>
                </div>
                
                <h3 style="margin: 1.5rem 0 0.5rem 0; font-size: 1.4rem;">🔌 Connected Tools</h3>
                <p style="margin: 0 0 1rem 0; font-size: 1.1rem;">
                    TrackMate AI uses a bunch of tools provided by the 
                    <a href="https://huggingface.co/spaces/Agents-MCP-Hackathon/TrackMate-AI-MCP-Server" target="_blank">
                        TrackMate AI MCP Server
                    </a>:
                </p>
                <table style="width: 100%; border-collapse: collapse; border: 1px solid white;">
                    <thead>
                        <tr>
                            <th style="border: 1px solid white; padding: 8px; background-color: #333; color: white;">Tool Name</th>
                            <th style="border: 1px solid white; padding: 8px; background-color: #333; color: white;">Description</th>
                        </tr>
                    </thead>
                    <tbody>
                        <tr>
                            <td style="border: 1px solid white; padding: 8px; color: white;"><code>list_tables</code></td>
                            <td style="border: 1px solid white; padding: 8px; color: white;">Lists all ERP database tables</td>
                        </tr>
                        <tr>
                            <td style="border: 1px solid white; padding: 8px; color: white;"><code>execute_query</code></td>
                            <td style="border: 1px solid white; padding: 8px; color: white;">Executes a custom SQL query</td>
                        </tr>
                        <tr>
                            <td style="border: 1px solid white; padding: 8px; color: white;"><code>get_order_status</code></td>
                            <td style="border: 1px solid white; padding: 8px; color: white;">Gets status of an order</td>
                        </tr>
                        <tr>
                            <td style="border: 1px solid white; padding: 8px; color: white;"><code>place_order</code></td>
                            <td style="border: 1px solid white; padding: 8px; color: white;">Places a new order</td>
                        </tr>
                        <tr>
                            <td style="border: 1px solid white; padding: 8px; color: white;"><code>cancel_order</code></td>
                            <td style="border: 1px solid white; padding: 8px; color: white;">Cancels an existing order</td>
                        </tr>
                        <tr>
                            <td style="border: 1px solid white; padding: 8px; color: white;"><code>get_active_disruptions</code></td>
                            <td style="border: 1px solid white; padding: 8px; color: white;">Returns geopolitical disruptions between nations</td>
                        </tr>
                        <tr>
                            <td style="border: 1px solid white; padding: 8px; color: white;"><code>add, subtract, multiple</code></td>
                            <td style="border: 1px solid white; padding: 8px; color: white;">Use a simple calculator</td>
                        </tr>
                    </tbody>
                </table>

            </div>
        </div>
        """)
        
        # Main Interface
        with gr.Row(equal_height=True):
            # Configuration Panel
            with gr.Column(scale=1):
                with gr.Group(elem_classes="config-panel"):
                    gr.HTML('<div class="panel-header"><h3 style="margin: 0;">⚙️ Agent Configuration</h3></div>')
                    
                    model_name = gr.Dropdown(
                        choices=["gpt-4o-mini", "gpt-4o", "gpt-4"],
                        value="gpt-4o",
                        label="🤖 Model",
                        info="Select the OpenAI model to use"
                    )
                    temperature = gr.Slider(
                        0.0, 1.0, 
                        value=0.7, 
                        step=0.1, 
                        label="🌡️ Temperature",
                        info="Controls randomness (0=deterministic, 1=creative)"
                    )
                    max_tokens = gr.Number(
                        value=1000, 
                        minimum=100, 
                        maximum=4000, 
                        step=100, 
                        label="📏 Max Tokens",
                        info="Maximum response length"
                    )
                    api_key = gr.Textbox(
                        value="",
                        label="🔑 API Key",
                        placeholder="Enter your OpenAI API key",
                        type="password",
                        interactive=True,
                        info="Your OpenAI API key (kept secure)"
                    )
                    
                    with gr.Accordion("📝 System Prompt", open=False):
                        system_prompt = gr.Textbox(
                            value=get_default_system_prompt(),
                            lines=10,
                            label="System Prompt",
                            interactive=True,
                            info="Customize the agent's behavior and instructions"
                        )

                    init_btn = gr.Button("🚀 Initialize Agent", variant="primary", size="lg")
                    
                    with gr.Group():
                        status_display = gr.Textbox(
                            label="📊 Status", 
                            interactive=False,
                            info="Current operation status"
                        )
                        agent_status = gr.Textbox(
                            label="🤖 Agent Status", 
                            interactive=False,
                            info="Agent readiness indicator"
                        )

            # Chat Panel
            with gr.Column(scale=3, min_width=700):
                with gr.Group(elem_classes="chat-panel"):
                    gr.HTML('<div class="panel-header"><h3 style="margin: 0;">💬 Chat Interface</h3></div>')
                    
                    chatbot = gr.Chatbot(
                        label="TrackMate AI Chat", 
                        height=650,
                        show_label=False,
                        show_copy_button=True,
                        elem_classes="chatbot"
                    )
                    
                    # Make the message input take full width
                    msg_input = gr.Textbox(
                        label="💭 Message",
                        placeholder="Ask me about orders, inventory, weather, or anything else...",
                        show_label=False
                    )

                    send_btn = gr.Button("📤 Send", variant="primary", scale=1)
                    
                    reset_btn = gr.Button("🔄 Reset Chat", variant="secondary", scale=1)

        # Event handlers
        init_btn.click(
            fn=initialize_agent_handler,
            inputs=[model_name, temperature, max_tokens, system_prompt, api_key],
            outputs=[agent_state, status_display, agent_status]
        )

        reset_btn.click(
            fn=reset_conversation,
            outputs=[chatbot, status_display, agent_status]
        )

        send_btn.click(
            fn=chat_with_agent,
            inputs=[agent_state, msg_input, chatbot],
            outputs=[chatbot, msg_input]
        )

        msg_input.submit(
            fn=chat_with_agent,
            inputs=[agent_state, msg_input, chatbot],
            outputs=[chatbot, msg_input]
        )

    return app

if __name__ == "__main__":
    app = create_interface()
    app.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False, share=True)