import gradio as gr
import asyncio
import os
from typing import List, Tuple, Optional, Dict, Any
from datetime import datetime
import logging
import signal
import sys
import json

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

try:
    from mcp_use import MCPClient
    from langchain_mcp_adapters.client import MultiServerMCPClient
    from langchain_community.tools.sleep.tool import SleepTool
    from langchain_mcp_adapters.tools import load_mcp_tools
    from langchain.agents import AgentExecutor, create_tool_calling_agent
    from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
    from langchain_mistralai import ChatMistralAI
except ImportError as e:
    logger.error(f"Import error: {e}")
    raise

class ConversationManager:
    """Manages conversation history with token optimization"""
    
    def __init__(self, max_history_pairs: int = 3, max_context_chars: int = 2000):
        self.max_history_pairs = max_history_pairs
        self.max_context_chars = max_context_chars
        self.session_context = {}  # Browser state context
        
    def update_session_context(self, action: str, result: str):
        """Update browser session context (current page, last actions, etc.)"""
        self.session_context.update({
            'last_action': action,
            'last_result': result[:500],  # Truncate long results
            'timestamp': datetime.now().isoformat()
        })
        
    def get_optimized_history(self, full_history: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
        """Get optimized history with recent messages + session context"""
        
        # Take only the last N conversation pairs
        recent_history = full_history[-self.max_history_pairs:] if full_history else []
        
        # Add session context as first "message" if we have browser state
        if self.session_context:
            context_msg = f"[SESSION_CONTEXT] Browser session active. Last action: {self.session_context.get('last_action', 'none')}"
            recent_history.insert(0, ("system", context_msg))
        
        return recent_history
    
    def get_context_summary(self) -> str:
        """Get a summary of current browser session state"""
        if not self.session_context:
            return "Browser session not active."
        
        return f"Browser session active. Last action: {self.session_context.get('last_action', 'none')} at {self.session_context.get('timestamp', 'unknown')}"

class BrowserAgent:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.client = None
        self.session = None
        self.session_context = None
        self.agent_executor = None
        self.model = None
        self.initialized = False
        self.available_tools = {}
        self.system_prompt = ""
        
        # Add conversation manager for token optimization
        self.conversation_manager = ConversationManager(
            max_history_pairs=3,  # Only keep last 3 exchanges
            max_context_chars=2000  # Limit context size
        )

    async def generate_tools_prompt(self):
        """Generate a detailed prompt section about available tools"""
        try:
            tools_prompt = "\n## 🛠️ AVAILABLE TOOLS\n"
            tools_prompt += "You have access to the following browser automation tools via MCP:\n\n"
            
            for tool_name, tool_info in self.available_tools.items():
                tools_prompt += f"### {tool_name}\n"
                
                # Add description from StructuredTool object
                description = getattr(tool_info, 'description', 'No description available')
                tools_prompt += f"**Description**: {description}\n"
                
                # Add parameters from args_schema if available
                if hasattr(tool_info, 'args_schema') and tool_info.args_schema:
                    try:
                        schema = tool_info.args_schema.model_json_schema()
                        if 'properties' in schema:
                            tools_prompt += "**Parameters**:\n"
                            for param_name, param_info in schema['properties'].items():
                                param_type = param_info.get('type', 'unknown')
                                param_desc = param_info.get('description', 'No description')
                                required = param_name in schema.get('required', [])
                                required_mark = " (required)" if required else " (optional)"
                                tools_prompt += f"- `{param_name}` ({param_type}){required_mark}: {param_desc}\n"
                    except Exception as schema_error:
                        logger.debug(f"Could not parse schema for {tool_name}: {schema_error}")
                        tools_prompt += "**Usage**: Call this tool when you need to perform this browser action\n"
                else:
                    tools_prompt += "**Usage**: Call this tool when you need to perform this browser action\n"
                
                tools_prompt += "\n"
            
            tools_prompt += """
🎯 Multi‑Step Workflow
Navigate & Snapshot
Load the target page
Capture a snapshot
Assess if further steps are needed—if so, proceed to the next action
Perform Action & Validate
if needed closes add or popups
Capture a snapshot
Verify results before moving on
Keep Browser Open
Never close the session unless explicitly instructed
Avoid Redundancy
Don’t repeat actions (e.g., clicking) when data is already collected
## 🚨 SESSION PERSISTENCE RULES
- Browser stays open for the entire conversation
- Each action builds on previous state
- Context is maintained between requests
"""
            return tools_prompt
        except Exception as e:
            logger.error(f"Failed to generate tools prompt: {e}")
            return "\n## 🛠️ TOOLS\nBrowser automation tools available but not detailed.\n"

    async def get_system_prompt_with_tools(self):
        base = """🌐 Browser Agent — Persistent Session & Optimized Memory
You are an intelligent browser automation agent (Playwright via MCP) tasked with keeping a lightweight, ongoing session:
🎯 Mission
Navigate pages, extract and analyze data without closing the browser
Handle pop‑ups and capture snapshots to validate each step
🔄 Session Management
Browser remains open across user requests
Only recent chat history is provided to save tokens
Session context (current page, recent actions) is maintained separately
⚡ Response Structure
For each action:
State → tool call
Snapshot → confirmation
Next plan (if needed)
💡 Best Practices
Use text selectors and wait for content
Pause 2 s between tool calls
Be concise and focused on the current task it s important as soon as you have the information you came for return it
If earlier context is needed, ask the user to clarify.
"""
        tools_section = await self.generate_tools_prompt()
        return base + tools_section

    async def initialize(self):
        """Initialize MCP client, model, session and agent"""
        try:
            logger.info("🚀 Initializing Browser Agent...")
            
            # LLM
            mistral_key = os.getenv("mistralkey") 
            if not mistral_key:
                raise ValueError("Mistral API key is required")
                
            self.model = ChatMistralAI(
                model="magistral-small-latest", 
                api_key=mistral_key,
        
            )
            logger.info("✅ Mistral LLM initialized with optimized settings")
            
            # MCP client
            self.client = MultiServerMCPClient({
                "browser": {
                    "command": "npx",
                    "args": ["@playwright/mcp@latest", "--browser", "chromium","--user-agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36"],
                    "transport": "stdio"
                }
            })
            logger.info("✅ MCP client created")

            # Start persistent session
            self.session_context = self.client.session("browser")
            self.session = await self.session_context.__aenter__()
            logger.info("✅ MCP session opened")
            
            # Load tools
            tools = await load_mcp_tools(self.session)
            tools.append(SleepTool(description="Wait 4 seconds between two calls"))
            logger.info(f"📥 Loaded {len(tools)} tools")
            self.available_tools = {t.name: t for t in tools}

            # Install browser if needed
            install_tool = self.available_tools.get("browser_install")
            if install_tool:
                try:
                    result = await install_tool.arun({})
                    logger.info(f"📥 Browser install: {result}")
                except Exception as e:
                    logger.warning(f"⚠️ Browser install failed: {e}, continuing.")

            # System prompt
            self.system_prompt = await self.get_system_prompt_with_tools()

            # Create agent
            prompt = ChatPromptTemplate.from_messages([
                ("system", self.system_prompt),
                MessagesPlaceholder(variable_name="chat_history"),
                ("human", "{input}"),
                MessagesPlaceholder(variable_name="agent_scratchpad"),
            ])
            agent = create_tool_calling_agent(
                llm=self.model,
                tools=tools,
                prompt=prompt
            )
            self.agent_executor = AgentExecutor(
                agent=agent,
                tools=tools,
                verbose=True,
                max_iterations=10,  # Reduced from 30
                early_stopping_method="generate",
                handle_parsing_errors=True,
                return_intermediate_steps=True,
                max_execution_time=180  # Reduced from 300
            )

            self.initialized = True
            logger.info("✅ Agent initialized with persistent session and optimized memory")
            return True

        except Exception as e:
            logger.error(f"❌ Initialization failed: {e}")
            await self.cleanup()
            raise

    async def process_query(self, query: str, chat_history: List[Tuple[str, str]]) -> str:
        if not self.initialized:
            return "❌ Agent not initialized. Please restart the application."
            
        try:
            # ✅ KEY OPTIMIZATION: Use only recent history instead of full history
            optimized_history = self.conversation_manager.get_optimized_history(chat_history)
            
            # Convert to message format
            history_messages = []
            for human, ai in optimized_history:
                if human: history_messages.append(("human", human))
                if ai: history_messages.append(("ai", ai))
            
            # Add session context
            context_summary = self.conversation_manager.get_context_summary()
            enhanced_query = f"{query}\n\n[SESSION_INFO]: {context_summary}"
            
            # Log token savings
            original_pairs = len(chat_history)
            optimized_pairs = len(optimized_history)
            logger.info(f"💰 Token optimization: {original_pairs} → {optimized_pairs} history pairs")
            
            # Execute with optimized history
            resp = await self.agent_executor.ainvoke({
                "input": enhanced_query,
                "chat_history": history_messages
            })
            
            # Update session context with this interaction
            self.conversation_manager.update_session_context(
                action=query,
                result=resp["output"]
            )
            
            return resp["output"]
            
        except Exception as e:
            logger.error(f"Error processing query: {e}")
            return f"❌ Error: {e}\n💡 Ask for a screenshot to diagnose."

    async def cleanup(self):
        """Cleanup resources properly"""
        try:
            if self.session_context:
                await self.session_context.__aexit__(None, None, None)
                logger.info("✅ MCP session closed")
                self.session_context = None
                self.session = None
                
            if self.client:
                await self.client.close()
                logger.info("✅ MCP client closed")
                self.client = None
                
            self.initialized = False
            
        except Exception as e:
            logger.error(f"Cleanup error: {e}")

    def get_token_usage_stats(self, full_history: List[Tuple[str, str]]) -> Dict[str, Any]:
        """Get statistics about token usage optimization"""
        original_pairs = len(full_history)
        optimized_pairs = len(self.conversation_manager.get_optimized_history(full_history))
        
        # Rough token estimation (1 token ≈ 4 characters)
        def estimate_tokens(text: str) -> int:
            return len(text) // 4
        
        original_tokens = sum(estimate_tokens(msg[0] + msg[1]) for msg in full_history)
        optimized_tokens = sum(estimate_tokens(msg[0] + msg[1]) for msg in self.conversation_manager.get_optimized_history(full_history))
        
        return {
            "original_pairs": original_pairs,
            "optimized_pairs": optimized_pairs,
            "pairs_saved": original_pairs - optimized_pairs,
            "estimated_original_tokens": original_tokens,
            "estimated_optimized_tokens": optimized_tokens,
            "estimated_tokens_saved": original_tokens - optimized_tokens,
            "savings_percentage": ((original_tokens - optimized_tokens) / original_tokens * 100) if original_tokens > 0 else 0
        }

# Global agent instance
agent: Optional[BrowserAgent] = None
event_loop: Optional[asyncio.AbstractEventLoop] = None

async def initialize_agent_async(api_key: str) -> str:
    """Initialize the agent asynchronously"""
    global agent
    
    if not api_key.strip():
        return "❌ Please provide a Mistral API key"
        
    try:
        # Cleanup existing agent
        if agent:
            await agent.cleanup()
            
        # Create new agent
        agent = BrowserAgent(api_key)
        await agent.initialize()
        
        info = await agent.get_system_prompt_with_tools()
        return f"✅ Agent Initialized Successfully with Token Optimization!\n\n{info[:1000]}..."
        
    except Exception as e:
        logger.error(f"Initialization error: {e}")
        return f"❌ Failed to initialize agent: {e}"

async def process_message_async(message: str, history: List[List[str]]) -> Tuple[str, List[List[str]]]:
    """Process message asynchronously with token optimization"""
    global agent
    
    if not agent or not agent.initialized:
        error_msg = "❌ Agent not initialized. Please initialize first with your API key."
        history.append([message, error_msg])
        return "", history
        
    if not message.strip():
        error_msg = "Please enter a message"
        history.append([message, error_msg])
        return "", history
        
    try:
        # Convert history format for the agent
        agent_history = [(msg[0], msg[1]) for msg in history]
        
        # Get token usage stats before processing
        stats = agent.get_token_usage_stats(agent_history)
        
        # Process the query with optimized history
        response = await agent.process_query(message, agent_history)
        
        # Add token savings info to response if significant savings
        if stats["savings_percentage"] > 50:
            response += f"\n\n💰 Token savings: {stats['savings_percentage']:.1f}% ({stats['estimated_tokens_saved']} tokens saved)"
        
        # Add to history
        history.append([message, response])
        
        return "", history
        
    except Exception as e:
        logger.error(f"Message processing error: {e}")
        error_msg = f"❌ Error: {e}\n💡 Try asking for a screenshot to diagnose."
        history.append([message, error_msg])
        return "", history

def run_in_event_loop(coro):
    """Run coroutine in the event loop"""
    global event_loop
    if event_loop and not event_loop.is_closed():
        return asyncio.run_coroutine_threadsafe(coro, event_loop).result()
    else:
        return asyncio.run(coro)

# Sync wrappers for Gradio
def initialize_agent_sync(api_key: str) -> str:
    """Sync wrapper for agent initialization"""
    return run_in_event_loop(initialize_agent_async(api_key))

def process_message_sync(message: str, history: List[List[str]]) -> Tuple[str, List[List[str]]]:
    """Sync wrapper for message processing"""
    return run_in_event_loop(process_message_async(message, history))

def get_token_stats_sync(history: List[List[str]]) -> str:
    """Get token usage statistics"""
    global agent
    if not agent or not agent.initialized:
        return "Agent not initialized"
    
    agent_history = [(msg[0], msg[1]) for msg in history]
    stats = agent.get_token_usage_stats(agent_history)
    
    return f"""📊 Token Usage Statistics:
• Original conversation pairs: {stats['original_pairs']}
• Optimized conversation pairs: {stats['optimized_pairs']}
• Pairs saved: {stats['pairs_saved']}
• Estimated original tokens: {stats['estimated_original_tokens']:,}
• Estimated optimized tokens: {stats['estimated_optimized_tokens']:,}
• Estimated tokens saved: {stats['estimated_tokens_saved']:,}
• Savings percentage: {stats['savings_percentage']:.1f}%"""

def create_interface():
    """Create the Gradio interface"""
    
    with gr.Blocks(
        title="MCP Browser Agent - Token Optimized",
        theme=gr.themes.Soft(),
        css="""
        .container { max-width: 1200px; margin: auto; }
        .header { text-align: center; margin-bottom: 2rem; }
        .status-box { padding: 1rem; border-radius: 8px; margin: 1rem 0; }
        .token-stats { background: #f0f8ff; padding: 1rem; border-radius: 8px; }
        """
    ) as interface:
        
        gr.HTML("""
        <div class="header">
            <h1>🌐 MCP Browser Agent - Token Optimized</h1>
            <p>AI-powered web browsing with persistent sessions and optimized token usage</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 🔧 Configuration")
                api_key_input = gr.Textbox(
                    label="Mistral API Key",
                    placeholder="Enter your Mistral API key...",
                    type="password",
                    lines=1
                )
                
                init_button = gr.Button("Initialize Agent", variant="primary")
                status_output = gr.Textbox(
                    label="Status & Available Tools",
                    interactive=False,
                    lines=6
                )
                
                gr.Markdown("### 💰 Token Optimization")
                token_stats_button = gr.Button("Show Token Stats", variant="secondary")
                token_stats_output = gr.Textbox(
                    label="Token Usage Statistics",
                    interactive=False,
                    lines=8
                )
                
                gr.Markdown("""
                ### 📝 Optimized Usage Tips
                **Token Savings Features:**
                - Only last 3 conversation pairs sent to API
                - Session context maintained separately
                - Reduced max tokens per response
                - Smart context summarization
                
                **Best Practices:**
                - Be specific in your requests
                - Use "take screenshot" to check current state
                - Ask for "browser status" if you need context
                - Long conversations automatically optimized
                """)
                
            with gr.Column(scale=2):
                gr.Markdown("### 💬 Chat with Browser Agent")
                
                chatbot = gr.Chatbot(
                    label="Conversation",
                    height=500,
                    show_copy_button=True
                )
                
                with gr.Row():
                    message_input = gr.Textbox(
                        label="Message",
                        placeholder="Enter your browsing request...",
                        lines=2,
                        scale=4
                    )
                    send_button = gr.Button("Send", variant="primary", scale=1)
                
                with gr.Row():
                    clear_button = gr.Button("Clear Chat", variant="secondary")
                    screenshot_button = gr.Button("Quick Screenshot", variant="secondary")
        
        # Event handlers
        init_button.click(
            fn=initialize_agent_sync,
            inputs=[api_key_input],
            outputs=[status_output]
        )
        
        send_button.click(
            fn=process_message_sync,
            inputs=[message_input, chatbot],
            outputs=[message_input, chatbot]
        )
        
        message_input.submit(
            fn=process_message_sync,
            inputs=[message_input, chatbot],
            outputs=[message_input, chatbot]
        )
        
        clear_button.click(
            fn=lambda: [],
            outputs=[chatbot]
        )
        
        screenshot_button.click(
            fn=lambda history: process_message_sync("Take a screenshot of the current page", history),
            inputs=[chatbot],
            outputs=[message_input, chatbot]
        )
        
        token_stats_button.click(
            fn=get_token_stats_sync,
            inputs=[chatbot],
            outputs=[token_stats_output]
        )
        
        # Add helpful information
        with gr.Accordion("ℹ️ Token Optimization Guide", open=False):
            gr.Markdown("""
            ## 💰 How Token Optimization Works
            
            **The Problem with Original Code:**
            - Every API call sent complete conversation history
            - Token usage grew exponentially with conversation length
            - Costs could explode for long sessions
            
            **Our Optimization Solutions:**
            
            1. **Limited History Window**: Only last 3 conversation pairs sent to API
            2. **Session Context**: Browser state maintained separately from chat history
            3. **Smart Summarization**: Key session info added to each request
            4. **Reduced Limits**: Lower max_tokens and max_iterations
            5. **Token Tracking**: Real-time savings statistics
            
            **Token Savings Example:**
            ```
            Original: 10 messages = 5,000 tokens per API call
            Optimized: 10 messages = 500 tokens per API call
            Savings: 90% reduction in token usage!
            ```
            
            **What This Means:**
            - ✅ Persistent browser sessions still work
            - ✅ 90%+ reduction in API costs
            - ✅ Faster response times
            - ✅ Better performance for long conversations
            - ⚠️ Agent has limited memory of old messages
            
            **If Agent Needs Earlier Context:**
            - Use "browser status" to check current state
            - Take screenshots to show current page
            - Re-explain context if needed
            - Clear chat periodically for fresh start
            """)
    
    return interface

async def cleanup_agent():
    """Cleanup agent resources"""
    global agent
    if agent:
        await agent.cleanup()
        logger.info("🧹 Agent cleaned up")

def signal_handler(signum, frame):
    """Handle shutdown signals"""
    logger.info(f"📡 Received signal {signum}, cleaning up...")
    global event_loop
    if event_loop and not event_loop.is_closed():
        event_loop.create_task(cleanup_agent())
    sys.exit(0)

async def main():
    """Main async function to run everything"""
    global event_loop
    
    # Set up signal handlers
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    # Get the current event loop
    event_loop = asyncio.get_event_loop()
    
    try:
        logger.info("🚀 Starting MCP Browser Agent Application with Token Optimization...")
        
        # Create and launch interface
        interface = create_interface()
        
        # Launch interface (this will block)
       
        interface.launch(
            server_name="0.0.0.0",
            server_port=7860,
            debug=False,
            show_error=True,
            quiet=False
        )
       
        
    except Exception as e:
        logger.error(f"Application error: {e}")
    finally:
        await cleanup_agent()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("🛑 Application stopped by user")
    except Exception as e:
        logger.error(f"Fatal error: {e}")
    finally:
        logger.info("👋 Application shutdown complete")