import os
import sys
import json
import asyncio
import traceback
import uuid
import shutil
import base64
import multiprocessing
from pathlib import Path
import glob
from PIL import Image
import io
import sqlite3
import pandas as pd
import gradio as gr
from dotenv import load_dotenv

from dynamic_agent import AgentFactory
from agno_kb import AgnoKnowledgeBase
from agno.tools.mcp import MCPTools

# Load environment variables
load_dotenv()

# Set event loop policy for Windows
if sys.platform.startswith("win"):
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

# MCP server runner from mcp_tools.py
def run_mcp_server():
    import mcp_tools
    
# Create a directory for storing session files
SESSION_FILES_DIR = Path("session_files")
SESSION_FILES_DIR.mkdir(exist_ok=True)

# Define the images folder path
IMAGES_FOLDER_PATH = Path("./plots")

# Available tools for agent configuration
available_tools = ["MCP Server Tool", "KB Configuration"]


DB_PATH = "flipkart_mobiles.db"
TABLE_NAME = "mobiles"

def load_entire_database():
    try:
        conn = sqlite3.connect(DB_PATH)
        df = pd.read_sql_query(f"SELECT * FROM {TABLE_NAME}", conn)
        conn.close()
        return df
    except Exception as e:
        return pd.DataFrame({"Error": [str(e)]})
    
def clear_database_view():
    # Return an empty DataFrame to clear/hide the display
    return pd.DataFrame()
    
# Print the last assistant message from response
async def last_message(response):
    last_msg = next(
        (m.content for m in reversed(response.messages or []) if m.role == "assistant"),
        None
    )
    return last_msg


# === Deployment Handler (FIXED) ===
def handle_deploy_and_close(platform):
    """Handle deployment and close modal in one action"""
    if not platform:
        # Return error message but keep modal open
        return "Please select a platform.", gr.update(visible=True)
    
    # Get current host URL (placeholder logic)
    current_url = os.getenv("HOST_URL", "http://127.0.0.1:7860")
    success_msg = f"✅ Successfully deployed into {platform}!\n🔗 {current_url}"
    # success_msg = f"✅ Successfully deployed into {platform}!\n🔗 [{current_url}]({current_url})"

    # Return success message and close modal
    return success_msg, gr.update(visible=False)

def show_deploy_modal():
    """Show the deploy modal"""
    return gr.update(visible=True)

def hide_deploy_modal():
    """Hide the deploy modal"""
    return gr.update(visible=False)

# === Utility Functions ===

def toggle_visibility(current):
    new_state = not current
    return gr.update(visible=new_state), new_state

def handle_tool_selection(selected_tools):
    # Always ensure "MCP Server Tool" is included
    if "MCP Server Tool" not in selected_tools:
        selected_tools = ["MCP Server Tool"] + selected_tools
    return gr.update(value=selected_tools)

def save_agent(name, desc, instructions, tools, session_data):
    if name:
        agent = {
            "name": name, 
            "desc": desc, 
            "instructions": instructions, 
            "tools": tools,
            "kb_enabled": "KB Configuration" in tools  # Track if KB is enabled
        }
        session_data["agents"].append(agent)
        return (
            gr.update(choices=[a["name"] for a in session_data["agents"]], value=name),
            "Agent saved successfully!",  # This is the missing return value for the textbox
            session_data,
        )
    return gr.update(), "Agent name is required!", session_data

def store_files_in_session(uploaded_files, session_data):
    """Store uploaded files in session storage and return their paths."""
    if not uploaded_files:
        return []
    
    # Initialize session_files if not exists
    if "session_files" not in session_data:
        session_data["session_files"] = {}
    
    stored_file_paths = []
    
    for file in uploaded_files:
        # Generate unique filename to avoid conflicts
        file_id = str(uuid.uuid4())
        original_name = Path(file.name).name
        file_extension = Path(file.name).suffix
        stored_filename = f"{file_id}_{original_name.replace(' ', '_')}"
        
        # Create session-specific directory
        session_dir = SESSION_FILES_DIR / file_id[:8]  # Use first 8 chars of UUID
        session_dir.mkdir(exist_ok=True)
        
        # Copy file to session storage
        stored_file_path = session_dir / stored_filename
        shutil.copy2(file.name, stored_file_path)
        
        # Store file metadata in session
        session_data["session_files"][file_id] = {
            "original_name": original_name,
            "stored_path": str(stored_file_path),
            "file_extension": file_extension,
            "file_size": os.path.getsize(stored_file_path)
        }
        
        stored_file_paths.append({
            "file_id": file_id,
            "original_name": original_name,
            "stored_path": str(stored_file_path)
        })
    
    return stored_file_paths

def save_kb(uploaded_files, name, session_data):
    if not name:
        return (
            gr.update(),
            gr.update(),
            gr.update(),
            "Knowledge Base name required.",
            session_data
        )
    
    if uploaded_files:
        # Store files in session storage
        stored_files = store_files_in_session(uploaded_files, session_data)
        
        # Create knowledge base entry with stored file references
        kb_entry = {
            "name": name,
            "files": stored_files,  # Store file metadata instead of paths
            "created_at": str(uuid.uuid4())  # Add timestamp/id for uniqueness
        }
        
        session_data["kb"].append(kb_entry)
        
        file_names = ", ".join([f["original_name"] for f in stored_files])
        message = f"""✅ Knowledge Base Saved!\nName: {name}\nFiles: {file_names}"""
        
        # Disable upload and save after saving
        return (
            gr.update(interactive=False),
            gr.update(interactive=False),
            gr.update(interactive=False),
            message,
            session_data
        )
    else:
        return (
            gr.update(),
            gr.update(),
            gr.update(),
            "⚠ No files uploaded. Please upload files before saving.",
            session_data
        )

def get_session_file_paths(session_data, kb_name):
    """Retrieve actual file paths for a given knowledge base from session storage."""
    kb_entries = session_data.get("kb", [])
    
    for kb in kb_entries:
        if kb["name"] == kb_name:
            file_paths = []
            for file_info in kb.get("files", []):
                if isinstance(file_info, dict) and "stored_path" in file_info:
                    # Verify file still exists
                    if os.path.exists(file_info["stored_path"]):
                        file_paths.append(file_info["stored_path"])
                    else:
                        print(f"Warning: File {file_info['original_name']} no longer exists at {file_info['stored_path']}")
                elif isinstance(file_info, str):
                    # Handle legacy format
                    if os.path.exists(file_info):
                        file_paths.append(file_info)
            return file_paths
    
    return []

def build_configs_from_session(session_data):
    kb_entries = session_data.get("kb", [])
    agents = session_data.get("agents", [])
    
    # Get the selected/active agent
    selected_agent = agents[-1] if agents else {}
    
    # Check if the selected agent has KB Configuration enabled
    kb_enabled = selected_agent.get("kb_enabled", False)
    
    if not kb_enabled or not kb_entries:
        # No KB configuration or agent doesn't use KB
        agno_kb_config = {
            "knowledge_base": {
                "collection_name": "default_collection",
                "chunk_size": 1000,
                "overlap": 200,
                "num_documents": 6,
                "chunking_strategy": "fixed",
                "recreate": False,
                "input_data": {
                    "type": "pdf",
                    "source": []
                }
            },
            "instructions": {
                "collections_to_search": "default_collection"
            }
        }
    else:
        # Use KB configuration since agent has KB enabled
        selected_kb = kb_entries[0]  # Use the first KB entry
        kb_name = selected_kb["name"]
        
        # Get actual file paths from session storage
        file_paths = get_session_file_paths(session_data, kb_name)
        
        agno_kb_config = {
            "knowledge_base": {
                "collection_name": kb_name,
                "chunk_size": 1000,
                "overlap": 200,
                "num_documents": 6,
                "chunking_strategy": "fixed",
                "recreate": False,
                "input_data": {
                    "type": "pdf",
                    "source": [{"path": file_paths}] if file_paths else []
                }
            },
            "instructions": {
                "collections_to_search": kb_name
            }
        }

    # Base agent config from session
    agent_description = selected_agent.get("desc", "Agent Description")
    agent_instructions = selected_agent.get("instructions", [])

    # Extended defaults
    default_description = (
        "This agent dynamically selects and uses tools based on the user's query. "
        "It can access two main resources: an MCP tools with a SQL database and a knowledge base."
    )

    default_instructions = [
        "Carefully analyze the user's question to determine the context.",
        "If a query is related to the knowledge base, use the knowledge base to answer or perform the action.",
        "If the question is related to SQL database operations—such as retrieving data—use the MCP tools to answer or perform the action.",
        "If the question is related to generation of graph, plot, chart, or visualization, use the MCP tools 'get_columns_info_from_database' tool to get the necessary columns information then use 'generate_python_code' tool to generate the python code and finally use 'visualization_tool' tool to execute the python code and generate the visualization. While generating python code for chart or plot or graph use different and attractive color combinations for visualization. It should be multi-color and attractive for better visualization.",
        "If the query requires both, prioritize extracting structured data from the SQL database first, then supplement with information from the knowledge base as needed.",
        "Always respond with accurate, concise, and relevant information based on the selected tool."
    ]

    # Merge instructions
    if isinstance(agent_instructions, str):
        agent_instructions = [agent_instructions]
    elif not isinstance(agent_instructions, list):
        agent_instructions = []

    agent_config = {
        "name": selected_agent.get("name", "Default Agent"),
        "description": f"{agent_description}\n\n{default_description}",
        "instructions": agent_instructions + default_instructions,
        "kb_enabled": kb_enabled
    }

    return agno_kb_config, agent_config

def cleanup_session_files(session_data):
    """Clean up session files when session ends (optional)."""
    session_files = session_data.get("session_files", {})
    
    for file_id, file_info in session_files.items():
        file_path = file_info.get("stored_path")
        if file_path and os.path.exists(file_path):
            try:
                os.remove(file_path)
                # Also try to remove the directory if it's empty
                parent_dir = Path(file_path).parent
                if parent_dir.exists() and not any(parent_dir.iterdir()):
                    parent_dir.rmdir()
            except OSError as e:
                print(f"Error cleaning up file {file_path}: {e}")

# === Backend query handler with routing functionality ===
async def handle_query(query, chat_history, session_data):
    try:
        user_id = str(uuid.uuid4())
        thread_id = str(uuid.uuid4())

        agno_kb_config, agent_config = build_configs_from_session(session_data)

        print("Knowledge base config:", json.dumps(agno_kb_config, indent=2))
        print("Agent config:", json.dumps(agent_config, indent=2))

        if agent_config.get("kb_enabled") and agno_kb_config["knowledge_base"]["input_data"]["source"]:
            file_paths = agno_kb_config["knowledge_base"]["input_data"]["source"][0].get("path", [])
            existing_files = [f for f in file_paths if os.path.exists(f)]
            print(f"Files to process: {len(existing_files)} out of {len(file_paths)}")
            if len(existing_files) != len(file_paths):
                print("Warning: Some files are missing!")

        agno_kb = None
        if agent_config.get("kb_enabled"):
            agno_kb_module = AgnoKnowledgeBase(
                query=query, 
                user_id=user_id, 
                thread_id=thread_id, 
                agno_kb_config=agno_kb_config
            )
            agno_kb = agno_kb_module.setup_knowledge_base()

        async with MCPTools(url="http://127.0.0.1:7863/gradio_api/mcp/sse", transport="sse") as mcp_tool:
            agent_factory = AgentFactory(user_id, thread_id, agent_config, knowledge_base=agno_kb)
            router = await agent_factory.routing_agent()
            router_event = await router.arun(query, stream=False)
            router_event_resp = await last_message(router_event)

            if router_event_resp and router_event_resp.lower().strip() == "visualization":
                model_name = "Qwen/Qwen3-235B-A22B"
            else:
                model_name = "meta-llama/Meta-Llama-3.1-405B-Instruct"

            agent = await agent_factory.normal_and_reasoning_agent(
                tools=[mcp_tool],
                model_name=model_name
            )

            response = await agent.arun(query, stream=False)

            last_msg = next(
                (m.content for m in reversed(response.messages or []) if m.role == "assistant"),
                "⚠️ No response from agent."
            )

            image_paths = get_images_from_folder()

            updated_chat_history = chat_history + [{"role": "user", "content": query}]

            if image_paths:
                print(f"Found {len(image_paths)} images to process")

                formatted_response = format_response_with_images(last_msg, len(image_paths))
                updated_chat_history.append({
                    "role": "assistant", 
                    "content": formatted_response
                })

                for img_path in image_paths:
                    try:
                        img_html = add_image_as_base64_html(img_path)
                        if img_html:
                            updated_chat_history.append({
                                "role": "assistant",
                                "content": img_html
                            })
                            print(f"Added image via base64 HTML: {os.path.basename(img_path)}")
                        else:
                            updated_chat_history.append({
                                "role": "assistant", 
                                "content": f"🖼️ Generated visualization: {os.path.basename(img_path)} (Error displaying image)"
                            })

                    except Exception as e:
                        print(f"Error processing image {img_path}: {e}")
                        updated_chat_history.append({
                            "role": "assistant", 
                            "content": f"⚠️ Error displaying image: {os.path.basename(img_path)}"
                        })

                # Delete images after displaying them
                delete_images_from_folder(image_paths)

            else:
                updated_chat_history.append({
                    "role": "assistant", 
                    "content": last_msg
                })

            return updated_chat_history, session_data

    except Exception as e:
        traceback.print_exc()
        error_msg = f"Error processing query: {str(e)}"
        updated_chat_history = chat_history + [
            {"role": "user", "content": query},
            {"role": "assistant", "content": error_msg}
        ]
        return updated_chat_history, session_data


# === Image Handling Functions ===
def get_images_from_folder():
    """Get all image files from the specified folder."""
    if not IMAGES_FOLDER_PATH.exists():
        print(f"Images folder does not exist: {IMAGES_FOLDER_PATH}")
        return []
    
    # Look for common image file extensions
    image_extensions = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.svg', '.webp'}
    image_files = []
    
    try:
        for file_path in IMAGES_FOLDER_PATH.iterdir():
            if file_path.is_file() and file_path.suffix.lower() in image_extensions:
                image_files.append(str(file_path))
        
        print(f"Found {len(image_files)} images in {IMAGES_FOLDER_PATH}")
        return image_files
        
    except Exception as e:
        print(f"Error scanning images folder: {e}")
        return []


def format_response_with_images(text_response, image_count):
    """Format the response to include text and mention images."""
    if not image_count:
        return text_response
    
    # Add a clear separator and mention of visualizations
    separator = "\n\n" + "="*50 + "\n"
    
    if image_count == 1:
        return f"{text_response}{separator}📊 **Generated Visualization** (displayed below):"
    else:
        return f"{text_response}{separator}📊 **Generated {image_count} Visualizations** (displayed below):"


def add_image_as_base64_html(image_path):
    """Convert image to base64 HTML for direct embedding with larger size and better quality."""
    try:
        import io, base64, os
        from PIL import Image

        with Image.open(image_path) as img:
            # Resize for web display while maintaining aspect ratio (larger size)
            img.thumbnail((1200, 900), Image.Resampling.LANCZOS)

            # Convert to RGB if necessary
            if img.mode in ('RGBA', 'LA', 'P'):
                background = Image.new('RGB', img.size, (255, 255, 255))
                if img.mode == 'P':
                    img = img.convert('RGBA')
                background.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None)
                img = background

            # Convert to base64
            buffer = io.BytesIO()
            img.save(buffer, format='PNG')
            img_bytes = buffer.getvalue()
            img_base64 = base64.b64encode(img_bytes).decode('utf-8')

            # Create styled HTML img tag with bigger max width and height
            img_html = f'''
            <div style="text-align: center; margin: 15px 0; padding: 10px; border: 1px solid #e0e0e0; border-radius: 10px; background-color: #fafafa;">
                <img src="data:image/png;base64,{img_base64}" 
                     alt="{os.path.basename(image_path)}" 
                     style="max-width: 90%; max-height: 900px; border-radius: 8px; box-shadow: 0 4px 12px rgba(0,0,0,0.15); display: block; margin: 0 auto;">
                <p style="margin-top: 10px; font-size: 16px; color: #555; text-align: center; font-weight: 600;">
                    📊 {os.path.basename(image_path)}
                </p>
            </div>
            '''
            return img_html

    except Exception as e:
        print(f"Error converting image {image_path} to base64: {e}")
        return None

            
    except Exception as e:
        print(f"Error converting image {image_path} to base64: {e}")
        return None


def delete_images_from_folder(image_paths):
    """Delete specified image files from the folder."""
    deleted_count = 0
    for image_path in image_paths:
        try:
            if os.path.exists(image_path):
                os.remove(image_path)
                deleted_count += 1
                print(f"Deleted image: {os.path.basename(image_path)}")
        except OSError as e:
            print(f"Error deleting image {image_path}: {e}")

    if deleted_count > 0:
        print(f"Successfully deleted {deleted_count} images from folder")

    
with gr.Blocks(theme=gr.themes.Soft(),
    css="""
#agent-popup, #kb-popup {
  position: fixed !important;
  top: 50% !important;
  left: 50% !important;
  transform: translate(-50%, -50%) !important;
  background-color: #1e1e1e !important;
  border: 2px solid #ccc !important;
  border-radius: 10px !important;
  padding: 20px !important;
  z-index: 1000 !important;
  box-shadow: 0 0 20px rgba(0,0,0,0.5);
}

#tool-popup {
  position: fixed !important;
  top: 0 !important;
  left: 0 !important;
  right: 0 !important;
  bottom: 0 !important;
  width: 100vw !important;
  height: 100vh !important;
  max-width: none !important;
  max-height: none !important;
  margin: 0 !important;
  padding: 40px !important;
  background-color: #1e1e1e !important;
  border: none !important;
  border-radius: 0 !important;
  z-index: 9999 !important;
  box-shadow: none !important;
  overflow-y: auto !important;
  transform: none !important;
}

#tool-popup > * {
  max-width: none !important;
}

.deploy-modal {
  background-color: #2d2d2d !important;
  border: 2px solid #444 !important;
  border-radius: 15px !important;
  padding: 30px !important;
  max-width: 450px !important;
  width: 90% !important;
  box-shadow: 0 10px 30px rgba(0,0,0,0.8) !important;
  color: white !important;
}

.modal-overlay {
  position: fixed !important;
  top: 0 !important;
  left: 0 !important;
  right: 0 !important;
  bottom: 0 !important;
  background-color: rgba(0, 0, 0, 0.5) !important;
  z-index: 999 !important;
  display: flex !important;
  align-items: center !important;
  justify-content: center !important;
}

.modal {
  background-color: #1e1e1e !important;
  border: 2px solid #ccc !important;
  border-radius: 10px !important;
  padding: 20px !important;
  max-width: 500px !important;
  width: 90% !important;
  box-shadow: 0 0 20px rgba(0,0,0,0.5) !important;
}

.small-upload .wrap-inner {
    padding: 6px !important;
    font-size: 12px !important;
}
.closing-btn {
    font-size: 14px !important;
    color: #e74c3c;
    background: transparent;
    border: none;
    text-align: left;
    cursor: pointer;
}
""")as app:

    # Initialize session state
    session_data = gr.State({
        "agents": [], 
        "tools": [], 
        "kb": [],
        "session_files": {}  # Add session files storage
    })
    
    kb_visible = gr.State(False)
    tool_visible = gr.State(False)
    agent_visible = gr.State(False)
    deploy_visible = gr.State(False)

    with gr.Group(visible=True):
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("### 📁 Knowledge Base")
                kb_btn = gr.Button("Add KB ➕")
                
                # Display current KB status
                with gr.Group():
                    kb_status_display = gr.Textbox(
                        label="Current KB Status", 
                        value="No Knowledge Base loaded",
                        interactive=False
                    )
                
                gr.Markdown("### 🧰 Tools")
                tool_btn = gr.Button("Add Tool")
                gr.Markdown("### 🤖 Agent")
                agent_btn = gr.Button("Add Agent")
                gr.Markdown("### 🧠 Active Agent")
                active_agent = gr.Radio(choices=[], label="Selected Agent", interactive=True)

            # Style for Send Button
            gr.HTML("""
            <style>
            #send-btn {
                height: 50px;
                font-size: 16px;
                padding: 10px 20px;
                background-color: #6262E2;
                color: white;
                border: none;
                border-radius: 8px;
                cursor: pointer;
            }
            #send-btn:hover {
                background-color: #0056b3;
            }
            </style>
            """)

            with gr.Column(scale=7):
                # gr.Markdown("## Dynamic and Customizable AI Agent")
                gr.Markdown("<h2 style='color:#4CAF50;'>NoCoMind - Dynamic & Customizable AI Agents</h2>")
                deploy_output = gr.Textbox(label="Deployment Status", visible=False, lines=2)

                chatbot = gr.Chatbot(
                    label="Chat", 
                    height=400, 
                    type="messages",
                    elem_classes=["chatbot-container"],
                    show_copy_button=True,
                    bubble_full_width=False,
                    render_markdown=True,
                    value=[]
                )

                with gr.Row():
                    user_input = gr.Textbox(
                        placeholder="Type your message...",
                        show_label=False,
                        scale=4,
                        lines=1
                    )
                    send_btn = gr.Button(
                        "Send", 
                        scale=0.1,
                        elem_id="send-btn"
                    )


    # === Modals ===
    with gr.Group(visible=False, elem_id="kb-popup") as kb_popup:
        gr.Markdown("### Knowledge Base")
        kb_files = gr.File(
            label="Upload Files",
            file_types=[".pdf"],
            file_count="multiple",
            elem_classes=["small-upload"],
            interactive=True
        )
        kb_name = gr.Textbox(label="Name:", interactive=True)
        kb_save = gr.Button("Save", interactive=True)
        kb_output = gr.Textbox(label="Status", lines=4)
        kb_close = gr.Button("❌ Close", elem_classes=["closing-btn"])

    with gr.Group(visible=False, elem_id="tool-popup") as tool_popup:
        gr.Markdown("## Tool Configuration")
        tool_description = gr.Textbox(
            label="Description",
            value='''
            ✨ You're now using a Gradio MCP tools!
            This feature helps our smart agents assist you better with your tasks.
            Just sit back and let the Gradio MCP Agents do the heavy lifting for you!
            '''   ,
            lines = 4
        )

        with gr.Blocks() as demo:
            gr.Markdown(
                '''
                🚀 **Integrated Tools in Gradio MCP Agents**  

                1. 🗃 **SQL Database Engine Tool** – Currenty, Flipkart sample database is already configured. We can add any database in the future.
                Query and interact with structured data effortlessly. Whether you're retrieving, filtering, or analyzing records, this tool lets the agents access your database just like a data expert.

                2. 🐍 **Python Code Generator Tool**  
                Instantly generate and run Python code for your needs. From data processing to simple logic, let the agents write and execute scripts without you touching a single line of code.

                3. 📊 **Visualization Tool**  
                Transform data into beautiful, interactive charts and plots. This tool helps agents present insights in a clear and visually appealing way, making it easier for you to understand and decide.

                📂 **Use the buttons below to open or close the Flipkart database preview (`mobiles` table):**
                '''
            )

            show_db_btn = gr.Button("Show Database Preview")
            hide_db_btn = gr.Button("Hide Database Preview")

            data_preview = gr.Dataframe(
                value=None,
                label="📊 Flipkart DB (Full View)",
                interactive=False,
                wrap=True
            )

            # Button clicks update the dataframe accordingly
            show_db_btn.click(fn=load_entire_database, inputs=None, outputs=data_preview)
            hide_db_btn.click(fn=clear_database_view, inputs=None, outputs=data_preview)

        tool_close = gr.Button("❌ Close", elem_classes=["closing-btn"], variant="secondary")

    with gr.Group(visible=False, elem_id="agent-popup") as agent_popup:
        gr.HTML('<div style="max-height: 400px; overflow-y: auto; padding: 10px;">')
        gr.Markdown("#### Agent Configuration")
        agent_name = gr.Textbox(label="Name", placeholder="Enter agent name")
        agent_desc = gr.Textbox(label="Description", placeholder="Enter description", lines=1)
        agent_instructions = gr.Textbox(label="Instructions", placeholder="Enter specific instructions", lines=2)
        tool_select = gr.CheckboxGroup(
            choices=available_tools,
            label="Tools",
            value=["MCP Server Tool"],  # Always start with MCP Server Tool selected
            info="Note: MCP Server Tool is always required and will be automatically selected"
        )
        agent_save = gr.Button("💾 Save Agent")
        agent_status = gr.Textbox(label="Status", interactive=False,lines = 1)
        agent_close = gr.Button("❌ Close")
        gr.HTML('</div>')

# === Deploy Modal ===

    with gr.Group(visible=False, elem_id="deploy-modal-overlay") as deploy_overlay:
        with gr.Group(elem_id="deploy-modal"):
            gr.Markdown("## 🌍 Choose Deployment Option")
            deploy_choice = gr.Radio(
                choices=["GCP", "AWS", "Azure", "Gradio Cloud"], 
                label="Select Platform",
                value=None  # No default selection
            )
            with gr.Row():
                confirm_deploy = gr.Button("✅ Confirm", variant="primary")
                close_deploy = gr.Button("❌ Close", elem_id="closing-btn", variant="secondary")

    trigger_btn = gr.Button("🚀 Deploy AI Agent")

    # === Event Bindings ===

    kb_btn.click(toggle_visibility, inputs=kb_visible, outputs=[kb_popup, kb_visible])
    kb_close.click(toggle_visibility, inputs=kb_visible, outputs=[kb_popup, kb_visible])

    tool_btn.click(toggle_visibility, inputs=tool_visible, outputs=[tool_popup, tool_visible])
    tool_close.click(toggle_visibility, inputs=tool_visible, outputs=[tool_popup, tool_visible])

    agent_btn.click(toggle_visibility, inputs=agent_visible, outputs=[agent_popup, agent_visible])
    agent_close.click(toggle_visibility, inputs=agent_visible, outputs=[agent_popup, agent_visible])


    # Show modal on trigger
    trigger_btn.click(toggle_visibility, inputs=deploy_visible, outputs=[deploy_overlay, deploy_visible])
    confirm_deploy.click(toggle_visibility, inputs=deploy_visible, outputs=[deploy_overlay, deploy_visible])
    close_deploy.click(toggle_visibility, inputs=deploy_visible, outputs=[deploy_overlay, deploy_visible])
    
    close_deploy.click(
        fn=lambda: gr.update(visible=False),
        outputs=deploy_overlay
    )
    
    # FIXED: Confirm deploy now properly handles the modal state
    confirm_deploy.click(
        fn=handle_deploy_and_close,
        inputs=deploy_choice,
        outputs=[deploy_output, deploy_overlay]
    ).then(
        fn=lambda: gr.update(visible=True),  # Show the deploy_output after successful deployment
        outputs=deploy_output
    )

    # Add event handler to ensure MCP Server Tool stays selected
    tool_select.change(
        fn=handle_tool_selection,
        inputs=[tool_select],
        outputs=[tool_select]
    )

    agent_save.click(
        save_agent,
        inputs=[agent_name, agent_desc, agent_instructions, tool_select, session_data],
        outputs=[active_agent,agent_status, session_data]
    )

    # Updated KB save with session file handling and disabling functionality
    def save_kb_and_update_status(uploaded_files, name, session_data):
        kb_files_update, name_update, save_btn_update, result_msg, updated_session = save_kb(uploaded_files, name, session_data)
        
        # Update KB status display
        kb_list = updated_session.get("kb", [])
        if kb_list:
            latest_kb = kb_list[-1]
            file_count = len(latest_kb.get("files", []))
            status_msg = f"✅ KB '{latest_kb['name']}' loaded with {file_count} files"
        else:
            status_msg = "No Knowledge Base loaded"
        
        return kb_files_update, name_update, save_btn_update, result_msg, updated_session, status_msg

    kb_save.click(
        save_kb_and_update_status,
        inputs=[kb_files, kb_name, session_data],
        outputs=[kb_files, kb_name, kb_save, kb_output, session_data, kb_status_display]
    )

    # UPDATED: Send button with proper chat history management
    def send_message_wrapper(query, current_chat, session_data):
        if not query.strip():
            return current_chat + [{"role": "assistant", "content": "Please enter a question."}], session_data, ""
        
        # Check if agent has KB enabled and KB is available
        agents = session_data.get("agents", [])
        if agents:
            selected_agent = agents[-1]  # Get the last/active agent
            if selected_agent.get("kb_enabled", False):
                kb_list = session_data.get("kb", [])
                if not kb_list:
                    error_response = current_chat + [
                        {"role": "user", "content": query},
                        {"role": "assistant", "content": "⚠️ Agent requires KB Configuration but no Knowledge Base is loaded. Please upload and save a KB first."}
                    ]
                    return error_response, session_data, ""
        
        # Call the async handler with current chat history
        try:
            updated_chat, updated_session = asyncio.run(handle_query(query, current_chat, session_data))
            return updated_chat, updated_session, ""  # Clear input
        except Exception as e:
            print(f"Error in send_message_wrapper: {e}")
            error_response = current_chat + [
                {"role": "user", "content": query},
                {"role": "assistant", "content": f"Error: {str(e)}"}
            ]
            return error_response, session_data, ""

    # NEW: Clear chat function
    def clear_chat_history(session_data):
        """Clear the chat history"""
        session_data["chat_history"] = []
        return [], session_data

    # Updated event handlers for sending messages
    send_btn.click(
        fn=send_message_wrapper,
        inputs=[user_input, chatbot, session_data],
        outputs=[chatbot, session_data, user_input]
    )
    
    # Allow Enter key to send message
    user_input.submit(
        fn=send_message_wrapper,
        inputs=[user_input, chatbot, session_data],
        outputs=[chatbot, session_data, user_input]
    )

if __name__ == "__main__":
    mcp_process = multiprocessing.Process(target=run_mcp_server)
    mcp_process.start()
    
    # Give MCP server time to start
    import time
    time.sleep(3)
    
    app.launch(server_port=7860, debug=True)