import os
import asyncio
import logging
import sys
import time
import json
import httpx
from datetime import datetime
from typing import Any, Dict, List, Optional, Union
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse
from pydantic import BaseModel, Field
import uvicorn
from contextlib import asynccontextmanager

# Import our LLM client
from llm_client import LLMToolCaller
from init_ollama import init_ollama

# Configure comprehensive logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('/app/inference_app.log')
    ]
)

# Set external library log levels
logging.getLogger('httpx').setLevel(logging.WARNING)
logging.getLogger('uvicorn.access').setLevel(logging.WARNING)

logger = logging.getLogger(__name__)

# Configuration
DEFAULT_PRESTASHOP_URL = "https://hub.mcp.integration.ambris.com"
HUB_SERVER_URL = "https://hub.mcp.integration.ambris.com/index.php?fc=module&module=ambmcpclient&controller=api"

# Global MCP client and LLM client instances
mcp_client = None
llm_client = None

@asynccontextmanager 
async def lifespan(app: FastAPI):
    """Manage application lifespan - better for HF Spaces"""
    global mcp_client, llm_client
    logger.info("🚀 Starting HF Inference Point application")
    logger.info(f"📦 Environment: {os.getenv('ENVIRONMENT', 'unknown')}")
    logger.info(f"🔗 Default PrestaShop URL: {DEFAULT_PRESTASHOP_URL}")
    
    # Startup
    try:
        # Initialize Ollama first
        logger.info("🤖 Initializing Ollama LLM...")
        if init_ollama():
            llm_client = LLMToolCaller()
            logger.info("✅ LLM client initialized successfully")
        else:
            logger.warning("⚠️ LLM initialization failed, using fallback")
            llm_client = None
            
        logger.info("🔧 Attempting to initialize MCP client...")
        start_time = time.time()
        mcp_client = await get_mcp_client()
        init_time = time.time() - start_time
        logger.info(f"✅ MCP client initialized successfully in {init_time:.2f}s")
        logger.info(f"🛠️ MCP client type: {type(mcp_client)}")
    except Exception as e:
        logger.error(f"❌ Failed to initialize MCP client: {e}")
        logger.error(f"📋 Exception details: {type(e).__name__}: {str(e)}")
        logger.error(f"🔍 MCP client will be None, falling back to graceful degradation")
        mcp_client = None
    
    logger.info("🎯 Application startup complete, ready to serve requests")
    yield  # App runs here
    
    # Shutdown
    logger.info("🛑 Application shutdown initiated")
    try:
        await close_mcp_client()
        logger.info("✅ MCP client closed successfully")
    except Exception as e:
        logger.error(f"⚠️ Error closing MCP client: {e}")
    logger.info("👋 Application shutdown complete")

# Initialize FastAPI with lifespan
app = FastAPI(
    title="PrestaShop MCP Inference Point",
    description="HuggingFace-compatible inference endpoint for PrestaShop MCP tools with LLM-driven tool selection",
    version="2.2.0",
    lifespan=lifespan
)

# Data Models
class ChatMessage(BaseModel):
    """Individual chat message"""
    role: str = Field(..., description="Role: user or assistant")
    content: str = Field(..., description="Message content")

class ChatInput(BaseModel):
    """Input for chat interface"""
    message: str = Field(..., description="User message")
    conversation_history: List[ChatMessage] = Field(default=[], description="Previous conversation")
    client_token: Optional[str] = Field(default=None, description="Optional client authentication token")

class ChatOutput(BaseModel):
    """Output from chat interface"""
    response: str = Field(..., description="Assistant response")
    tool_used: Optional[str] = None
    conversation_history: List[ChatMessage] = Field(..., description="Updated conversation history")

class InferenceInput(BaseModel):
    """Input model for direct MCP tool calls"""
    tool: str = Field(..., description="Name of the MCP tool to call")
    arguments: Dict[str, Any] = Field(..., description="Arguments for the tool")

class InferenceOutput(BaseModel):
    """Output model for direct MCP tool calls"""
    ok: bool
    tool: str
    result: str
    error: Optional[str] = None

# MCP Client Management
async def get_mcp_client():
    """Get or create MCP client connection"""
    global mcp_client
    if mcp_client is not None:
        return mcp_client
    
    try:
        # Use httpx client for HTTP-based MCP communication
        mcp_client = HTTPMCPClient(base_url=HUB_SERVER_URL)
        logger.info("✅ HTTP MCP client created successfully")
        return mcp_client
    except Exception as e:
        logger.error(f"❌ Failed to create MCP client: {e}")
        raise

async def close_mcp_client():
    """Close MCP client connection"""
    global mcp_client
    if mcp_client is not None:
        try:
            if hasattr(mcp_client, 'close'):
                await mcp_client.close()
            mcp_client = None
            logger.info("✅ MCP client closed")
        except Exception as e:
            logger.error(f"⚠️ Error closing MCP client: {e}")

class HTTPMCPClient:
    """Simple HTTP client for MCP communication via hub"""
    
    def __init__(self, base_url: str):
        self.base_url = base_url
        self.client = httpx.AsyncClient(timeout=30.0)
    
    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str:
        """Call an MCP tool via HTTP"""
        try:
            # Map tool name to action
            action_map = {
                "search_products": "search_products",
                "get_product_details": "get_product",
                "get_product_features": "get_product_features", 
                "get_product_images": "get_product_images"
            }
            
            action = action_map.get(tool_name, tool_name)
            
            # Prepare payload for hub
            payload = {
                "action": action,
                "client_token": arguments.get("client_token", "test8f4e2b9c7a1d6e3f5a8b2c4d7e9f1a3b6c8e")
            }
            
            # Add specific parameters based on tool
            if tool_name == "search_products":
                payload["query"] = arguments.get("query", "")
            elif tool_name in ["get_product_details", "get_product_features", "get_product_images"]:
                payload["product_id"] = arguments.get("product_id", 1)
            
            logger.info(f"🌐 Calling hub with: {payload}")
            
            # Make HTTP request to hub
            response = await self.client.post(
                self.base_url,
                json=payload,
                headers={"Content-Type": "application/json"}
            )
            
            if response.status_code == 200:
                try:
                    result = response.json()
                    logger.info(f"✅ Hub response: {result}")
                    return json.dumps(result, indent=2)
                except:
                    return response.text
            else:
                logger.error(f"❌ Hub error {response.status_code}: {response.text}")
                return f"Error: {response.status_code} - {response.text}"
                
        except Exception as e:
            logger.error(f"❌ HTTP MCP call failed: {e}")
            return f"Connection error: {str(e)}"
    
    async def close(self):
        """Close the HTTP client"""
        await self.client.aclose()

# Web Interface
@app.get("/", response_class=HTMLResponse)
async def chat_interface():
    """Serve the chat interface"""
    html_content = """
    <!DOCTYPE html>
    <html>
    <head>
        <title>PrestaShop Assistant</title>
        <style>
            body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
            .chat-container { border: 1px solid #ddd; height: 400px; overflow-y: auto; padding: 10px; margin-bottom: 10px; }
            .message { margin-bottom: 10px; padding: 8px; border-radius: 4px; }
            .user { background-color: #e3f2fd; margin-left: 20%; }
            .assistant { background-color: #f3e5f5; margin-right: 20%; }
            .input-container { display: flex; gap: 10px; }
            input[type="text"] { flex: 1; padding: 8px; }
            button { padding: 8px 16px; background-color: #1976d2; color: white; border: none; border-radius: 4px; cursor: pointer; }
        </style>
    </head>
    <body>
        <h1>🛍️ PrestaShop Assistant</h1>
        <p>Ask me about products! Try: "tell me about product 12" or "search for shoes"</p>
        
        <div id="chatContainer" class="chat-container"></div>
        
        <div class="input-container">
            <input type="text" id="messageInput" placeholder="Ask about products..." />
            <button onclick="sendMessage()">Send</button>
        </div>

        <script>
            let conversationHistory = [];
            
            async function sendMessage() {
                const input = document.getElementById('messageInput');
                const message = input.value.trim();
                if (!message) return;
                
                // Add user message to chat
                addMessageToChat('user', message);
                input.value = '';
                
                try {
                    const response = await fetch('/chat', {
                        method: 'POST',
                        headers: { 'Content-Type': 'application/json' },
                        body: JSON.stringify({
                            message: message,
                            conversation_history: conversationHistory
                        })
                    });
                    
                    const data = await response.json();
                    addMessageToChat('assistant', data.response);
                    conversationHistory = data.conversation_history;
                    
                } catch (error) {
                    addMessageToChat('assistant', 'Sorry, I encountered an error: ' + error.message);
                }
            }
            
            function addMessageToChat(role, content) {
                const container = document.getElementById('chatContainer');
                const messageDiv = document.createElement('div');
                messageDiv.className = `message ${role}`;
                messageDiv.innerHTML = `<strong>${role === 'user' ? 'You' : 'Assistant'}:</strong> ${content}`;
                container.appendChild(messageDiv);
                container.scrollTop = container.scrollHeight;
            }
            
            // Send message on Enter key
            document.getElementById('messageInput').addEventListener('keypress', function(e) {
                if (e.key === 'Enter') sendMessage();
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

@app.post("/chat", response_model=ChatOutput)
async def chat_endpoint(chat_input: ChatInput) -> ChatOutput:
    """
    Chat endpoint for the web interface - uses LLM for natural tool selection
    """
    global mcp_client, llm_client
    logger.info(f"💬 Direct chat request: {chat_input.message[:100]}...")
    
    try:
        # Use LLM to determine which tool to call (if any)
        if llm_client:
            tool_name, arguments = await llm_client.parse_message_for_tool_call(chat_input.message)
        else:
            # Fallback to simple detection if LLM not available
            tool_name, arguments = None, {}
        
        if tool_name and mcp_client:
            # Add client token for authentication
            arguments["client_token"] = chat_input.client_token or "test8f4e2b9c7a1d6e3f5a8b2c4d7e9f1a3b6c8e"
            
            # Call the MCP tool
            logger.info(f"🔧 Calling MCP tool: {tool_name} with {arguments}")
            result = await mcp_client.call_tool(tool_name, arguments)
            
            # Format the response naturally
            if tool_name == "search_products":
                response = f"I searched for '{arguments['query']}' and found: {result}"
            elif tool_name == "get_product_details":
                response = f"Here are the details for product {arguments['product_id']}: {result}"
            elif tool_name == "get_product_features":
                response = f"Here are the features for product {arguments['product_id']}: {result}"
            elif tool_name == "get_product_images":
                response = f"Here are the images for product {arguments['product_id']}: {result}"
            else:
                response = result
                
            tool_used = tool_name
        else:
            # Handle as conversational message
            response = f"I understand you said: '{chat_input.message}'. I'm designed to help with PrestaShop products. You can ask me to search for products, get details about specific products (like 'tell me about product 12'), or ask for features and images."
            tool_used = None
        
        # Update conversation history
        updated_history = list(chat_input.conversation_history)
        updated_history.append(ChatMessage(role="user", content=chat_input.message))
        updated_history.append(ChatMessage(role="assistant", content=response))
        
        return ChatOutput(
            response=response,
            conversation_history=updated_history,
            tool_used=tool_used
        )
        
    except Exception as e:
        logger.error(f"❌ Chat endpoint error: {e}")
        return ChatOutput(
            response=f"Sorry, I encountered an error: {str(e)}",
            conversation_history=list(chat_input.conversation_history)
        )

class EndpointHandler:
    """Simple endpoint handler for MCP tool calls"""
    
    def __init__(self):
        """Initialize the handler - lightweight setup only"""
        pass
    
    async def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Main inference handler - calls MCP tools and returns results
        """
        global mcp_client
        
        try:
            # Validate input
            payload = InferenceInput(**data)
        except Exception as e:
            return {
                "ok": False, 
                "tool": "unknown", 
                "result": f"I'm sorry, but I couldn't understand your request format. Please make sure you're providing a valid 'tool' name and 'arguments' object. The error was: {str(e)}",
                "error": f"Invalid input: {e}"
            }
        
        # Initialize MCP client if not available (fallback)
        if not mcp_client:
            try:
                mcp_client = await get_mcp_client()
            except Exception as e:
                return {
                    "ok": False, 
                    "tool": payload.tool, 
                    "result": f"I apologize, but I'm currently unable to access the PrestaShop MCP tools required for the '{payload.tool}' operation. The MCP server appears to be unavailable at the moment. Please try again later or contact support if this issue persists.",
                    "error": f"MCP server unavailable: {e}"
                }
        
        try:
            # Call the MCP tool
            result = await mcp_client.call_tool(payload.tool, payload.arguments)
            
            return {
                "ok": True,
                "tool": payload.tool,
                "result": result,
                "error": None
            }
            
        except Exception as e:
            error_details = f"Tool '{payload.tool}' failed: {str(e)}"
            logger.error(f"❌ MCP tool call error: {error_details}")
            
            return {
                "ok": False,
                "tool": payload.tool,
                "result": f"I encountered an error while trying to {payload.tool}: {str(e)}. This might be a temporary issue with the PrestaShop system or the specific operation you requested.",
                "error": error_details
            }

# Create handler instance
handler = EndpointHandler()

# API Routes (Programmatic Access)
@app.post("/api/inference", response_model=InferenceOutput)
async def api_inference(data: Dict[str, Any]) -> InferenceOutput:
    """
    API inference endpoint - handles both MCP tool calls and LLM prompt-based requests
    """
    global mcp_client, llm_client
    logger.info(f"🔧 API inference request: {data}")
    
    try:
        # Check if this is a prompt-based request from LLM controller
        if "prompt" in data and "tool" not in data:
            logger.info("📝 Processing LLM prompt-based request")
            
            # Use LLM to convert prompt to tool call
            if llm_client:
                tool_name, arguments = await llm_client.parse_message_for_tool_call(data["prompt"])
            else:
                # Fallback: simple keyword detection
                prompt = data["prompt"].lower()
                if "product" in prompt and any(word in prompt for word in ["details", "about", "information"]):
                    # Try to extract product ID
                    words = data["prompt"].split()
                    product_id = None
                    for word in words:
                        if word.isdigit():
                            product_id = word
                            break
                    if product_id:
                        tool_name, arguments = "get_product_details", {"product_id": product_id}
                    else:
                        tool_name, arguments = "search_products", {"query": data["prompt"]}
                else:
                    tool_name, arguments = "search_products", {"query": data["prompt"]}
            
            # Add client token from LLM request
            if "client_token" in data:
                arguments["client_token"] = data["client_token"]
            else:
                arguments["client_token"] = "test8f4e2b9c7a1d6e3f5a8b2c4d7e9f1a3b6c8e"
            
            # Create the expected payload
            payload = {
                "tool": tool_name,
                "arguments": arguments
            }
            
            logger.info(f"🔄 Converted prompt to tool call: {payload}")
        else:
            # Standard tool-based request
            payload = data
        
        # Process through the handler
        result = await handler(payload)
        return InferenceOutput(**result)
        
    except Exception as e:
        logger.error(f"❌ API inference error: {e}")
        return InferenceOutput(
            ok=False,
            tool="unknown",
            result=f"Error processing request: {str(e)}",
            error=str(e)
        )

# Health check
@app.get("/health")
async def health_check():
    """Health check endpoint"""
    global mcp_client, llm_client
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "mcp_client": "connected" if mcp_client else "disconnected",
        "llm_client": "available" if llm_client else "unavailable",
        "version": "2.2.0"
    }

# Run the application
if __name__ == "__main__":
    port = int(os.getenv("PORT", 7860))
    logger.info(f"🚀 Starting server on port {port}")
    uvicorn.run(
        app, 
        host="0.0.0.0", 
        port=port,
        log_level="info"
    )