import gradio as gr
import requests
from smolagents import CodeAgent
import logging
from tenacity import retry, stop_after_attempt, wait_exponential

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

# Correct URL based on your Space
HF_SPACE_URL = "https://manavraj-troubleshoot-mcp.hf.space"

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def call_mcp_server(message, tool_type="knowledge_base"):
    """Call MCP server with proper endpoint routing"""
    try:
        endpoint_map = {
            "knowledge_base": "api/knowledge_base",
            "web_search": "api/web_search",
            "formatter": "api/formatter"
        }
        
        endpoint = endpoint_map.get(tool_type, "api/knowledge_base")
        url = f"{HF_SPACE_URL}/{endpoint}"
        
        response = requests.post(
            url,
            json={"data": [message]},
            timeout=30
        )
        
        if response.status_code != 200:
            raise Exception(f"MCP server returned {response.status_code}")
            
        data = response.json()
        if not isinstance(data, dict) or 'data' not in data:
            raise Exception("Invalid MCP response format")
            
        return data['data'][0]
        
    except Exception as e:
        logger.error(f"MCP call failed: {str(e)}")
        return f"Error: {str(e)}"

# Initialize CodeAgent with proper system prompt including required tags
agent = CodeAgent(
    tools=[],  # No tools defined here as they're in MCP
    model="microsoft/DialoGPT-medium",
    system_prompt="""{{authorized_imports}}
- requests for API calls
- standard Python libraries

{{managed_agents_descriptions}}

You are an advanced Technical Support Assistant with these capabilities:
1. Troubleshooting technical issues (wifi, hardware, software)
2. Finding information through web search
3. Formatting instructions into clear steps

Tools available via MCP server:
- knowledge_base: For technical issues (wifi, screen, sound problems)
- web_search: For finding latest information or non-technical queries
- formatter: To organize raw steps into numbered lists

Response Guidelines:
1. INTERNALLY follow THOUGHT/ACTION/OBSERVATION cycle
2. Use knowledge_base for technical troubleshooting
3. Use web_search for information lookup
4. Use formatter when user provides unorganized steps
5. Provide FINAL RESPONSE in clear, helpful language

Example workflow:
THOUGHT: User has wifi issue, should check knowledge base
ACTION: knowledge_base
OBSERVATION: Found 4 troubleshooting steps
FINAL RESPONSE: Here's how to fix your wifi: [steps]

Never show THOUGHT/ACTION/OBSERVATION to user - only final response."""
)

def chat_interface(message, history):
    """Simplified chat interface"""
    try:
        # Let agent handle the complete process
        response = agent.run(message)
        return str(response).split("FINAL RESPONSE:")[-1].strip()
    except Exception as e:
        logger.error(f"Chat error: {str(e)}")
        return f"Error processing request: {str(e)}"

# Gradio interface
demo = gr.ChatInterface(
    fn=chat_interface,
    title="🔧 Technical Support Agent",
    examples=[
        "My wifi keeps disconnecting",
        "Find the latest Windows 11 update",
        "Format: Restart. Check connections. Update drivers"
    ]
)

if __name__ == "__main__":
    demo.launch()