mcp_server.py

"""
╔══════════════════════════════════════════════════════════════════╗
║          NVIDIA AI — MCP SERVER STANDALONE                      ║
║          Model Context Protocol 2024-11-05                      ║
║                                                                  ║
║  Transport supportati:                                           ║
║    • HTTP  → POST /v1/mcp  (JSON-RPC 2.0)                       ║
║    • SSE   → GET  /v1/mcp/sse  (Server-Sent Events)             ║
║    • STDIO → python mcp_server.py --stdio  (Claude Desktop)     ║
║                                                                  ║
║  Tools:                                                          ║
║    • nvidia_chat           Chat con memoria                      ║
║    • nvidia_chat_stream    Info streaming SSE                    ║
║    • list_models           Catalogo modelli                      ║
║    • get_model_info        Info singolo modello                  ║
║    • memory_stats          Statistiche sessione                  ║
║    • memory_clear          Pulisci memoria                       ║
║    • memory_save_summary   Salva riassunto                       ║
║    • get_history           Cronologia conversazione              ║
║    • list_sessions         Lista sessioni attive                 ║
║    • delete_session        Elimina sessione                      ║
║    • summarize_session     Riassumi sessione con AI              ║
║    • multi_model_compare   Confronta risposte più modelli        ║
╚══════════════════════════════════════════════════════════════════╝
"""

from __future__ import annotations

import os
import sys
import json
import uuid
import time
import base64
import asyncio
import logging
import argparse
from typing import Any, AsyncGenerator, Optional

import httpx
import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from sse_starlette.sse import EventSourceResponse

# Importa il memory manager locale
try:
    from memory_manager import MemoryManager
except ImportError:
    # Se non trovato, definisce una versione minimale inline
    class MemoryManager:  # type: ignore
        def __init__(self, **kw):
            self.short_memory: dict[str, list[dict]] = {}

        def add_short(self, sid, role, content):
            self.short_memory.setdefault(sid, []).append(
                {"role": role, "content": content, "timestamp": time.time()}
            )

        def get_short(self, sid):
            return [{"role": m["role"], "content": m["content"]}
                    for m in self.short_memory.get(sid, [])]

        def get_full_context(self, sid, system=""):
            msgs = []
            if system:
                msgs.append({"role": "system", "content": system})
            msgs.extend(self.get_short(sid))
            return msgs

        def get_long_context(self, sid, **kw):
            return ""

        def get_stats(self, sid):
            short = self.short_memory.get(sid, [])
            return {"session_id": sid,
                    "short_memory": {"messages": len(short)},
                    "long_memory": {"messages": 0},
                    "total_messages": len(short)}

        def clear_short(self, sid):
            self.short_memory[sid] = []

        def delete_session(self, sid):
            self.short_memory.pop(sid, None)

        def save_session_summary(self, sid, summary):
            pass

        def list_sessions(self):
            return list(self.short_memory.keys())

load_dotenv()

# ══════════════════════════════════════════════════════════════
# LOGGING
# ══════════════════════════════════════════════════════════════

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [MCP] %(levelname)s %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("mcp_server")


# ══════════════════════════════════════════════════════════════
# CONFIGURAZIONE
# ══════════════════════════════════════════════════════════════

def _decode_key() -> str:
    b64 = os.getenv(
        "NVIDIA_API_KEY_B64",
        "bnZhcGktUXdVV29UMlJPSlR4X2luTWtraWVSeF9sTl9TdkxRM2FiNnBVYTFmREg4c1hOYUFST0tyWWhDM3NWRU5MYU03dg=="
    )
    try:
        return base64.b64decode(b64.encode()).decode("utf-8")
    except Exception:
        return b64


NVIDIA_API_KEY: str  = os.getenv("NVIDIA_API_KEY") or _decode_key()
NVIDIA_BASE_URL: str = "https://integrate.api.nvidia.com/v1"
DEFAULT_MODEL: str   = os.getenv("DEFAULT_MODEL", "meta/llama-3.1-405b-instruct")
MCP_PORT: int        = int(os.getenv("MCP_PORT", "8765"))
MCP_HOST: str        = os.getenv("MCP_HOST", "0.0.0.0")

memory = MemoryManager(
    max_short=int(os.getenv("MAX_SHORT_MEMORY", "20")),
    max_long_tokens=int(os.getenv("MAX_LONG_MEMORY_TOKENS", "50000")),
)


# ══════════════════════════════════════════════════════════════
# CATALOGO MODELLI
# ══════════════════════════════════════════════════════════════

MODELS: dict[str, dict] = {
    # ── LLaMA ──────────────────────────────────────────────────
    "meta/llama-3.3-70b-instruct": {
        "name": "LLaMA 3.3 70B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "Meta LLaMA 3.3 – 128K context",
    },
    "meta/llama-3.1-405b-instruct": {
        "name": "LLaMA 3.1 405B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "Largest open-source LLM",
    },
    "meta/llama-3.1-70b-instruct": {
        "name": "LLaMA 3.1 70B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "High-quality 70B – 128K",
    },
    "meta/llama-3.1-8b-instruct": {
        "name": "LLaMA 3.1 8B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "Fast and efficient 8B",
    },
    "meta/llama-3.2-3b-instruct": {
        "name": "LLaMA 3.2 3B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "Compact 3B model",
    },
    "meta/llama-3.2-1b-instruct": {
        "name": "LLaMA 3.2 1B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "Ultra-lightweight 1B",
    },
    # ── DeepSeek ───────────────────────────────────────────────
    "deepseek-ai/deepseek-r1": {
        "name": "DeepSeek R1",
        "category": "Reasoning",
        "context": 65536,
        "description": "Advanced reasoning model",
    },
    "deepseek-ai/deepseek-r1-distill-llama-70b": {
        "name": "DeepSeek R1 Distill LLaMA 70B",
        "category": "Reasoning",
        "context": 131072,
        "description": "R1 reasoning distilled into LLaMA 70B",
    },
    # ── Qwen ───────────────────────────────────────────────────
    "qwen/qwen2.5-72b-instruct": {
        "name": "Qwen 2.5 72B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "Alibaba Qwen 72B",
    },
    "qwen/qwen2.5-7b-instruct": {
        "name": "Qwen 2.5 7B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "Efficient Qwen 7B",
    },
    "qwen/qwq-32b": {
        "name": "QwQ 32B",
        "category": "Reasoning",
        "context": 131072,
        "description": "Qwen reasoning 32B",
    },
    "qwen/qwen2.5-coder-32b-instruct": {
        "name": "Qwen 2.5 Coder 32B",
        "category": "Code",
        "context": 131072,
        "description": "Specialized coding model",
    },
    # ── Mistral ────────────────────────────────────────────────
    "mistralai/mistral-large-2-instruct": {
        "name": "Mistral Large 2",
        "category": "LLM",
        "context": 131072,
        "description": "Mistral flagship",
    },
    "mistralai/mixtral-8x22b-instruct-v0.1": {
        "name": "Mixtral 8x22B",
        "category": "LLM",
        "context": 65536,
        "description": "MoE 8 experts × 22B",
    },
    "mistralai/mixtral-8x7b-instruct-v0.1": {
        "name": "Mixtral 8x7B",
        "category": "LLM",
        "context": 32768,
        "description": "Efficient MoE",
    },
    "mistralai/mistral-7b-instruct-v0.3": {
        "name": "Mistral 7B v0.3",
        "category": "LLM",
        "context": 32768,
        "description": "Compact 7B",
    },
    # ── Google ─────────────────────────────────────────────────
    "google/gemma-2-27b-it": {
        "name": "Gemma 2 27B IT",
        "category": "LLM",
        "context": 8192,
        "description": "Google Gemma 2 instruction-tuned",
    },
    "google/gemma-2-9b-it": {
        "name": "Gemma 2 9B IT",
        "category": "LLM",
        "context": 8192,
        "description": "Efficient Gemma 9B",
    },
    # ── Microsoft ──────────────────────────────────────────────
    "microsoft/phi-3.5-mini-instruct": {
        "name": "Phi 3.5 Mini",
        "category": "LLM",
        "context": 131072,
        "description": "Small but powerful",
    },
    "microsoft/phi-3-medium-128k-instruct": {
        "name": "Phi 3 Medium 128K",
        "category": "LLM",
        "context": 131072,
        "description": "Medium Phi with 128K",
    },
    # ── NVIDIA ─────────────────────────────────────────────────
    "nvidia/llama-3.1-nemotron-70b-instruct": {
        "name": "Nemotron 70B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "NVIDIA fine-tuned LLaMA 70B",
    },
    "nvidia/llama-3.1-nemotron-70b-reward": {
        "name": "Nemotron 70B Reward",
        "category": "Reward",
        "context": 131072,
        "description": "Reward model for RLHF",
    },
    "nvidia/nemotron-mini-4b-instruct": {
        "name": "Nemotron Mini 4B",
        "category": "LLM",
        "context": 8192,
        "description": "Ultra-efficient 4B by NVIDIA",
    },
    # ── Vision ─────────────────────────────────────────────────
    "meta/llama-3.2-90b-vision-instruct": {
        "name": "LLaMA 3.2 90B Vision",
        "category": "Vision",
        "context": 131072,
        "description": "Vision + Language 90B",
    },
    "meta/llama-3.2-11b-vision-instruct": {
        "name": "LLaMA 3.2 11B Vision",
        "category": "Vision",
        "context": 131072,
        "description": "Efficient vision-language",
    },
    # ── Code ───────────────────────────────────────────────────
    "meta/codellama-70b": {
        "name": "Code LLaMA 70B",
        "category": "Code",
        "context": 16384,
        "description": "Meta code specialist",
    },
    "ibm/granite-34b-code-instruct": {
        "name": "Granite 34B Code",
        "category": "Code",
        "context": 8192,
        "description": "IBM code model",
    },
    # ── Safety ─────────────────────────────────────────────────
    "meta/llama-guard-3-8b": {
        "name": "LLaMA Guard 3 8B",
        "category": "Safety",
        "context": 8192,
        "description": "Content safety classifier",
    },
    # ── Finance ────────────────────────────────────────────────
    "writer/palmyra-fin-70b-32k": {
        "name": "Palmyra Fin 70B",
        "category": "Finance",
        "context": 32768,
        "description": "Financial domain specialist",
    },
    # ── Embedding ──────────────────────────────────────────────
    "nvidia/nv-embed-v2": {
        "name": "NV Embed V2",
        "category": "Embedding",
        "context": 32768,
        "description": "NVIDIA embedding model",
    },
}


# ══════════════════════════════════════════════════════════════
# NVIDIA NIM CLIENT
# ══════════════════════════════════════════════════════════════

async def nvidia_chat(
    messages: list[dict],
    model: str = DEFAULT_MODEL,
    temperature: float = 0.7,
    max_tokens: int = 4096,
    top_p: float = 0.9,
) -> str:
    """Chiamata non-streaming a NVIDIA NIM."""
    headers = {
        "Authorization": f"Bearer {NVIDIA_API_KEY}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "top_p": top_p,
        "stream": False,
    }
    async with httpx.AsyncClient(timeout=120.0) as client:
        resp = await client.post(
            f"{NVIDIA_BASE_URL}/chat/completions",
            headers=headers,
            json=payload,
        )
        if resp.status_code != 200:
            raise RuntimeError(
                f"NVIDIA API {resp.status_code}: {resp.text[:300]}"
            )
        return resp.json()["choices"][0]["message"]["content"]


async def nvidia_chat_stream(
    messages: list[dict],
    model: str = DEFAULT_MODEL,
    temperature: float = 0.7,
    max_tokens: int = 4096,
    top_p: float = 0.9,
) -> AsyncGenerator[str, None]:
    """Streaming chunk-by-chunk da NVIDIA NIM."""
    headers = {
        "Authorization": f"Bearer {NVIDIA_API_KEY}",
        "Content-Type": "application/json",
        "Accept": "text/event-stream",
    }
    payload = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "top_p": top_p,
        "stream": True,
    }
    async with httpx.AsyncClient(timeout=120.0) as client:
        async with client.stream(
            "POST",
            f"{NVIDIA_BASE_URL}/chat/completions",
            headers=headers,
            json=payload,
        ) as resp:
            if resp.status_code != 200:
                body = await resp.aread()
                raise RuntimeError(
                    f"NVIDIA API {resp.status_code}: {body.decode()[:300]}"
                )
            async for line in resp.aiter_lines():
                if line.startswith("data: "):
                    chunk_str = line[6:]
                    if chunk_str.strip() == "[DONE]":
                        break
                    try:
                        chunk = json.loads(chunk_str)
                        delta = chunk["choices"][0].get("delta", {})
                        if "content" in delta:
                            yield delta["content"]
                    except (json.JSONDecodeError, KeyError, IndexError):
                        continue


# ══════════════════════════════════════════════════════════════
# DEFINIZIONE TOOLS MCP
# ══════════════════════════════════════════════════════════════

MCP_TOOLS: list[dict] = [
    {
        "name": "nvidia_chat",
        "description": (
            "Chatta con qualsiasi modello NVIDIA NIM. "
            "Supporta memoria breve (RAM) e lunga (SQLite) per sessioni persistenti."
        ),
        "inputSchema": {
            "type": "object",
            "required": ["message"],
            "properties": {
                "message": {
                    "type": "string",
                    "description": "Messaggio dell'utente"
                },
                "model": {
                    "type": "string",
                    "default": DEFAULT_MODEL,
                    "description": "ID modello NVIDIA (es. meta/llama-3.1-405b-instruct)"
                },
                "session_id": {
                    "type": "string",
                    "description": "ID sessione per la memoria. Auto-generato se omesso."
                },
                "system_prompt": {
                    "type": "string",
                    "default": "You are a helpful AI assistant.",
                    "description": "System prompt da usare nella conversazione"
                },
                "temperature": {
                    "type": "number",
                    "default": 0.7,
                    "minimum": 0.0,
                    "maximum": 2.0,
                    "description": "Creatività della risposta"
                },
                "max_tokens": {
                    "type": "integer",
                    "default": 4096,
                    "minimum": 1,
                    "maximum": 16384,
                    "description": "Numero massimo di token nella risposta"
                },
                "top_p": {
                    "type": "number",
                    "default": 0.9,
                    "minimum": 0.0,
                    "maximum": 1.0,
                    "description": "Nucleus sampling"
                },
                "use_memory": {
                    "type": "boolean",
                    "default": True,
                    "description": "Se True usa il sistema di memoria breve/lunga"
                },
            },
        },
    },
    {
        "name": "nvidia_chat_stream",
        "description": (
            "Restituisce le informazioni per connettersi all'endpoint SSE streaming "
            "di NVIDIA NIM. Il client deve aprire una connessione SSE all'URL fornito."
        ),
        "inputSchema": {
            "type": "object",
            "required": ["message"],
            "properties": {
                "message": {"type": "string"},
                "model": {"type": "string", "default": DEFAULT_MODEL},
                "session_id": {"type": "string"},
                "temperature": {"type": "number", "default": 0.7},
                "max_tokens": {"type": "integer", "default": 4096},
            },
        },
    },
    {
        "name": "list_models",
        "description": "Lista tutti i modelli NVIDIA NIM disponibili, filtrabili per categoria.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "category": {
                    "type": "string",
                    "description": (
                        "Categoria filtro: LLM | Code | Vision | "
                        "Reasoning | Embedding | Safety | Finance | Reward"
                    ),
                },
                "search": {
                    "type": "string",
                    "description": "Cerca per nome o descrizione (case-insensitive)"
                },
            },
        },
    },
    {
        "name": "get_model_info",
        "description": "Restituisce informazioni dettagliate su un singolo modello NVIDIA NIM.",
        "inputSchema": {
            "type": "object",
            "required": ["model_id"],
            "properties": {
                "model_id": {
                    "type": "string",
                    "description": "ID completo del modello (es. meta/llama-3.1-405b-instruct)"
                }
            },
        },
    },
    {
        "name": "memory_stats",
        "description": "Statistiche di memoria breve e lunga per una sessione.",
        "inputSchema": {
            "type": "object",
            "required": ["session_id"],
            "properties": {
                "session_id": {"type": "string"}
            },
        },
    },
    {
        "name": "memory_clear",
        "description": "Pulisce la memoria breve o tutta la memoria di una sessione.",
        "inputSchema": {
            "type": "object",
            "required": ["session_id"],
            "properties": {
                "session_id": {"type": "string"},
                "clear_type": {
                    "type": "string",
                    "enum": ["short", "all"],
                    "default": "short",
                    "description": "short = solo memoria breve; all = elimina tutta la sessione",
                },
            },
        },
    },
    {
        "name": "memory_save_summary",
        "description": "Salva un riassunto testuale persistente per una sessione.",
        "inputSchema": {
            "type": "object",
            "required": ["session_id", "summary"],
            "properties": {
                "session_id": {"type": "string"},
                "summary": {
                    "type": "string",
                    "description": "Testo del riassunto da salvare"
                },
            },
        },
    },
    {
        "name": "get_history",
        "description": (
            "Restituisce la cronologia completa della conversazione: "
            "memoria breve (messaggi recenti) + contesto dalla memoria lunga."
        ),
        "inputSchema": {
            "type": "object",
            "required": ["session_id"],
            "properties": {
                "session_id": {"type": "string"},
                "include_long": {
                    "type": "boolean",
                    "default": True,
                    "description": "Includi contesto dalla memoria lunga"
                },
            },
        },
    },
    {
        "name": "list_sessions",
        "description": "Lista tutte le sessioni attive (con memoria breve o lunga).",
        "inputSchema": {
            "type": "object",
            "properties": {},
        },
    },
    {
        "name": "delete_session",
        "description": "Elimina completamente una sessione (memoria breve + lunga).",
        "inputSchema": {
            "type": "object",
            "required": ["session_id"],
            "properties": {
                "session_id": {"type": "string"}
            },
        },
    },
    {
        "name": "summarize_session",
        "description": (
            "Usa un modello AI per generare automaticamente un riassunto "
            "della conversazione nella sessione e lo salva in memoria lunga."
        ),
        "inputSchema": {
            "type": "object",
            "required": ["session_id"],
            "properties": {
                "session_id": {"type": "string"},
                "model": {
                    "type": "string",
                    "default": "meta/llama-3.1-8b-instruct",
                    "description": "Modello da usare per il riassunto"
                },
                "language": {
                    "type": "string",
                    "default": "italiano",
                    "description": "Lingua del riassunto"
                },
            },
        },
    },
    {
        "name": "multi_model_compare",
        "description": (
            "Invia lo stesso messaggio a più modelli NVIDIA in parallelo "
            "e restituisce tutte le risposte per confronto."
        ),
        "inputSchema": {
            "type": "object",
            "required": ["message", "models"],
            "properties": {
                "message": {"type": "string"},
                "models": {
                    "type": "array",
                    "items": {"type": "string"},
                    "minItems": 2,
                    "maxItems": 6,
                    "description": "Lista di 2-6 model IDs da confrontare",
                },
                "system_prompt": {"type": "string", "default": "You are a helpful assistant."},
                "temperature": {"type": "number", "default": 0.7},
                "max_tokens": {"type": "integer", "default": 1024},
            },
        },
    },
]


# ══════════════════════════════════════════════════════════════
# ESECUZIONE TOOLS
# ══════════════════════════════════════════════════════════════

async def _tool_nvidia_chat(args: dict) -> dict:
    session_id  = args.get("session_id") or str(uuid.uuid4())
    message     = args["message"]
    model       = args.get("model", DEFAULT_MODEL)
    system      = args.get("system_prompt", "You are a helpful AI assistant.")
    temperature = float(args.get("temperature", 0.7))
    max_tokens  = int(args.get("max_tokens", 4096))
    top_p       = float(args.get("top_p", 0.9))
    use_memory  = bool(args.get("use_memory", True))

    if use_memory:
        memory.add_short(session_id, "user", message)
        messages = memory.get_full_context(session_id, system)
    else:
        messages = [
            {"role": "system", "content": system},
            {"role": "user",   "content": message},
        ]

    log.info("nvidia_chat  model=%s  session=%s", model, session_id)
    response = await nvidia_chat(
        messages=messages,
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
        top_p=top_p,
    )

    if use_memory:
        memory.add_short(session_id, "assistant", response)

    return {
        "response":     response,
        "model":        model,
        "session_id":   session_id,
        "memory_stats": memory.get_stats(session_id),
        "timestamp":    time.time(),
    }


async def _tool_nvidia_chat_stream(args: dict) -> dict:
    session_id  = args.get("session_id") or str(uuid.uuid4())
    message     = args["message"]
    model       = args.get("model", DEFAULT_MODEL)
    temperature = float(args.get("temperature", 0.7))
    max_tokens  = int(args.get("max_tokens", 4096))

    # Costruiamo i parametri query per l'URL SSE
    params = (
        f"message={message}&model={model}"
        f"&session_id={session_id}"
        f"&temperature={temperature}"
        f"&max_tokens={max_tokens}"
    )
    return {
        "info": "Connettiti all'endpoint SSE per ricevere la risposta in streaming.",
        "sse_endpoint": f"GET /v1/mcp/sse/chat?{params}",
        "events": {
            "start": "Inizio generazione (session_id, model, timestamp)",
            "token": "Singolo token generato {'token': '...'}",
            "done":  "Fine risposta {'full_response': '...', 'memory_stats': {...}}",
            "error": "Errore {'error': '...'}",
        },
        "session_id": session_id,
    }


async def _tool_list_models(args: dict) -> dict:
    category = args.get("category", "").lower()
    search   = args.get("search",   "").lower()

    result = {}
    for mid, info in MODELS.items():
        if category and info["category"].lower() != category:
            continue
        if search and search not in info["name"].lower() \
                  and search not in info["description"].lower() \
                  and search not in mid.lower():
            continue
        result[mid] = info

    categories = sorted({m["category"] for m in MODELS.values()})
    return {
        "models":     result,
        "total":      len(result),
        "categories": categories,
    }


async def _tool_get_model_info(args: dict) -> dict:
    mid = args["model_id"]
    if mid not in MODELS:
        raise ValueError(f"Modello '{mid}' non trovato nel catalogo.")
    return {"model_id": mid, **MODELS[mid]}


async def _tool_memory_stats(args: dict) -> dict:
    return memory.get_stats(args["session_id"])


async def _tool_memory_clear(args: dict) -> dict:
    sid        = args["session_id"]
    clear_type = args.get("clear_type", "short")
    if clear_type == "all":
        memory.delete_session(sid)
        msg = "Tutta la memoria della sessione eliminata."
    else:
        memory.clear_short(sid)
        msg = "Memoria breve pulita (messaggi spostati in memoria lunga)."
    return {"status": "ok", "message": msg, "session_id": sid}


async def _tool_memory_save_summary(args: dict) -> dict:
    sid     = args["session_id"]
    summary = args["summary"]
    memory.save_session_summary(sid, summary)
    return {"status": "ok", "session_id": sid, "summary_saved": summary[:100] + "..."}


async def _tool_get_history(args: dict) -> dict:
    sid          = args["session_id"]
    include_long = bool(args.get("include_long", True))
    result: dict = {
        "session_id":   sid,
        "short_memory": memory.get_short(sid),
        "stats":        memory.get_stats(sid),
    }
    if include_long:
        result["long_context"] = memory.get_long_context(sid)
    return result


async def _tool_list_sessions(args: dict) -> dict:  # noqa: ARG001
    sessions = memory.list_sessions()
    return {
        "sessions": sessions,
        "total":    len(sessions),
    }


async def _tool_delete_session(args: dict) -> dict:
    sid = args["session_id"]
    memory.delete_session(sid)
    return {"status": "deleted", "session_id": sid}


async def _tool_summarize_session(args: dict) -> dict:
    sid      = args["session_id"]
    model    = args.get("model", "meta/llama-3.1-8b-instruct")
    language = args.get("language", "italiano")

    short  = memory.get_short(sid)
    long_c = memory.get_long_context(sid)

    if not short and not long_c:
        return {"status": "no_data", "session_id": sid,
                "message": "Nessun dato da riassumere per questa sessione."}

    # Prepara il testo da riassumere
    history_text = ""
    if long_c:
        history_text += f"[Contesto precedente]\n{long_c}\n\n"
    if short:
        history_text += "[Conversazione recente]\n"
        for m in short:
            history_text += f"{m['role'].upper()}: {m['content']}\n"

    messages = [
        {
            "role": "system",
            "content": (
                f"Sei un assistente specializzato nel riassumere conversazioni. "
                f"Rispondi sempre in {language}."
            ),
        },
        {
            "role": "user",
            "content": (
                f"Riassumi la seguente conversazione in modo conciso, "
                f"evidenziando i punti chiave, le decisioni prese e i temi principali.\n\n"
                f"{history_text}"
            ),
        },
    ]

    log.info("summarize_session  model=%s  session=%s", model, sid)
    summary = await nvidia_chat(
        messages=messages,
        model=model,
        temperature=0.3,
        max_tokens=1024,
    )

    memory.save_session_summary(sid, summary)
    return {
        "status":     "ok",
        "session_id": sid,
        "summary":    summary,
        "model_used": model,
        "saved":      True,
    }


async def _tool_multi_model_compare(args: dict) -> dict:
    message     = args["message"]
    models      = args["models"]
    system      = args.get("system_prompt", "You are a helpful assistant.")
    temperature = float(args.get("temperature", 0.7))
    max_tokens  = int(args.get("max_tokens", 1024))

    messages_base = [
        {"role": "system", "content": system},
        {"role": "user",   "content": message},
    ]

    # Invalidi
    unknown = [m for m in models if m not in MODELS]
    if unknown:
        return {
            "error":   f"Modelli non nel catalogo: {unknown}",
            "hint":    "Usa il tool list_models per vedere i modelli disponibili.",
        }

    # Chiamate in parallelo
    log.info("multi_model_compare  models=%s", models)
    tasks = [
        nvidia_chat(
            messages=messages_base,
            model=m,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        for m in models
    ]

    results_raw = await asyncio.gather(*tasks, return_exceptions=True)

    results = {}
    for model_id, res in zip(models, results_raw):
        if isinstance(res, Exception):
            results[model_id] = {"error": str(res)}
        else:
            results[model_id] = {
                "response":   res,
                "model_name": MODELS.get(model_id, {}).get("name", model_id),
            }

    return {
        "message":   message,
        "results":   results,
        "models":    models,
        "timestamp": time.time(),
    }


# Dispatch table
_TOOL_HANDLERS: dict[str, Any] = {
    "nvidia_chat":          _tool_nvidia_chat,
    "nvidia_chat_stream":   _tool_nvidia_chat_stream,
    "list_models":          _tool_list_models,
    "get_model_info":       _tool_get_model_info,
    "memory_stats":         _tool_memory_stats,
    "memory_clear":         _tool_memory_clear,
    "memory_save_summary":  _tool_memory_save_summary,
    "get_history":          _tool_get_history,
    "list_sessions":        _tool_list_sessions,
    "delete_session":       _tool_delete_session,
    "summarize_session":    _tool_summarize_session,
    "multi_model_compare":  _tool_multi_model_compare,
}


async def execute_tool(name: str, args: dict) -> dict:
    handler = _TOOL_HANDLERS.get(name)
    if handler is None:
        raise ValueError(f"Tool sconosciuto: '{name}'")
    return await handler(args)


# ══════════════════════════════════════════════════════════════
# JSON-RPC 2.0 HANDLER (core)
# ══════════════════════════════════════════════════════════════

async def handle_jsonrpc(body: dict) -> dict:
    """
    Processa un singolo messaggio JSON-RPC 2.0 MCP.
    Ritorna il dict della risposta (o errore).
    """
    req_id  = body.get("id", 0)
    method  = body.get("method", "")
    params  = body.get("params", {})

    def ok(result: Any) -> dict:
        return {"jsonrpc": "2.0", "id": req_id, "result": result}

    def err(code: int, msg: str) -> dict:
        return {"jsonrpc": "2.0", "id": req_id,
                "error": {"code": code, "message": msg}}

    log.info("MCP method=%s  id=%s", method, req_id)

    try:
        # ── Lifecycle ─────────────────────────────────────────
        if method == "initialize":
            return ok({
                "protocolVersion": "2024-11-05",
                "serverInfo": {
                    "name":        "nvidia-ai-mcp",
                    "version":     "2.0.0",
                    "description": "NVIDIA NIM Multi-Model MCP Server con memoria breve/lunga",
                    "author":      "nvidia-ai-space",
                },
                "capabilities": {
                    "tools":     {"listChanged": False},
                    "resources": {"subscribe": False, "listChanged": False},
                    "prompts":   {"listChanged": False},
                },
            })

        elif method == "initialized":
            # Notifica, non richiede risposta
            return ok({})

        elif method == "ping":
            return ok({"pong": True, "timestamp": time.time()})

        # ── Tools ─────────────────────────────────────────────
        elif method == "tools/list":
            cursor = params.get("cursor")        # paginazione (ignorata, ritorna tutto)
            return ok({"tools": MCP_TOOLS, "nextCursor": None})

        elif method == "tools/call":
            tool_name = params.get("name", "")
            tool_args = params.get("arguments", {})

            if not tool_name:
                return err(-32602, "Parametro 'name' mancante in tools/call")

            result = await execute_tool(tool_name, tool_args)
            return ok({
                "content": [
                    {
                        "type": "text",
                        "text": json.dumps(result, ensure_ascii=False, indent=2),
                    }
                ],
                "isError": False,
            })

        # ── Resources ─────────────────────────────────────────
        elif method == "resources/list":
            return ok({
                "resources": [
                    {
                        "uri":         "nvidia://models",
                        "name":        "NVIDIA NIM Models Catalog",
                        "description": "Catalogo completo dei modelli NVIDIA NIM disponibili",
                        "mimeType":    "application/json",
                    },
                    {
                        "uri":         "nvidia://sessions",
                        "name":        "Active Sessions",
                        "description": "Lista delle sessioni di memoria attive",
                        "mimeType":    "application/json",
                    },
                    {
                        "uri":         "nvidia://health",
                        "name":        "Server Health",
                        "description": "Stato del server MCP e della connessione NVIDIA",
                        "mimeType":    "application/json",
                    },
                ],
                "nextCursor": None,
            })

        elif method == "resources/read":
            uri = params.get("uri", "")

            if uri == "nvidia://models":
                content = json.dumps(MODELS, ensure_ascii=False, indent=2)
            elif uri == "nvidia://sessions":
                sessions = memory.list_sessions()
                stats    = {s: memory.get_stats(s) for s in sessions}
                content  = json.dumps({"sessions": sessions, "stats": stats}, ensure_ascii=False, indent=2)
            elif uri == "nvidia://health":
                content = json.dumps({
                    "status":          "healthy",
                    "timestamp":       time.time(),
                    "models_in_catalog": len(MODELS),
                    "active_sessions": len(memory.list_sessions()),
                    "nvidia_endpoint": NVIDIA_BASE_URL,
                    "api_key_set":     bool(NVIDIA_API_KEY),
                }, indent=2)
            else:
                return err(-32602, f"Risorsa sconosciuta: {uri}")

            return ok({
                "contents": [
                    {"uri": uri, "mimeType": "application/json", "text": content}
                ]
            })

        # ── Prompts ───────────────────────────────────────────
        elif method == "prompts/list":
            return ok({
                "prompts": [
                    {
                        "name":        "code_review",
                        "description": "Revisione del codice con spiegazione dettagliata",
                        "arguments": [
                            {"name": "code",     "description": "Codice da revisionare", "required": True},
                            {"name": "language", "description": "Linguaggio di programmazione", "required": False},
                        ],
                    },
                    {
                        "name":        "summarize",
                        "description": "Riassumi un testo lungo",
                        "arguments": [
                            {"name": "text",     "description": "Testo da riassumere", "required": True},
                            {"name": "language", "description": "Lingua output",       "required": False},
                        ],
                    },
                    {
                        "name":        "translate",
                        "description": "Traduci un testo",
                        "arguments": [
                            {"name": "text",        "description": "Testo da tradurre", "required": True},
                            {"name": "target_lang", "description": "Lingua target",     "required": True},
                        ],
                    },
                ],
                "nextCursor": None,
            })

        elif method == "prompts/get":
            prompt_name = params.get("name", "")
            arguments   = params.get("arguments", {})

            if prompt_name == "code_review":
                code = arguments.get("code", "")
                lang = arguments.get("language", "auto-detect")
                return ok({
                    "description": "Revisione del codice",
                    "messages": [
                        {
                            "role": "user",
                            "content": {
                                "type": "text",
                                "text": (
                                    f"Fai una revisione dettagliata del seguente codice {lang}. "
                                    f"Identifica bug, problemi di performance, sicurezza e best practices.\n\n"
                                    f"```{lang}\n{code}\n```"
                                ),
                            },
                        }
                    ],
                })

            elif prompt_name == "summarize":
                text = arguments.get("text", "")
                lang = arguments.get("language", "italiano")
                return ok({
                    "description": "Riassunto testo",
                    "messages": [
                        {
                            "role": "user",
                            "content": {
                                "type": "text",
                                "text": (
                                    f"Riassumi il seguente testo in {lang}, "
                                    f"evidenziando i punti chiave:\n\n{text}"
                                ),
                            },
                        }
                    ],
                })

            elif prompt_name == "translate":
                text = arguments.get("text", "")
                tgt  = arguments.get("target_lang", "inglese")
                return ok({
                    "description": "Traduzione",
                    "messages": [
                        {
                            "role": "user",
                            "content": {
                                "type": "text",
                                "text": f"Traduci il seguente testo in {tgt}:\n\n{text}",
                            },
                        }
                    ],
                })

            return err(-32602, f"Prompt sconosciuto: {prompt_name}")

        # ── Shutdown ──────────────────────────────────────────
        elif method == "shutdown":
            log.info("MCP shutdown richiesto.")
            return ok({"message": "Server in chiusura..."})

        else:
            return err(-32601, f"Metodo non trovato: '{method}'")

    except ValueError as ve:
        return err(-32602, str(ve))
    except RuntimeError as re:
        return err(-32000, str(re))
    except Exception as exc:
        log.exception("Errore inatteso nel handler MCP: %s", exc)
        return err(-32603, f"Errore interno: {exc}")


# ══════════════════════════════════════════════════════════════
# FASTAPI APP — HTTP + SSE
# ══════════════════════════════════════════════════════════════

mcp_app = FastAPI(
    title="NVIDIA AI MCP Server",
    description="Model Context Protocol server per NVIDIA NIM – HTTP + SSE + STDIO",
    version="2.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

mcp_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@mcp_app.get("/")
async def root():
    return {
        "name":         "NVIDIA AI MCP Server",
        "version":      "2.0.0",
        "protocol":     "MCP 2024-11-05",
        "transport":    ["HTTP POST /mcp", "SSE GET /mcp/sse", "STDIO"],
        "tools":        len(MCP_TOOLS),
        "models":       len(MODELS),
        "docs":         "/docs",
        "health":       "/health",
    }


@mcp_app.get("/health")
async def health():
    return {
        "status":          "healthy",
        "timestamp":       time.time(),
        "tools":           len(MCP_TOOLS),
        "models_catalog":  len(MODELS),
        "active_sessions": len(memory.list_sessions()),
        "nvidia_endpoint": NVIDIA_BASE_URL,
        "api_key_ok":      bool(NVIDIA_API_KEY),
    }


# ── Transport 1: HTTP POST ─────────────────────────────────────

@mcp_app.post("/mcp")
async def mcp_http(request: Request):
    """
    MCP via HTTP POST (JSON-RPC 2.0).

    Supporta sia singolo messaggio che batch (lista).
    """
    try:
        body = await request.json()
    except Exception:
        return JSONResponse(
            {"jsonrpc": "2.0", "id": None,
             "error": {"code": -32700, "message": "JSON non valido"}},
            status_code=400,
        )

    # Batch request (lista di messaggi)
    if isinstance(body, list):
        responses = await asyncio.gather(*[handle_jsonrpc(b) for b in body])
        return JSONResponse(list(responses))

    # Singola request
    response = await handle_jsonrpc(body)
    return JSONResponse(response)


# ── Transport 2: SSE (notifiche + streaming chat) ─────────────

@mcp_app.get("/mcp/sse")
async def mcp_sse_endpoint(request: Request):
    """
    MCP SSE transport.

    Il client invia il metodo come query param `method` e
    i params come JSON in `params`.

    Esempio:
      GET /mcp/sse?method=tools/list
      GET /mcp/sse?method=tools/call&params={"name":"list_models","arguments":{}}
    """
    method     = request.query_params.get("method", "tools/list")
    params_raw = request.query_params.get("params", "{}")
    req_id     = request.query_params.get("id", "1")

    try:
        params = json.loads(params_raw)
    except json.JSONDecodeError:
        params = {}

    body = {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params}

    async def event_gen():
        # Evento di connessione
        yield {
            "event": "connected",
            "data": json.dumps({
                "server": "nvidia-ai-mcp",
                "protocol": "2024-11-05",
                "timestamp": time.time(),
            }),
        }
        # Elabora il metodo
        result = await handle_jsonrpc(body)
        yield {
            "event": "message",
            "data": json.dumps(result),
        }
        # Fine sessione SSE
        yield {
            "event": "close",
            "data": json.dumps({"reason": "request completed"}),
        }

    return EventSourceResponse(event_gen())


@mcp_app.get("/mcp/sse/chat")
async def mcp_sse_chat(
    message:      str   = "Hello!",
    model:        str   = DEFAULT_MODEL,
    session_id:   Optional[str] = None,
    system_prompt: str  = "You are a helpful AI assistant.",
    temperature:  float = 0.7,
    max_tokens:   int   = 4096,
    use_memory:   bool  = True,
):
    """
    SSE streaming chat dedicato.

    Questo endpoint apre uno stream SSE e invia token-by-token.

    Events:
      - start  → {session_id, model, timestamp}
      - token  → {token: "..."}
      - done   → {full_response, session_id, memory_stats, timestamp}
      - error  → {error: "..."}
    """
    sid = session_id or str(uuid.uuid4())

    if use_memory:
        memory.add_short(sid, "user", message)
        messages = memory.get_full_context(sid, system_prompt)
    else:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": message},
        ]

    async def event_gen():
        full: list[str] = []
        try:
            yield {
                "event": "start",
                "data": json.dumps({
                    "session_id": sid,
                    "model":      model,
                    "timestamp":  time.time(),
                }),
            }

            async for chunk in nvidia_chat_stream(
                messages=messages,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
            ):
                full.append(chunk)
                yield {
                    "event": "token",
                    "data": json.dumps({"token": chunk}),
                }

            complete = "".join(full)
            if use_memory:
                memory.add_short(sid, "assistant", complete)

            yield {
                "event": "done",
                "data": json.dumps({
                    "full_response": complete,
                    "session_id":    sid,
                    "memory_stats":  memory.get_stats(sid),
                    "timestamp":     time.time(),
                }),
            }

        except Exception as exc:
            log.exception("SSE chat error: %s", exc)
            yield {
                "event": "error",
                "data": json.dumps({"error": str(exc)}),
            }

    return EventSourceResponse(event_gen())


# ── Transport 3: STDIO (Claude Desktop, Cursor, …) ────────────

async def stdio_server():
    """
    Loop STDIO per client MCP (Claude Desktop, Cursor, ecc.).

    Legge richieste JSON-RPC da stdin (una per riga),
    scrive le risposte su stdout.
    """
    log.info("MCP STDIO server avviato. In attesa di messaggi su stdin...")

    reader = asyncio.StreamReader()
    proto  = asyncio.StreamReaderProtocol(reader)
    loop   = asyncio.get_event_loop()

    await loop.connect_read_pipe(lambda: proto, sys.stdin)

    writer_transport, writer_proto = await loop.connect_write_pipe(
        lambda: asyncio.BaseProtocol(), sys.stdout
    )

    async def write_response(data: dict):
        line = json.dumps(data, ensure_ascii=False) + "\n"
        writer_transport.write(line.encode("utf-8"))

    while True:
        try:
            line_bytes = await reader.readline()
            if not line_bytes:
                break  # EOF

            line = line_bytes.decode("utf-8").strip()
            if not line:
                continue

            try:
                body = json.loads(line)
            except json.JSONDecodeError as e:
                await write_response({
                    "jsonrpc": "2.0", "id": None,
                    "error": {"code": -32700, "message": f"JSON parse error: {e}"},
                })
                continue

            # Gestisci batch
            if isinstance(body, list):
                responses = await asyncio.gather(*[handle_jsonrpc(b) for b in body])
                for r in responses:
                    await write_response(r)
            else:
                response = await handle_jsonrpc(body)
                # Le notifiche (no id) non ricevono risposta
                if body.get("id") is not None:
                    await write_response(response)

        except asyncio.CancelledError:
            break
        except Exception as exc:
            log.exception("STDIO loop error: %s", exc)


# ══════════════════════════════════════════════════════════════
# ENTRY POINT
# ══════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(
        description="NVIDIA AI MCP Server – Model Context Protocol 2024-11-05"
    )
    parser.add_argument(
        "--stdio",
        action="store_true",
        help="Avvia in modalità STDIO (per Claude Desktop / Cursor)",
    )
    parser.add_argument(
        "--host", default=MCP_HOST,
        help=f"Host HTTP (default: {MCP_HOST})"
    )
    parser.add_argument(
        "--port", type=int, default=MCP_PORT,
        help=f"Porta HTTP (default: {MCP_PORT})"
    )
    parser.add_argument(
        "--log-level", default="info",
        choices=["debug", "info", "warning", "error"],
    )
    args = parser.parse_args()

    logging.getLogger().setLevel(args.log_level.upper())

    if args.stdio:
        log.info("Avvio MCP in modalità STDIO")
        asyncio.run(stdio_server())
    else:
        log.info("Avvio MCP HTTP server su %s:%d", args.host, args.port)
        log.info("Endpoints:")
        log.info("  POST http://%s:%d/mcp        (JSON-RPC HTTP)", args.host, args.port)
        log.info("  GET  http://%s:%d/mcp/sse    (SSE transport)", args.host, args.port)
        log.info("  GET  http://%s:%d/mcp/sse/chat (SSE streaming chat)", args.host, args.port)
        log.info("  GET  http://%s:%d/docs        (Swagger UI)", args.host, args.port)
        log.info("Tools disponibili: %d", len(MCP_TOOLS))
        log.info("Modelli nel catalogo: %d", len(MODELS))
        uvicorn.run(
            mcp_app,
            host=args.host,
            port=args.port,
            log_level=args.log_level,
        )


if __name__ == "__main__":
    main()