"""
IBM watsonx.ai MCP Server
Official IBM SDK Integration for Model Context Protocol

Based on IBM watsonx.ai Python SDK documentation:
https://ibm.github.io/watsonx-ai-python-sdk/

Usage:
    python ibm_watsonx_mcp_server.py

Requirements:
    - IBM Cloud API key
    - watsonx.ai project ID
    - .env file with credentials
"""

import json
import logging
import os
from typing import Dict, Any
from dotenv import load_dotenv

# MCP imports (using FastMCP - the correct current API)
from mcp.server.fastmcp import FastMCP

# IBM watsonx.ai imports (official SDK)
from ibm_watsonx_ai import APIClient, Credentials
from ibm_watsonx_ai.foundation_models import ModelInference
from ibm_watsonx_ai.foundation_models.utils.enums import ModelTypes
from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams

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

# Load environment variables
load_dotenv()

# IBM watsonx.ai configuration
WATSONX_APIKEY = os.getenv("IBM_API_KEY")
WATSONX_URL = os.getenv("IBM_URL")
WATSONX_PROJECT_ID = os.getenv("IBM_PROJECT_ID")

def validate_credentials():
    """Validate IBM watsonx.ai credentials"""
    required_vars = [
        ("WATSONX_APIKEY", WATSONX_APIKEY),
        ("WATSONX_URL", WATSONX_URL), 
        ("WATSONX_PROJECT_ID", WATSONX_PROJECT_ID)
    ]
    
    missing_vars = [name for name, value in required_vars if not value]
    if missing_vars:
        raise ValueError(
            f"Missing required IBM credentials: {', '.join(missing_vars)}\n"
            "Please set these in your .env file:\n"
            "WATSONX_APIKEY=your-ibm-cloud-api-key\n"
            "WATSONX_URL=https://us-south.ml.cloud.ibm.com\n"
            "WATSONX_PROJECT_ID=your-project-id"
        )
    
    logger.info("✅ IBM credentials validated successfully")

# Validate credentials on startup
validate_credentials()

# Initialize IBM API client
credentials = Credentials(
    api_key=WATSONX_APIKEY,
    url=WATSONX_URL
)
api_client = APIClient(credentials)

# Available IBM models for the team
AVAILABLE_MODELS = {
    "granite-13b-instruct": ModelTypes.GRANITE_13B_INSTRUCT_V2,
    "granite-13b-chat": ModelTypes.GRANITE_13B_CHAT_V2,
    "granite-20b-multilingual": ModelTypes.GRANITE_20B_MULTILINGUAL,
    "llama-3-8b-instruct": ModelTypes.LLAMA_3_8B_INSTRUCT,
    "llama-3-70b-instruct": ModelTypes.LLAMA_3_70B_INSTRUCT
}

# Default model for team usage
DEFAULT_MODEL = ModelTypes.LLAMA_3_70B_INSTRUCT

# Create the FastMCP server
mcp = FastMCP("IBM watsonx.ai MCP Server")

logger.info("IBM watsonx.ai MCP Server initialized successfully")

# === TOOLS ===

@mcp.tool()
def generate_text(
    prompt: str, 
    model_name: str = "granite-13b-instruct",
    max_tokens: int = 200,
    temperature: float = 0.7
) -> Dict[str, Any]:
    """
    Generate text using IBM watsonx.ai foundation models (Granite, Llama, Mistral)
    
    Args:
        prompt: Text prompt for generation
        model_name: IBM model to use (granite-13b-instruct, granite-13b-chat, llama-3-8b-instruct, etc.)
        max_tokens: Maximum tokens to generate (1-4000)
        temperature: Sampling temperature (0.0-2.0)
    
    Returns:
        Dictionary with generated text and metadata
    """
    try:
        # Get model ID
        model_id = AVAILABLE_MODELS.get(model_name, DEFAULT_MODEL)
        
        # Configure generation parameters using IBM SDK
        generation_params = {
            GenParams.MAX_NEW_TOKENS: max_tokens,
            GenParams.TEMPERATURE: temperature,
            GenParams.TOP_P: 0.9,
            GenParams.TOP_K: 50
        }
        
        # Initialize IBM model inference
        model = ModelInference(
            model_id=model_id,
            params=generation_params,
            credentials=credentials,
            project_id=WATSONX_PROJECT_ID
        )
        
        # Generate response using IBM watsonx.ai
        response = model.generate_text(prompt=prompt)
        
        # Return structured response
        return {
            "generated_text": response,
            "model_used": model_name,
            "model_id": model_id,
            "prompt_tokens": len(prompt.split()),  # Approximate
            "completion_tokens": len(response.split()) if response else 0,
            "status": "success"
        }
        
    except Exception as e:
        logger.error(f"❌ IBM watsonx.ai generation error: {e}")
        return {
            "error": str(e),
            "model_used": model_name,
            "status": "error"
        }

@mcp.tool()
def analyze_sentiment(text: str) -> Dict[str, Any]:
    """
    Analyze sentiment of text using IBM Granite models
    
    Args:
        text: Text to analyze for sentiment
    
    Returns:
        Dictionary with sentiment analysis results
    """
    prompt = f"""
        Analyze the sentiment of the following text and provide a detailed analysis:

        Text: "{text}"

        Please provide:
        1. Overall sentiment (positive, negative, neutral)
        2. Confidence score (0-1)
        3. Key phrases that influenced the sentiment
        4. Brief explanation

        Format your response as structured analysis.
        """
    
    result = generate_text(
        prompt=prompt,
        model_name="granite-13b-instruct",
        max_tokens=300,
        temperature=0.3
    )
    
    if result["status"] == "success":
        return {
            "text_analyzed": text,
            "sentiment_analysis": result["generated_text"],
            "model_used": "IBM Granite 13B Instruct",
            "status": "success"
        }
    else:
        return result

@mcp.tool()
def summarize_text(text: str, max_length: int = 150) -> Dict[str, Any]:
    """
    Summarize text using IBM foundation models
    
    Args:
        text: Text to summarize
        max_length: Maximum length of summary in words (50-500)
    
    Returns:
        Dictionary with summary and metadata
    """
    prompt = f"""
        Summarize the following text in approximately {max_length} words or less:

        Text: {text}

        Summary:
        """
    
    result = generate_text(
        prompt=prompt,
        model_name="granite-13b-instruct", 
        max_tokens=max_length + 50,
        temperature=0.5
    )
    
    if result["status"] == "success":
        return {
            "original_text": text[:200] + "..." if len(text) > 200 else text,
            "summary": result["generated_text"],
            "model_used": "IBM Granite 13B Instruct",
            "requested_length": max_length,
            "status": "success"
        }
    else:
        return result

@mcp.tool()
def get_model_info(model_name: str = "granite-13b-instruct") -> Dict[str, Any]:
    """
    Get detailed information about a specific IBM watsonx.ai model
    
    Args:
        model_name: Name of the model to get information about
    
    Returns:
        Dictionary with model specifications and capabilities
    """
    try:
        if model_name not in AVAILABLE_MODELS:
            return {
                "error": f"Model '{model_name}' not available",
                "available_models": list(AVAILABLE_MODELS.keys()),
                "status": "error"
            }
        
        model_id = AVAILABLE_MODELS[model_name]
        
        # Get model specifications from IBM API
        specs = api_client.foundation_models.get_model_specs(model_id)
        
        return {
            "model_name": model_name,
            "model_id": model_id,
            "description": specs.get("short_description", "IBM Foundation Model"),
            "supported_tasks": specs.get("supported_tasks", []),
            "model_limits": specs.get("model_limits", {}),
            "provider": "IBM",
            "status": "success"
        }
        
    except Exception as e:
        logger.error(f"❌ Error getting model info for {model_name}: {e}")
        return {
            "error": str(e),
            "model_name": model_name,
            "status": "error"
        }

# === RESOURCES ===

@mcp.resource("ibm://models/available")
def get_available_models() -> str:
    """List of available IBM watsonx.ai foundation models for the team"""
    try:
        # Use IBM API to get model specs
        model_info = {}
        for model_name, model_type in AVAILABLE_MODELS.items():
            try:
                specs = api_client.foundation_models.get_model_specs(model_type)
                model_info[model_name] = {
                    "model_id": model_type,
                    "description": specs.get("short_description", "IBM Foundation Model"),
                    "supported_tasks": specs.get("supported_tasks", []),
                    "model_limits": specs.get("model_limits", {}),
                    "provider": "IBM"
                }
            except Exception as e:
                logger.warning(f"⚠️  Could not get specs for {model_name}: {e}")
                model_info[model_name] = {
                    "model_id": model_type,
                    "description": "IBM Foundation Model",
                    "status": "available",
                    "provider": "IBM"
                }
        
        return json.dumps(model_info, indent=2)
        
    except Exception as e:
        logger.error(f"❌ Error fetching IBM model information: {e}")
        return json.dumps({"error": f"Failed to fetch model info: {str(e)}"}, indent=2)

@mcp.resource("ibm://config/server-info")  
def get_server_info() -> str:
    """IBM watsonx.ai MCP server configuration and status"""
    config_info = {
        "server_name": "IBM watsonx.ai MCP Server",
        "version": "1.0.0",
        "watsonx_url": WATSONX_URL,
        "project_id": WATSONX_PROJECT_ID[:8] + "..." if WATSONX_PROJECT_ID else "Not set",
        "available_models": list(AVAILABLE_MODELS.keys()),
        "default_model": "granite-13b-instruct",
        "total_models": len(AVAILABLE_MODELS),
        "provider": "IBM watsonx.ai",
        "status": "operational",
        "features": [
            "Text Generation",
            "Sentiment Analysis", 
            "Text Summarization",
            "Model Information"
        ]
    }
    return json.dumps(config_info, indent=2)

@mcp.resource("ibm://examples/prompts")
def get_example_prompts() -> str:
    """Example prompts optimized for IBM Granite models"""
    examples = {
        "text_generation": [
            "Explain the concept of artificial intelligence in simple terms",
            "Write a professional email requesting a meeting",
            "Create a summary of quarterly business results",
            "Generate code comments for a Python function"
        ],
        "sentiment_analysis": [
            "Our team is excited about the new IBM watsonx.ai features!",
            "The project deadline is too tight and causing stress",
            "Neutral statement about weather conditions today",
            "Mixed feelings about the recent company reorganization"
        ],
        "summarization": [
            "Long technical documentation that needs condensing",
            "Meeting notes that require executive summary",
            "Research paper abstract creation",
            "Customer feedback compilation"
        ],
        "enterprise_use_cases": [
            "Generate a risk assessment for cloud migration",
            "Create talking points for investor presentation",
            "Draft a compliance report summary",
            "Analyze customer satisfaction survey results"
        ]
    }
    return json.dumps(examples, indent=2)

# Run the server
if __name__ == "__main__":
    logger.info("🚀 Starting IBM watsonx.ai MCP Server...")
    mcp.run()