app.py

"""

🚀 NVIDIA AI Multi-Model Space v3.1

- 38 modelli NVIDIA NIM VERIFICATI e FUNZIONANTI

- Auto-adattamento messaggi per modelli senza 'system role'

- Memoria breve (RAM) e lunga (SQLite)

- API REST GET/POST + SSE Streaming

- Server MCP integrato

- Test parallelo async ultra-veloce

"""

import os
import json
import time
import uuid
import base64
import asyncio
from pathlib import Path
from typing import Optional, AsyncGenerator

import gradio as gr
import httpx
from dotenv import load_dotenv
from fastapi import FastAPI, Request, Query, HTTPException
from fastapi.responses import (
    StreamingResponse, JSONResponse, HTMLResponse, FileResponse
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from sse_starlette.sse import EventSourceResponse
from pydantic import BaseModel, Field, ConfigDict

from memory_manager import MemoryManager

load_dotenv()

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

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

NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY") or _decode_key()
NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1"
DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "meta/llama-3.2-3b-instruct")

HTTP_LIMITS = httpx.Limits(max_connections=100, max_keepalive_connections=20)
HTTP_TIMEOUT = httpx.Timeout(120.0, connect=10.0)

# ══════════════════════════════════════════════════════════════
# CATALOGO — SOLO MODELLI VERIFICATI FUNZIONANTI (test 2026-07-04)
# ══════════════════════════════════════════════════════════════

NVIDIA_MODELS = {
    # ═══════════════════════════════════════════════════════════
    # ⚡ TOP FASTEST (< 300ms)
    # ═══════════════════════════════════════════════════════════
    "meta/llama-3.2-3b-instruct": {
        "name": "⚡ LLaMA 3.2 3B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "🥇 FASTEST — 153ms — DEFAULT"
    },
    "nvidia/nemotron-mini-4b-instruct": {
        "name": "⚡ Nemotron Mini 4B",
        "category": "LLM",
        "context": 8192,
        "description": "Ultra-fast 4B — 186ms"
    },
    "meta/llama-3.2-11b-vision-instruct": {
        "name": "⚡ LLaMA 3.2 11B Vision",
        "category": "Vision",
        "context": 131072,
        "description": "Fast vision-language — 196ms"
    },
    "microsoft/phi-4-mini-instruct": {
        "name": "⚡ Phi 4 Mini Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "🆕 Latest Phi 4 Mini — 206ms"
    },
    "google/gemma-2-2b-it": {
        "name": "⚡ Gemma 2 2B",
        "category": "LLM",
        "context": 8192,
        "description": "Compact Gemma — 212ms (no system role)"
    },
    "nvidia/llama-3.1-nemotron-nano-8b-v1": {
        "name": "⚡ Nemotron Nano 8B v1",
        "category": "LLM",
        "context": 131072,
        "description": "Efficient Nemotron — 239ms"
    },
    "stockmark/stockmark-2-100b-instruct": {
        "name": "⚡ Stockmark 2 100B",
        "category": "LLM",
        "context": 32768,
        "description": "Japanese specialist 100B — 261ms"
    },
    "upstage/solar-10.7b-instruct": {
        "name": "⚡ Solar 10.7B",
        "category": "LLM",
        "context": 4096,
        "description": "Upstage Solar — 261ms"
    },
    "meta/llama-3.1-70b-instruct": {
        "name": "⚡ LLaMA 3.1 70B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "High-quality 70B — 263ms"
    },
    "sarvamai/sarvam-m": {
        "name": "⚡ Sarvam M",
        "category": "LLM",
        "context": 32768,
        "description": "Indian multilingual — 285ms"
    },

    # ═══════════════════════════════════════════════════════════
    # 🚀 FAST (300-500ms)
    # ═══════════════════════════════════════════════════════════
    "nvidia/ising-calibration-1-35b-a3b": {
        "name": "🚀 Ising Calibration 35B",
        "category": "LLM",
        "context": 131072,
        "description": "Calibration model — 331ms"
    },
    "nvidia/nemotron-nano-12b-v2-vl": {
        "name": "🚀 Nemotron Nano 12B v2 VL",
        "category": "Vision",
        "context": 131072,
        "description": "Vision Nemotron 12B — 334ms"
    },
    "mistralai/mistral-small-4-119b-2603": {
        "name": "🚀 Mistral Small 4 119B",
        "category": "LLM",
        "context": 131072,
        "description": "Mistral Small 4 — 357ms"
    },
    "qwen/qwen3-next-80b-a3b-instruct": {
        "name": "🚀 Qwen3 Next 80B",
        "category": "LLM",
        "context": 131072,
        "description": "🆕 Latest Qwen 3 — 371ms"
    },
    "google/gemma-3n-e2b-it": {
        "name": "🚀 Gemma 3n E2B",
        "category": "LLM",
        "context": 8192,
        "description": "Gemma 3 Nano — 380ms (no system role)"
    },
    "mistralai/mistral-nemotron": {
        "name": "🚀 Mistral Nemotron",
        "category": "LLM",
        "context": 131072,
        "description": "NVIDIA-Mistral collab — 383ms"
    },
    "meta/llama-3.1-8b-instruct": {
        "name": "🚀 LLaMA 3.1 8B Instruct",
        "category": "LLM",
        "context": 131072,
        "description": "Popular 8B — 404ms"
    },
    "minimaxai/minimax-m2.7": {
        "name": "🚀 MiniMax M2.7",
        "category": "LLM",
        "context": 131072,
        "description": "MiniMax model — 409ms"
    },
    "nvidia/nemotron-3-nano-30b-a3b": {
        "name": "🚀 Nemotron 3 Nano 30B",
        "category": "LLM",
        "context": 131072,
        "description": "🆕 Nemotron 3 MoE — 410ms"
    },
    "meta/llama-3.2-90b-vision-instruct": {
        "name": "🚀 LLaMA 3.2 90B Vision",
        "category": "Vision",
        "context": 131072,
        "description": "Large vision-language — 419ms"
    },
    "mistralai/mistral-medium-3.5-128b": {
        "name": "🚀 Mistral Medium 3.5 128B",
        "category": "LLM",
        "context": 131072,
        "description": "Mistral Medium 3.5 — 419ms"
    },

    # ═══════════════════════════════════════════════════════════
    # 💪 POWERFUL (500ms-1s)
    # ═══════════════════════════════════════════════════════════
    "qwen/qwen3.5-122b-a10b": {
        "name": "💪 Qwen 3.5 122B",
        "category": "LLM",
        "context": 131072,
        "description": "🔥 Qwen 3.5 MoE — 505ms"
    },
    "nvidia/nemotron-3-super-120b-a12b": {
        "name": "💪 Nemotron 3 Super 120B",
        "category": "LLM",
        "context": 131072,
        "description": "🔥 Nemotron 3 Super MoE — 556ms"
    },
    "mistralai/mixtral-8x7b-instruct-v0.1": {
        "name": "💪 Mixtral 8x7B",
        "category": "LLM",
        "context": 32768,
        "description": "Classic MoE — 685ms"
    },
    "nvidia/nemotron-3-ultra-550b-a55b": {
        "name": "💪 Nemotron 3 Ultra 550B",
        "category": "LLM",
        "context": 131072,
        "description": "🚀 MASSIVE 550B MoE — 707ms"
    },
    "deepseek-ai/deepseek-v4-flash": {
        "name": "💪 DeepSeek V4 Flash",
        "category": "LLM",
        "context": 131072,
        "description": "🆕 Fast DeepSeek V4 — 712ms"
    },
    "abacusai/dracarys-llama-3.1-70b-instruct": {
        "name": "💪 Dracarys LLaMA 3.1 70B",
        "category": "LLM",
        "context": 131072,
        "description": "AbacusAI fine-tuned 70B — 774ms"
    },
    "google/diffusiongemma-26b-a4b-it": {
        "name": "💪 DiffusionGemma 26B",
        "category": "LLM",
        "context": 131072,
        "description": "Diffusion-based Gemma — 843ms (no system role)"
    },

    # ═══════════════════════════════════════════════════════════
    # 🐢 SLOWER but WORKING (>1s)
    # ═══════════════════════════════════════════════════════════
    "moonshotai/kimi-k2.6": {
        "name": "🐢 Kimi K2.6",
        "category": "LLM",
        "context": 131072,
        "description": "🆕 Moonshot Kimi — 1.12s"
    },
    "z-ai/glm-5.2": {
        "name": "🐢 GLM 5.2",
        "category": "LLM",
        "context": 131072,
        "description": "🆕 Zhipu GLM 5.2 — 1.5s"
    },
    "deepseek-ai/deepseek-v4-pro": {
        "name": "🐢 DeepSeek V4 Pro",
        "category": "Reasoning",
        "context": 131072,
        "description": "🔥 Advanced reasoning — 1.7s"
    },
    "meta/llama-3.2-1b-instruct": {
        "name": "🐢 LLaMA 3.2 1B",
        "category": "LLM",
        "context": 131072,
        "description": "Ultra-lightweight — 1.9s"
    },
    "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning": {
        "name": "🐢 Nemotron 3 Nano Omni (Reasoning)",
        "category": "Reasoning",
        "context": 131072,
        "description": "🧠 Reasoning model — 2.6s"
    },
    "minimaxai/minimax-m3": {
        "name": "🐢 MiniMax M3",
        "category": "LLM",
        "context": 131072,
        "description": "🆕 Latest MiniMax — 3.6s"
    },
    "mistralai/ministral-14b-instruct-2512": {
        "name": "🐢 Ministral 14B",
        "category": "LLM",
        "context": 131072,
        "description": "Mistral Ministral — 4.6s"
    },
    "nvidia/llama-3.1-nemotron-nano-vl-8b-v1": {
        "name": "🐢 Nemotron Nano VL 8B",
        "category": "Vision",
        "context": 131072,
        "description": "Vision Nemotron — 5.9s"
    },
    "mistralai/mistral-large-3-675b-instruct-2512": {
        "name": "🐢 Mistral Large 3 675B",
        "category": "LLM",
        "context": 131072,
        "description": "🚀 MASSIVE 675B — 8.2s"
    },

    # ═══════════════════════════════════════════════════════════
    # 🌍 TRANSLATION
    # ═══════════════════════════════════════════════════════════
    "nvidia/riva-translate-4b-instruct-v1.1": {
        "name": "🌍 Riva Translate 4B v1.1",
        "category": "Translation",
        "context": 8192,
        "description": "NVIDIA translation — 178ms"
    },
}

# ══════════════════════════════════════════════════════════════
# MODELLI CHE NON SUPPORTANO IL ROLE "system"
# (verranno auto-convertiti: system → prefisso del primo user msg)
# ══════════════════════════════════════════════════════════════

MODELS_NO_SYSTEM_ROLE = {
    # Gemma family
    "google/gemma-2-2b-it",
    "google/gemma-3n-e2b-it",
    "google/diffusiongemma-26b-a4b-it",
    # Aggiornato dinamicamente in caso di errore "system role not supported"
}


def normalize_messages(messages: list[dict], model: str) -> list[dict]:
    """

    Adatta i messaggi al modello:

    - Se il modello non supporta 'system', converte il system prompt

      come prefisso del primo messaggio user.

    - Rimuove messaggi vuoti.

    """
    if not messages:
        return messages

    # Se il modello supporta system, non fare nulla
    if model not in MODELS_NO_SYSTEM_ROLE:
        return messages

    # Estrai tutti i system prompt
    system_parts = []
    non_system = []
    for m in messages:
        if m.get("role") == "system":
            content = (m.get("content") or "").strip()
            if content:
                system_parts.append(content)
        else:
            non_system.append(m)

    if not system_parts:
        return non_system if non_system else messages

    # Prependi il system al primo user message
    system_text = "\n\n".join(system_parts)
    prepended = False
    for i, m in enumerate(non_system):
        if m.get("role") == "user":
            non_system[i] = {
                "role": "user",
                "content": f"[System instructions]\n{system_text}\n\n[User message]\n{m.get('content', '')}"
            }
            prepended = True
            break

    if not prepended:
        # Nessun user msg trovato, aggiungi come user
        non_system.insert(0, {"role": "user", "content": system_text})

    return non_system


# ══════════════════════════════════════════════════════════════
# MEMORIA
# ══════════════════════════════════════════════════════════════

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


# ══════════════════════════════════════════════════════════════
# CLIENT NVIDIA (con auto-adattamento system role)
# ══════════════════════════════════════════════════════════════

def _is_system_role_error(err_text: str) -> bool:
    """Rileva errori dovuti a system role non supportato"""
    if not err_text:
        return False
    err_lower = err_text.lower()
    keywords = [
        "system role not supported",
        "does not support 'system'",
        "does not support system",
        "system role is not supported",
        "role 'system' not allowed",
    ]
    return any(k in err_lower for k in keywords)


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]:
    if not model or not str(model).strip():
        model = DEFAULT_MODEL

    # 🎯 Auto-adatta i messaggi al modello
    original_messages = messages
    messages = normalize_messages(messages, model)

    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=HTTP_TIMEOUT, limits=HTTP_LIMITS) as client:
        async with client.stream(
            "POST",
            f"{NVIDIA_BASE_URL}/chat/completions",
            headers=headers, json=payload,
        ) as response:

            # 🎯 AUTO-RETRY se system role non supportato
            if response.status_code != 200:
                err = (await response.aread()).decode(errors="ignore")[:500]
                if _is_system_role_error(err) and model not in MODELS_NO_SYSTEM_ROLE:
                    print(f"[AUTO-FIX] Aggiunto {model} a MODELS_NO_SYSTEM_ROLE")
                    MODELS_NO_SYSTEM_ROLE.add(model)
                    # Ri-normalizza e riprova (con un nuovo stream)
                    payload["messages"] = normalize_messages(original_messages, model)
                    async with client.stream(
                        "POST",
                        f"{NVIDIA_BASE_URL}/chat/completions",
                        headers=headers, json=payload,
                    ) as retry_response:
                        if retry_response.status_code != 200:
                            err2 = (await retry_response.aread()).decode(errors="ignore")[:500]
                            raise Exception(f"NVIDIA API {retry_response.status_code}: {err2}")
                        async for line in retry_response.aiter_lines():
                            if line.startswith("data: "):
                                data = line[6:]
                                if data.strip() == "[DONE]":
                                    return
                                try:
                                    chunk = json.loads(data)
                                    delta = chunk["choices"][0].get("delta", {})
                                    if "content" in delta and delta["content"]:
                                        yield delta["content"]
                                except (json.JSONDecodeError, KeyError, IndexError):
                                    continue
                    return

                # Altri errori
                if response.status_code == 404:
                    raise Exception(f"Modello '{model}' non disponibile (404)")
                raise Exception(f"NVIDIA API {response.status_code}: {err}")

            # Success normale
            async for line in response.aiter_lines():
                if line.startswith("data: "):
                    data = line[6:]
                    if data.strip() == "[DONE]":
                        break
                    try:
                        chunk = json.loads(data)
                        delta = chunk["choices"][0].get("delta", {})
                        if "content" in delta and delta["content"]:
                            yield delta["content"]
                    except (json.JSONDecodeError, KeyError, IndexError):
                        continue


async def nvidia_chat(

    messages: list[dict],

    model: str = DEFAULT_MODEL,

    temperature: float = 0.7,

    max_tokens: int = 4096,

    top_p: float = 0.9,

    client: Optional[httpx.AsyncClient] = None,

) -> str:
    if not model or not str(model).strip():
        model = DEFAULT_MODEL

    # 🎯 Auto-adatta i messaggi al modello
    original_messages = messages
    messages = normalize_messages(messages, model)

    headers = {
        "Authorization": f"Bearer {NVIDIA_API_KEY}",
        "Content-Type": "application/json",
    }

    async def _do_request(c: httpx.AsyncClient, msgs: list[dict]) -> httpx.Response:
        return await c.post(
            f"{NVIDIA_BASE_URL}/chat/completions",
            headers=headers,
            json={
                "model": model,
                "messages": msgs,
                "temperature": temperature,
                "max_tokens": max_tokens,
                "top_p": top_p,
                "stream": False,
            },
        )

    async def _call(c: httpx.AsyncClient) -> str:
        resp = await _do_request(c, messages)

        # 🎯 AUTO-RETRY se system role non supportato
        if resp.status_code != 200:
            err_body = resp.text[:500]
            if _is_system_role_error(err_body) and model not in MODELS_NO_SYSTEM_ROLE:
                print(f"[AUTO-FIX] Aggiunto {model} a MODELS_NO_SYSTEM_ROLE")
                MODELS_NO_SYSTEM_ROLE.add(model)
                # Retry con messaggi rinormalizzati
                retry_msgs = normalize_messages(original_messages, model)
                resp = await _do_request(c, retry_msgs)
                if resp.status_code == 200:
                    return resp.json()["choices"][0]["message"]["content"]

            if resp.status_code == 404:
                raise Exception(f"Modello '{model}' non disponibile")
            raise Exception(f"NVIDIA API {resp.status_code}: {resp.text[:500]}")

        try:
            return resp.json()["choices"][0]["message"]["content"]
        except (KeyError, IndexError, TypeError):
            raise Exception(f"Risposta NVIDIA non valida per {model}: {resp.text[:200]}")

    if client:
        return await _call(client)
    else:
        async with httpx.AsyncClient(timeout=HTTP_TIMEOUT, limits=HTTP_LIMITS) as c:
            return await _call(c)


# ══════════════════════════════════════════════════════════════
# FASTAPI
# ══════════════════════════════════════════════════════════════

app = FastAPI(
    title="🚀 NVIDIA AI Multi-Model API",
    description=f"{len(NVIDIA_MODELS)} modelli verificati e funzionanti",
    version="3.1.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], allow_credentials=True,
    allow_methods=["*"], allow_headers=["*"],
)

STATIC_DIR = Path(__file__).parent / "static"
if STATIC_DIR.exists():
    app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


class ChatRequest(BaseModel):
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "message": "Ciao!",
                "model": DEFAULT_MODEL,
                "session_id": "test",
            }
        }
    )
    message: str
    model: str = Field(default=DEFAULT_MODEL)
    session_id: Optional[str] = None
    system_prompt: str = "Sei un assistente AI utile. Rispondi in italiano se l'utente scrive in italiano."
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int = Field(default=4096, ge=1, le=16384)
    top_p: float = Field(default=0.9, ge=0.0, le=1.0)
    use_memory: bool = True
    stream: bool = False


class MemoryRequest(BaseModel):
    session_id: str
    action: str
    summary: Optional[str] = None


# ══════════════════════════════════════════════════════════════
# ROOT + INFO
# ══════════════════════════════════════════════════════════════

@app.get("/", response_class=HTMLResponse)
async def root():
    ui_file = STATIC_DIR / "index.html"
    if ui_file.exists():
        return FileResponse(str(ui_file))
    return HTMLResponse("""

    <html><body style="font-family:sans-serif;background:#0a0e14;color:#e4e6eb;padding:40px">

      <h1>🚀 NVIDIA AI Multi-Model</h1>

      <ul>

        <li><a href="/ui" style="color:#76b900">Gradio</a></li>

        <li><a href="/docs" style="color:#76b900">Swagger</a></li>

        <li><a href="/api-docs" style="color:#76b900">API Docs</a></li>

      </ul>

    </body></html>

    """)


@app.get("/info")
async def info():
    return {
        "service": "🚀 NVIDIA AI Multi-Model",
        "version": "3.1.0",
        "default_model": DEFAULT_MODEL,
        "models_available": len(NVIDIA_MODELS),
        "models_no_system_role": sorted(MODELS_NO_SYSTEM_ROLE),
        "all_models_verified": True,
        "last_verified": "2026-07-04",
    }


@app.get("/v1/health")
async def health():
    return {
        "status": "healthy",
        "timestamp": time.time(),
        "models_count": len(NVIDIA_MODELS),
        "default_model": DEFAULT_MODEL,
        "active_sessions": len(memory.short_memory),
    }


# ══════════════════════════════════════════════════════════════
# MODELS
# ══════════════════════════════════════════════════════════════

@app.get("/v1/models")
async def list_models(category: Optional[str] = None):
    models = {}
    for mid, mi in NVIDIA_MODELS.items():
        if category and mi["category"].lower() != category.lower():
            continue
        models[mid] = mi
    return {
        "models": models,
        "total": len(models),
        "categories": sorted(set(m["category"] for m in NVIDIA_MODELS.values())),
        "default": DEFAULT_MODEL,
        "models_no_system_role": sorted(MODELS_NO_SYSTEM_ROLE),
    }


@app.get("/v1/models/{model_id:path}")
async def get_model_info(model_id: str):
    if model_id not in NVIDIA_MODELS:
        raise HTTPException(404, f"Modello '{model_id}' non nel catalogo")
    return {
        "model_id": model_id,
        **NVIDIA_MODELS[model_id],
        "supports_system_role": model_id not in MODELS_NO_SYSTEM_ROLE,
    }


# ══════════════════════════════════════════════════════════════
# DEBUG
# ══════════════════════════════════════════════════════════════

@app.get("/v1/debug/test-model")
async def debug_test_model(model: str = Query(DEFAULT_MODEL)):
    """Testa un modello con system + user (come nella UI)"""
    try:
        start = time.time()
        response = await nvidia_chat(
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "Say 'ok'."},
            ],
            model=model, max_tokens=10,
        )
        return {
            "status": "✅ OK",
            "model": model,
            "response": response,
            "elapsed_ms": round((time.time() - start) * 1000),
            "supports_system_role": model not in MODELS_NO_SYSTEM_ROLE,
        }
    except Exception as e:
        return {"status": "❌ ERROR", "model": model, "error": str(e)}


@app.get("/v1/debug/no-system-models")
async def debug_no_system_models():
    """Lista modelli che non supportano il role 'system'"""
    return {
        "models_no_system_role": sorted(MODELS_NO_SYSTEM_ROLE),
        "total": len(MODELS_NO_SYSTEM_ROLE),
        "note": "Questi modelli ricevono automaticamente il system_prompt come prefisso del primo user message."
    }


async def _test_single(mid: str, client: httpx.AsyncClient, sem: asyncio.Semaphore) -> dict:
    async with sem:
        start = time.time()
        try:
            # Test realistico: system + user (come la UI)
            resp = await nvidia_chat(
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "Say 'ok'"},
                ],
                model=mid, max_tokens=5, client=client,
            )
            return {
                "model": mid, "status": "✅",
                "elapsed_ms": round((time.time() - start) * 1000),
                "preview": (resp or "")[:30],
                "supports_system_role": mid not in MODELS_NO_SYSTEM_ROLE,
            }
        except Exception as e:
            return {
                "model": mid, "status": "❌",
                "elapsed_ms": round((time.time() - start) * 1000),
                "error": str(e)[:150],
            }


@app.get("/v1/debug/test-all-parallel")
async def debug_test_all_parallel(concurrency: int = Query(15, ge=1, le=50)):
    """⚡ Test parallelo async di tutti i modelli del catalogo"""
    start = time.time()
    testable = list(NVIDIA_MODELS.keys())
    sem = asyncio.Semaphore(concurrency)
    async with httpx.AsyncClient(timeout=HTTP_TIMEOUT, limits=HTTP_LIMITS) as client:
        tasks = [_test_single(m, client, sem) for m in testable]
        results = await asyncio.gather(*tasks)
    working = [r for r in results if r["status"] == "✅"]
    broken = [r for r in results if r["status"] == "❌"]
    working.sort(key=lambda x: x["elapsed_ms"])
    return {
        "summary": {
            "total_tested": len(results),
            "working": len(working),
            "broken": len(broken),
            "total_time_seconds": round(time.time() - start, 2),
            "models_no_system_role": sorted(MODELS_NO_SYSTEM_ROLE),
        },
        "working_models_by_speed": working,
        "broken_models": broken,
    }


# ══════════════════════════════════════════════════════════════
# CHAT
# ══════════════════════════════════════════════════════════════

@app.post("/v1/chat")
async def chat_post(req: ChatRequest):
    sid = req.session_id or str(uuid.uuid4())
    if req.stream:
        return await _stream_response(req, sid)

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

    try:
        response = await nvidia_chat(
            messages=messages, model=req.model,
            temperature=req.temperature, max_tokens=req.max_tokens, top_p=req.top_p,
        )
        if req.use_memory:
            memory.add_short(sid, "assistant", response)
        return {
            "response": response, "model": req.model, "session_id": sid,
            "memory_stats": memory.get_stats(sid) if req.use_memory else None,
            "timestamp": time.time(),
        }
    except Exception as e:
        raise HTTPException(500, detail=str(e))


@app.get("/v1/chat/get")
async def chat_get(

    message: str = Query(...),

    model: str = Query(DEFAULT_MODEL),

    session_id: Optional[str] = Query(None),

    system_prompt: str = Query("Sei un assistente AI utile."),

    temperature: float = Query(0.7),

    max_tokens: int = Query(4096),

    use_memory: bool = Query(True),

    stream: bool = Query(False),

):
    req = ChatRequest(
        message=message, model=model, session_id=session_id,
        system_prompt=system_prompt, temperature=temperature,
        max_tokens=max_tokens, use_memory=use_memory, stream=stream,
    )
    if stream:
        return await _stream_response(req, session_id or str(uuid.uuid4()))
    return await chat_post(req)


@app.post("/v1/chat/stream")
async def chat_stream_post(req: ChatRequest):
    req.stream = True
    return await _stream_response(req, req.session_id or str(uuid.uuid4()))


@app.get("/v1/chat/stream")
async def chat_stream_get(

    message: str = Query(...),

    model: str = Query(DEFAULT_MODEL),

    session_id: Optional[str] = Query(None),

    system_prompt: str = Query("Sei un assistente AI utile."),

    temperature: float = Query(0.7),

    max_tokens: int = Query(4096),

    use_memory: bool = Query(True),

):
    req = ChatRequest(
        message=message, model=model, session_id=session_id,
        system_prompt=system_prompt, temperature=temperature,
        max_tokens=max_tokens, use_memory=use_memory, stream=True,
    )
    return await _stream_response(req, session_id or str(uuid.uuid4()))


async def _stream_response(req: ChatRequest, session_id: str):
    if req.use_memory:
        memory.add_short(session_id, "user", req.message)
        messages = memory.get_full_context(session_id, req.system_prompt)
    else:
        messages = [
            {"role": "system", "content": req.system_prompt},
            {"role": "user", "content": req.message},
        ]

    async def event_generator():
        full = []
        try:
            yield {"event": "start", "data": json.dumps({
                "session_id": session_id, "model": req.model, "timestamp": time.time(),
                "supports_system_role": req.model not in MODELS_NO_SYSTEM_ROLE,
            })}
            async for chunk in nvidia_chat_stream(
                messages=messages, model=req.model,
                temperature=req.temperature, max_tokens=req.max_tokens, top_p=req.top_p,
            ):
                full.append(chunk)
                yield {"event": "token", "data": json.dumps({"token": chunk})}

            complete = "".join(full)
            if req.use_memory and complete:
                memory.add_short(session_id, "assistant", complete)
            yield {"event": "done", "data": json.dumps({
                "full_response": complete, "session_id": session_id,
                "memory_stats": memory.get_stats(session_id) if req.use_memory else None,
                "timestamp": time.time(),
            })}
        except Exception as e:
            yield {"event": "error", "data": json.dumps({"error": str(e)})}

    return EventSourceResponse(event_generator())


# ══════════════════════════════════════════════════════════════
# MEMORY
# ══════════════════════════════════════════════════════════════

@app.post("/v1/memory")
async def manage_memory(req: MemoryRequest):
    if req.action == "get_stats":
        return memory.get_stats(req.session_id)
    elif req.action == "clear_short":
        memory.clear_short(req.session_id)
        return {"status": "cleared"}
    elif req.action == "clear_all":
        memory.delete_session(req.session_id)
        return {"status": "deleted"}
    elif req.action == "save_summary":
        if not req.summary:
            raise HTTPException(400, "summary required")
        memory.save_session_summary(req.session_id, req.summary)
        return {"status": "saved"}
    raise HTTPException(400, f"Unknown action: {req.action}")


@app.get("/v1/memory/{session_id}")
async def get_memory_stats(session_id: str):
    return memory.get_stats(session_id)


@app.get("/v1/memory/{session_id}/history")
async def get_memory_history(session_id: str):
    return {
        "session_id": session_id,
        "short_memory": memory.get_short(session_id),
        "long_context": memory.get_long_context(session_id),
        "stats": memory.get_stats(session_id),
    }


@app.get("/v1/sessions")
async def list_sessions():
    return {"sessions": memory.list_sessions()}


@app.delete("/v1/sessions/{session_id}")
async def delete_session(session_id: str):
    memory.delete_session(session_id)
    return {"status": "deleted", "session_id": session_id}


# ══════════════════════════════════════════════════════════════
# MCP
# ══════════════════════════════════════════════════════════════

MCP_TOOLS = [
    {
        "name": "nvidia_chat",
        "description": "Chat with any verified NVIDIA NIM model with memory (auto-adapts system role)",
        "inputSchema": {
            "type": "object",
            "properties": {
                "message": {"type": "string"},
                "model": {"type": "string", "default": DEFAULT_MODEL},
                "session_id": {"type": "string"},
                "system_prompt": {"type": "string", "default": "You are a helpful AI assistant."},
                "temperature": {"type": "number", "default": 0.7},
                "max_tokens": {"type": "integer", "default": 4096},
                "use_memory": {"type": "boolean", "default": True},
            },
            "required": ["message"]
        }
    },
    {
        "name": "multi_model_compare",
        "description": "Send same message to multiple models in PARALLEL and compare",
        "inputSchema": {
            "type": "object",
            "properties": {
                "message": {"type": "string"},
                "models": {"type": "array", "items": {"type": "string"}, "minItems": 2, "maxItems": 6},
                "system_prompt": {"type": "string", "default": "You are helpful."},
                "max_tokens": {"type": "integer", "default": 512},
            },
            "required": ["message", "models"]
        }
    },
    {
        "name": "list_models",
        "description": "List all verified NVIDIA models",
        "inputSchema": {
            "type": "object",
            "properties": {"category": {"type": "string"}}
        }
    },
    {
        "name": "memory_stats",
        "description": "Memory stats for session",
        "inputSchema": {
            "type": "object",
            "properties": {"session_id": {"type": "string"}},
            "required": ["session_id"]
        }
    },
    {
        "name": "memory_clear",
        "description": "Clear memory",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "clear_type": {"type": "string", "enum": ["short", "all"], "default": "short"}
            },
            "required": ["session_id"]
        }
    },
    {
        "name": "get_conversation_history",
        "description": "Get full history for session",
        "inputSchema": {
            "type": "object",
            "properties": {"session_id": {"type": "string"}},
            "required": ["session_id"]
        }
    }
]


@app.post("/v1/mcp")
async def mcp_endpoint(request: Request):
    body = await request.json()
    method = body.get("method", "")
    params = body.get("params", {})
    req_id = body.get("id", 1)

    def _ok(r): return JSONResponse({"jsonrpc": "2.0", "id": req_id, "result": r})
    def _err(c, m): return JSONResponse({"jsonrpc": "2.0", "id": req_id, "error": {"code": c, "message": m}})

    if method == "initialize":
        return _ok({
            "protocolVersion": "2024-11-05",
            "serverInfo": {"name": "nvidia-ai", "version": "3.1.0"},
            "capabilities": {
                "tools": {"listChanged": False},
                "resources": {"subscribe": False, "listChanged": False}
            }
        })
    elif method == "tools/list":
        return _ok({"tools": MCP_TOOLS})
    elif method == "tools/call":
        try:
            result = await _execute_mcp_tool(params.get("name", ""), params.get("arguments", {}))
            return _ok({"content": [{"type": "text", "text": json.dumps(result, ensure_ascii=False)}]})
        except Exception as e:
            return _err(-32000, str(e))
    elif method == "resources/list":
        return _ok({
            "resources": [
                {"uri": "nvidia://models", "name": "Models", "mimeType": "application/json"},
                {"uri": "nvidia://sessions", "name": "Sessions", "mimeType": "application/json"},
            ]
        })
    elif method == "resources/read":
        uri = params.get("uri", "")
        if uri == "nvidia://models":
            return _ok({"contents": [{"uri": uri, "mimeType": "application/json",
                                       "text": json.dumps(NVIDIA_MODELS, ensure_ascii=False)}]})
        elif uri == "nvidia://sessions":
            return _ok({"contents": [{"uri": uri, "mimeType": "application/json",
                                       "text": json.dumps(memory.list_sessions())}]})
        return _err(-32602, f"Unknown resource: {uri}")
    return _err(-32601, f"Method not found: {method}")


async def _execute_mcp_tool(name: str, args: dict) -> dict:
    if name == "nvidia_chat":
        sid = args.get("session_id") or str(uuid.uuid4())
        use_mem = args.get("use_memory", True)
        message = args["message"]
        model = args.get("model", DEFAULT_MODEL)
        sysp = args.get("system_prompt", "You are a helpful AI assistant.")

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

        response = await nvidia_chat(
            messages=messages, model=model,
            temperature=args.get("temperature", 0.7),
            max_tokens=args.get("max_tokens", 4096),
        )
        if use_mem:
            memory.add_short(sid, "assistant", response)
        return {"response": response, "model": model, "session_id": sid,
                "memory_stats": memory.get_stats(sid)}

    elif name == "multi_model_compare":
        message = args["message"]
        models = args["models"]
        sysp = args.get("system_prompt", "You are helpful.")
        max_tok = args.get("max_tokens", 512)
        messages = [{"role": "system", "content": sysp}, {"role": "user", "content": message}]

        async with httpx.AsyncClient(timeout=HTTP_TIMEOUT, limits=HTTP_LIMITS) as client:
            tasks = [nvidia_chat(messages=messages, model=m, max_tokens=max_tok, client=client) for m in models]
            responses = await asyncio.gather(*tasks, return_exceptions=True)

        results = {}
        for m, r in zip(models, responses):
            if isinstance(r, Exception):
                results[m] = {"error": str(r)}
            else:
                results[m] = {"response": r, "name": NVIDIA_MODELS.get(m, {}).get("name", m)}
        return {"message": message, "results": results}

    elif name == "list_models":
        category = args.get("category")
        models = {}
        for mid, mi in NVIDIA_MODELS.items():
            if category and mi["category"].lower() != category.lower():
                continue
            models[mid] = mi
        return {"models": models, "total": len(models)}

    elif name == "memory_stats":
        return memory.get_stats(args["session_id"])

    elif name == "memory_clear":
        sid = args["session_id"]
        if args.get("clear_type", "short") == "all":
            memory.delete_session(sid)
        else:
            memory.clear_short(sid)
        return {"status": "cleared", "session_id": sid}

    elif name == "get_conversation_history":
        sid = args["session_id"]
        return {
            "session_id": sid,
            "short_memory": memory.get_short(sid),
            "long_context": memory.get_long_context(sid),
            "stats": memory.get_stats(sid),
        }
    raise ValueError(f"Unknown tool: {name}")


# ══════════════════════════════════════════════════════════════
# API DOCS
# ══════════════════════════════════════════════════════════════

@app.get("/api-docs", response_class=HTMLResponse)
async def api_docs_page():
    return f"""

<!DOCTYPE html><html><head><meta charset="UTF-8"><title>API Docs</title>

<style>

body{{font-family:sans-serif;background:#0a0a0a;color:#e0e0e0;padding:20px;max-width:1200px;margin:auto}}

h1,h2{{color:#76b900}} pre{{background:#16213e;padding:15px;border-radius:8px;overflow-x:auto}}

code{{color:#7ec8e3}} table{{border-collapse:collapse;width:100%}}

th,td{{padding:8px;border:1px solid #333;text-align:left}}

th{{background:#16213e;color:#76b900}}

</style></head><body>

<h1>🚀 NVIDIA AI Multi-Model API v3.1</h1>

<p><b>{len(NVIDIA_MODELS)} modelli VERIFICATI</b> + auto-adattamento system role</p>



<h2>💬 Chat POST</h2>

<pre>curl -X POST /v1/chat -H "Content-Type: application/json" \\

  -d '{{"message":"Ciao!","model":"{DEFAULT_MODEL}"}}'</pre>



<h2>📡 SSE Streaming</h2>

<pre>curl -N "/v1/chat/stream?message=Ciao&model={DEFAULT_MODEL}"</pre>



<h2>⚡ Test Parallelo</h2>

<pre>curl "/v1/debug/test-all-parallel?concurrency=15"</pre>



<h2>🎯 Modelli senza system role (auto-adattati)</h2>

<pre>curl "/v1/debug/no-system-models"</pre>



<h2>🔌 MCP</h2>

<pre>curl -X POST /v1/mcp -H "Content-Type: application/json" \\

  -d '{{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{{}}}}'</pre>



</body></html>

"""


# ══════════════════════════════════════════════════════════════
# GRADIO UI
# ══════════════════════════════════════════════════════════════

def get_model_choices():
    return [
        (f"{mi['name']} [{mi['category']}]", mid)
        for mid, mi in NVIDIA_MODELS.items()
    ]


async def gradio_chat(message, history, model, system_prompt, temperature, max_tokens, session_id, use_memory):
    if not message.strip():
        yield ""
        return

    if not model or not str(model).strip():
        model = DEFAULT_MODEL
    if not session_id:
        session_id = str(uuid.uuid4())[:8]

    if use_memory:
        memory.add_short(session_id, "user", message)
        messages = memory.get_full_context(session_id, system_prompt)
    else:
        messages = [{"role": "system", "content": system_prompt}]
        for h in history:
            if isinstance(h, dict):
                messages.append({"role": h.get("role"), "content": h.get("content", "")})
        messages.append({"role": "user", "content": message})

    full = ""
    try:
        async for chunk in nvidia_chat_stream(
            messages=messages, model=model,
            temperature=temperature, max_tokens=max_tokens,
        ):
            full += chunk
            yield full
    except Exception as e:
        full = f"❌ {str(e)}"
        yield full

    if use_memory and full and not full.startswith("❌"):
        memory.add_short(session_id, "assistant", full)


def get_memory_info(sid):
    return memory.get_stats(sid) if sid else {"error": "Session ID missing"}


def clear_session_memory(sid):
    if not sid: return "⚠️ Session ID missing"
    memory.delete_session(sid)
    return f"🗑️ '{sid}' deleted"


with gr.Blocks(
    title="🚀 NVIDIA AI",
    theme=gr.themes.Soft(primary_hue="green", secondary_hue="blue"),
    css=".gradio-container{max-width:1400px !important;} footer{display:none !important;}"
) as demo:
    gr.Markdown(f"""

    # 🚀 NVIDIA AI Multi-Model v3.1

    **{len(NVIDIA_MODELS)} modelli VERIFICATI** + auto-adattamento system role • Memoria • SSE • MCP



    🏠 [UI](/) | 📖 [Docs](/api-docs) | 🔧 [Swagger](/docs) | 🔌 [MCP](/v1/mcp)



    **Default:** `{DEFAULT_MODEL}` (⚡ 153ms)

    """)

    with gr.Tab("💬 Chat"):
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(height=600, type="messages", show_copy_button=True)
                msg = gr.Textbox(placeholder="Scrivi...", show_label=False, container=False)
            with gr.Column(scale=1):
                model = gr.Dropdown(choices=get_model_choices(), value=DEFAULT_MODEL, label="🤖 Modello")
                session_id = gr.Textbox(value="default", label="🔑 Session ID")
                system_prompt = gr.Textbox(
                    value="Sei un assistente AI utile. Rispondi in italiano.",
                    label="📝 System Prompt", lines=3,
                )
                temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="🌡️ Temp")
                max_tokens = gr.Slider(256, 16384, value=4096, step=256, label="📏 Tokens")
                use_memory = gr.Checkbox(value=True, label="🧠 Memoria")
                with gr.Row():
                    clear_btn = gr.Button("🗑️", size="sm")
                    mem_btn = gr.Button("📊", size="sm")
                mem_output = gr.JSON(label="📊 Memoria")

        msg.submit(gradio_chat,
            inputs=[msg, chatbot, model, system_prompt, temperature, max_tokens, session_id, use_memory],
            outputs=[chatbot]).then(lambda: "", outputs=[msg])
        clear_btn.click(lambda: [], outputs=[chatbot])
        mem_btn.click(get_memory_info, inputs=[session_id], outputs=[mem_output])

    with gr.Tab("🤖 Modelli"):
        gr.Markdown(f"### 📋 Catalogo — {len(NVIDIA_MODELS)} modelli verificati")
        data = [[mi["category"], mi["name"], mid, mi["context"], mi["description"],
                 "❌" if mid in MODELS_NO_SYSTEM_ROLE else "✅"]
                for mid, mi in NVIDIA_MODELS.items()]
        gr.Dataframe(
            headers=["Cat", "Nome", "ID", "Context", "Descrizione", "System Role"],
            value=data, interactive=False, wrap=True
        )

    with gr.Tab("🚀 Test Parallelo"):
        gr.Markdown("### ⚡ Verifica live tutti i modelli in parallelo (con system prompt)")
        test_btn = gr.Button("🚀 Testa tutti i modelli", variant="primary")
        test_output = gr.JSON(label="Risultati")

        async def run_test():
            testable = list(NVIDIA_MODELS.keys())
            sem = asyncio.Semaphore(15)
            start = time.time()
            async with httpx.AsyncClient(timeout=HTTP_TIMEOUT, limits=HTTP_LIMITS) as client:
                tasks = [_test_single(m, client, sem) for m in testable]
                results = await asyncio.gather(*tasks)
            working = sorted([r for r in results if r["status"] == "✅"], key=lambda x: x["elapsed_ms"])
            broken = [r for r in results if r["status"] == "❌"]
            return {
                "⏱️ tempo_sec": round(time.time() - start, 2),
                "✅ funzionanti": len(working),
                "❌ non_funzionanti": len(broken),
                "totale": len(results),
                "models_no_system_role_detected": sorted(MODELS_NO_SYSTEM_ROLE),
                "working_by_speed": working,
                "broken": broken,
            }
        test_btn.click(run_test, outputs=[test_output])

    with gr.Tab("🧠 Memoria"):
        with gr.Row():
            mem_session = gr.Textbox(label="Session ID", value="default")
            with gr.Column():
                stats_btn = gr.Button("📊 Stats")
                clear_mem_btn = gr.Button("🗑️ Elimina", variant="stop")
        mem_display = gr.JSON(label="Risultato")
        stats_btn.click(get_memory_info, inputs=[mem_session], outputs=[mem_display])
        clear_mem_btn.click(clear_session_memory, inputs=[mem_session], outputs=[mem_display])

    with gr.Tab("🎯 System Role"):
        gr.Markdown(f"""

        ## 🎯 Auto-adattamento System Role



        Alcuni modelli (come **Gemma 2**, **Gemma 3n**, **DiffusionGemma**) non supportano il ruolo `system`.



        Questo Space **rileva automaticamente** l'errore e converte il `system_prompt` come **prefisso** del primo messaggio user.



        ### 📋 Modelli attualmente in auto-adattamento:

        """)
        no_sys_display = gr.JSON(label="Modelli senza system role")

        def get_no_sys():
            return {
                "models_no_system_role": sorted(MODELS_NO_SYSTEM_ROLE),
                "total": len(MODELS_NO_SYSTEM_ROLE),
                "note": "Aggiornato dinamicamente quando viene rilevato un errore 'system role not supported'"
            }

        refresh_btn = gr.Button("🔄 Aggiorna lista")
        refresh_btn.click(get_no_sys, outputs=[no_sys_display])
        demo.load(get_no_sys, outputs=[no_sys_display])

    with gr.Tab("📖 Docs"):
        gr.Markdown("""

        ## 📖 Documentazione



        - **[API HTML Docs](/api-docs)**

        - **[Swagger UI](/docs)**

        - **[Info JSON](/info)**



        ### Endpoints principali:

        - `GET /v1/chat/get?message=Ciao` — Chat rapida

        - `POST /v1/chat` — Chat JSON

        - `GET /v1/chat/stream` — SSE streaming

        - `POST /v1/mcp` — MCP JSON-RPC

        - `GET /v1/debug/test-all-parallel` — Test parallelo

        - `GET /v1/debug/no-system-models` — Modelli senza system role

        """)


app = gr.mount_gradio_app(app, demo, path="/ui")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)