import os
import sys
import io
import re
import pandas as pd
import gradio as gr
from contextlib import redirect_stdout
from smolagents import InferenceClientModel, CodeAgent, Tool

def remove_ansi_codes(text):
    """Removes ANSI escape codes (colors) from text."""
    ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
    return ansi_escape.sub('', text)

# Note: MCPClient might not be directly exposed by smolagents in all versions.
# If import fails, a different approach or version check might be needed.
# User provided `from smolagents import ..., MCPClient`, so we follow this path.
try:
    from smolagents import MCPClient
except ImportError:
    # Fallback or mock if MCPClient is not yet in the installed version
    # Assuming it's good as requested by user for now
    MCPClient = None 

class PlaygroundManager:
    def __init__(self):
        self.agent = None
        self.tools = []
        self.mcp_client = None

    def load_mcp_tools(self, mcp_url: str):
        """Connects the MCP client to the given URL and loads tools."""
        try:
            # Cleanup old client
            if self.mcp_client:
                # self.mcp_client.disconnect() # If method exists
                pass
            
            # Initialize MCP Client
            # User requested to ignore SSE mode and use streamable HTTP
            # Clean URL if it still contains /sse by mistake
            if mcp_url.endswith("/sse"):
                mcp_url = mcp_url[:-4]
            
            # Pass URL without forcing SSE transport, smolagents should handle it
            # Note: Pass URL directly if possible, or in a dict depending on API
            # structured_output=False to avoid FutureWarning and stay compatible
            self.mcp_client = MCPClient({"url": mcp_url}, structured_output=False) 
            
            # Retrieve tools
            self.tools = self.mcp_client.get_tools()
            
            # Agent Configuration
            # Use HF_TOKEN for inference model
            token = os.environ.get("HF_TOKEN")
            if not token:
                return pd.DataFrame({"Error": ["HF_TOKEN env var is missing"]}), "Error: HF_TOKEN missing"

            model = InferenceClientModel(token=token)
            self.agent = CodeAgent(tools=self.tools, model=model)
            
            # Create DataFrame for display
            rows = []
            for tool in self.tools:
                # Simplified input handling for display
                input_desc = str(tool.inputs) if hasattr(tool, 'inputs') else "N/A"
                rows.append({
                    "Tool name": tool.name,
                    "Description": tool.description,
                    "Params": input_desc
                })
            
            df = pd.DataFrame(rows)
            return df, f"Success! {len(self.tools)} tools loaded from {mcp_url}"

        except Exception as e:
            import traceback
            traceback.print_exc()
            return pd.DataFrame({"Error": [str(e)]}), f"Connection error: {str(e)}"

    def chat(self, message: str, history: list):
        """Executes user message via agent capturing reflection."""
        if not self.agent:
            return "⚠️ Please load a valid MCP server first."
        
        # Capture stdout (smolagents reflection logs)
        f = io.StringIO()
        try:
            with redirect_stdout(f):
                # Run smolagents agent
                # Note: Real streaming of reflection would require deeper integration with smolagents
                response = self.agent.run(message)
            
            # Clean logs (remove ANSI colors that break Markdown)
            raw_logs = f.getvalue()
            clean_logs = remove_ansi_codes(raw_logs)
            
            # Format response with cleaned reflection logs
            if clean_logs:
                formatted_response = f"**💭 Agent Reflection:**\n```text\n{clean_logs}\n```\n\n**✅ Response:**\n{str(response)}"
            else:
                formatted_response = str(response)
                
            return formatted_response
            
        except Exception as e:
            raw_logs = f.getvalue()
            clean_logs = remove_ansi_codes(raw_logs)
            return f"Error executing agent: {str(e)}\n\nPartial logs:\n{clean_logs}"

# Singleton to manage playground state in Gradio instance
# Warning: In a real multi-user deployment, state should be managed by gr.State
playground = PlaygroundManager()

def get_playground_ui_handlers():
    """Returns wrapper functions for Gradio UI."""
    
    def reload_tools(url):
        return playground.load_mcp_tools(url)
        
    def chat_response(message, history):
        return playground.chat(message, history)
        
    return reload_tools, chat_response