import sys
import traceback
import os
import json
import logging
from typing import Any, Dict, List, Optional, Union
from mcp.server.fastmcp import FastMCP
import httpx

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.StreamHandler(sys.stderr)
    ]
)

# Print startup message
logging.info("Starting LMStudio bridge server...")

# ===== Configuration =====
# Load from environment variables with defaults
LMSTUDIO_HOST = os.getenv("LMSTUDIO_HOST", "127.0.0.1")
LMSTUDIO_PORT = os.getenv("LMSTUDIO_PORT", "1234")
LMSTUDIO_API_URL = f"http://{LMSTUDIO_HOST}:{LMSTUDIO_PORT}/v1"
DEBUG = os.getenv("DEBUG", "false").lower() in ("true", "1", "yes")

# Set more verbose logging if debug mode is enabled
if DEBUG:
    logging.getLogger().setLevel(logging.DEBUG)
    logging.debug(f"Debug mode enabled")

logging.info(f"Configured LM Studio API URL: {LMSTUDIO_API_URL}")

# Initialize FastMCP server
mcp = FastMCP("lmstudio-bridge")

# ===== Helper Functions =====
async def call_lmstudio_api(endpoint: str, payload: Dict[str, Any], timeout: float = 60.0) -> Dict[str, Any]:
    """Unified API communication function with better error handling"""
    headers = {
        "Content-Type": "application/json",
        "User-Agent": "claude-lmstudio-bridge/1.0"
    }
    
    url = f"{LMSTUDIO_API_URL}/{endpoint}"
    logging.debug(f"Making request to {url}")
    logging.debug(f"Payload: {json.dumps(payload, indent=2)}")
    
    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(
                url,
                json=payload,
                headers=headers,
                timeout=timeout
            )
            
            # Better error handling with specific error messages
            if response.status_code != 200:
                error_message = f"LM Studio API error: {response.status_code}"
                try:
                    error_json = response.json()
                    if "error" in error_json:
                        if isinstance(error_json["error"], dict) and "message" in error_json["error"]:
                            error_message += f" - {error_json['error']['message']}"
                        else:
                            error_message += f" - {error_json['error']}"
                except:
                    error_message += f" - {response.text[:100]}"
                
                logging.error(f"Error response: {error_message}")
                return {"error": error_message}
            
            result = response.json()
            logging.debug(f"Response received: {json.dumps(result, indent=2, default=str)[:200]}...")
            return result
    except httpx.RequestError as e:
        logging.error(f"Request error: {str(e)}")
        return {"error": f"Connection error: {str(e)}"}
    except Exception as e:
        logging.error(f"Unexpected error: {str(e)}")
        return {"error": f"Unexpected error: {str(e)}"}

def prepare_chat_messages(messages_input: Union[str, List, Dict]) -> List[Dict[str, str]]:
    """Convert various input formats to what LMStudio expects"""
    try:
        # If messages_input is a string
        if isinstance(messages_input, str):
            # Try to parse it as JSON
            try:
                parsed = json.loads(messages_input)
                if isinstance(parsed, list):
                    return parsed
                else:
                    # If it's parsed but not a list, make it a user message
                    return [{"role": "user", "content": messages_input}]
            except json.JSONDecodeError:
                # If not valid JSON, assume it's a simple message
                return [{"role": "user", "content": messages_input}]
        
        # If it's a list already
        elif isinstance(messages_input, list):
            return messages_input
        
        # If it's a dict, assume it's a single message
        elif isinstance(messages_input, dict) and "content" in messages_input:
            if "role" not in messages_input:
                messages_input["role"] = "user"
            return [messages_input]
            
        # If it's some other format, convert to string and make it a user message
        else:
            return [{"role": "user", "content": str(messages_input)}]
    except Exception as e:
        logging.error(f"Error preparing chat messages: {str(e)}")
        # Fallback to simplest format
        return [{"role": "user", "content": str(messages_input)}]

# ===== MCP Tools =====
@mcp.tool()
async def check_lmstudio_connection() -> str:
    """Check if the LM Studio server is running and accessible.
    
    Returns:
        Connection status and model information
    """
    try:
        # Try to get the server status via models endpoint
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{LMSTUDIO_API_URL}/models", timeout=5.0)
            
        if response.status_code == 200:
            models_data = response.json()
            if "data" in models_data and len(models_data["data"]) > 0:
                active_model = models_data["data"][0]["id"]
                return f"✅ Connected to LM Studio. Active model: {active_model}"
            else:
                return "✅ Connected to LM Studio but no models are currently loaded"
        else:
            return f"❌ LM Studio returned an error: {response.status_code}"
    except Exception as e:
        return f"❌ Failed to connect to LM Studio: {str(e)}"

@mcp.tool()
async def list_lmstudio_models() -> str:
    """List available LLM models in LM Studio.
    
    Returns:
        A formatted list of available models with their details.
    """
    logging.info("list_lmstudio_models function called")
    try:
        # Use the API helper function
        models_response = await call_lmstudio_api("models", {}, timeout=10.0)
        
        # Check for errors from the API helper
        if "error" in models_response:
            return f"Error listing models: {models_response['error']}"
        
        if not models_response or "data" not in models_response:
            return "No models found or unexpected response format."
        
        models = models_response["data"]
        model_info = []
        
        for model in models:
            model_info.append(f"ID: {model.get('id', 'Unknown')}")
            model_info.append(f"Name: {model.get('name', 'Unknown')}")
            if model.get('description'):
                model_info.append(f"Description: {model.get('description')}")
            model_info.append("---")
        
        if not model_info:
            return "No models available in LM Studio."
        
        return "\n".join(model_info)
    except Exception as e:
        logging.error(f"Unexpected error in list_lmstudio_models: {str(e)}")
        traceback.print_exc(file=sys.stderr)
        return f"Unexpected error: {str(e)}"

@mcp.tool()
async def generate_text(
    prompt: str,
    model_id: str = "",
    max_tokens: int = 1000,
    temperature: float = 0.7
) -> str:
    """Generate text using a local LLM in LM Studio.
    
    Args:
        prompt: The text prompt to send to the model
        model_id: ID of the model to use (leave empty for default model)
        max_tokens: Maximum number of tokens in the response (default: 1000)
        temperature: Randomness of the output (0-1, default: 0.7)
    
    Returns:
        The generated text from the local LLM
    """
    logging.info("generate_text function called")
    try:
        # Validate inputs
        if not prompt or not prompt.strip():
            return "Error: Prompt cannot be empty."
        
        if max_tokens < 1:
            return "Error: max_tokens must be a positive integer."
        
        if temperature < 0 or temperature > 1:
            return "Error: temperature must be between 0 and 1."
        
        # Prepare payload
        payload = {
            "prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": False
        }
        
        # Add model if specified
        if model_id and model_id.strip():
            payload["model"] = model_id.strip()
        
        # Make request to LM Studio API using the helper function
        response = await call_lmstudio_api("completions", payload)
        
        # Check for errors from the API helper
        if "error" in response:
            return f"Error generating text: {response['error']}"
        
        # Extract and return the generated text
        if "choices" in response and len(response["choices"]) > 0:
            return response["choices"][0].get("text", "")
        
        return "No response generated."
    except Exception as e:
        logging.error(f"Unexpected error in generate_text: {str(e)}")
        traceback.print_exc(file=sys.stderr)
        return f"Unexpected error: {str(e)}"

@mcp.tool()
async def chat_completion(
    messages: str,
    model_id: str = "",
    max_tokens: int = 1000,
    temperature: float = 0.7
) -> str:
    """Generate a chat completion using a local LLM in LM Studio.
    
    Args:
        messages: JSON string of messages in the format [{"role": "user", "content": "Hello"}, ...]
          or a simple text string which will be treated as a user message
        model_id: ID of the model to use (leave empty for default model)
        max_tokens: Maximum number of tokens in the response (default: 1000)
        temperature: Randomness of the output (0-1, default: 0.7)
    
    Returns:
        The generated text from the local LLM
    """
    logging.info("chat_completion function called")
    try:
        # Standardize message format using the helper function
        messages_formatted = prepare_chat_messages(messages)
        
        logging.debug(f"Formatted messages: {json.dumps(messages_formatted, indent=2)}")
        
        # Validate inputs
        if not messages_formatted:
            return "Error: At least one message is required."
        
        if max_tokens < 1:
            return "Error: max_tokens must be a positive integer."
        
        if temperature < 0 or temperature > 1:
            return "Error: temperature must be between 0 and 1."
        
        # Prepare payload
        payload = {
            "messages": messages_formatted,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": False
        }
        
        # Add model if specified
        if model_id and model_id.strip():
            payload["model"] = model_id.strip()
        
        # Make request to LM Studio API using the helper function
        response = await call_lmstudio_api("chat/completions", payload)
        
        # Check for errors from the API helper
        if "error" in response:
            return f"Error generating chat completion: {response['error']}"
        
        # Extract and return the generated text
        if "choices" in response and len(response["choices"]) > 0:
            choice = response["choices"][0]
            if "message" in choice and "content" in choice["message"]:
                return choice["message"]["content"]
        
        return "No response generated."
    except Exception as e:
        logging.error(f"Unexpected error in chat_completion: {str(e)}")
        traceback.print_exc(file=sys.stderr)
        return f"Unexpected error: {str(e)}"

def main():
    """Main entry point for the MCP server."""
    logging.info("Starting server with stdio transport...")
    # Initialize and run the server
    mcp.run(transport='stdio')

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logging.critical(f"CRITICAL ERROR: {str(e)}")
        logging.critical("Traceback:")
        traceback.print_exc(file=sys.stderr)
        sys.exit(1)