colabllm.py

"""
Colab LLM - Servidor de LLM Privado no Google Colab

Orquestrador que levanta:
  1. Ollama (Backend de Inferência)
  2. FastAPI (Gateway de Segurança e Proxy)
  3. Cloudflare Tunnel (Exposição Pública)

Implementa autenticação via Bearer Token e repassa streams compatíveis com OpenAI.

AUTOR: Gabriel Masson
REPOSITORIO: https://github.com/gmasson/colabllm
LICENÇA: MIT
"""

import os
import sys
import time
import hmac
import json
import subprocess
import threading
import logging
import re
from datetime import datetime
from collections import defaultdict
import asyncio

# --- CONFIGURAÇÕES ---

# Defina uma senha segura para proteger seu servidor
API_SECRET_KEY = "sua-senha-super-segura"

# Versão do projeto
VERSION = "0.2"

# Versão do Ollama (atualize quando necessário)
# Releases: https://github.com/ollama/ollama/releases
OLLAMA_VERSION = "0.15.5"

# Parâmetros padrão de inferência
# Estes valores são usados quando não especificados na requisição
DEFAULT_TEMPERATURE = 0.7       # Criatividade (0.0 = determinístico, 1.0 = criativo)
DEFAULT_TOP_P = 0.9             # Nucleus sampling (0.1 = conservador, 1.0 = diverso)
DEFAULT_REPEAT_PENALTY = 1.05   # Penalidade de repetição (1.0 = sem penalidade, evite > 1.1)

# Modelos balanceados para T4 (15GB GPU)
# Formato: nome: {tamanho_aprox_gb, descrição}
MODELS_CONFIG = {
    # --- GENERALISTAS ---
    "llama3.1:8b": {
        "size_gb": 4.9,
        "description": "Meta Llama 3.1 8B - O padrão da indústria. Sólido, confiável e ótimo em PT-BR."
    },
    "qwen2.5:7b": {
        "size_gb": 4.7,
        "description": "Qwen 2.5 7B - O melhor da categoria. Raciocínio lógico superior e PT-BR muito natural."
    },
    "gemma2:9b": {
        "size_gb": 5.4,
        "description": "Google Gemma 2 9B - Criativo e inteligente. Alucina menos e tem raciocínio complexo."
    },

    # --- ESPECIALISTAS EM CÓDIGO E LÓGICA ---
    "qwen2.5-coder:7b": {
        "size_gb": 4.7,
        "description": "Qwen Coder 2.5 7B - O melhor assistente de programação leve. Essencial para devs."
    },
    "deepseek-llm:7b": {
        "size_gb": 4.0,
        "description": "DeepSeek LLM 7B - Focado em matemática e lógica pesada. Ótimo para exatas."
    },
    "yi-coder:9b": {
        "size_gb": 5.0,
        "description": "Yi-Coder 9B - Rival direto do Qwen. Ótimo balanço entre tamanho e performance."
    },
    "deepseek-coder-v2:16b": {
        "size_gb": 8.9,
        "description": "DeepSeek V2 Lite 16B - O mais poderoso. MoE (Mixture of Experts). Atenção: Usa muita VRAM."
    },
    "codellama:13b": {
        "size_gb": 7.4,
        "description": "CodeLlama 13B - Modelo maior e mais estável para arquitetura de software."
    },

    # --- CRIATIVOS, RPG E SEM CENSURA ---
    "hermes3:8b": {
        "size_gb": 4.7,
        "description": "Hermes 3 - Versão 'destravada' do Llama 3.1. Ótimo para RPG e escrita criativa."
    },
    "aya:8b": {
        "size_gb": 4.8,
        "description": "Cohere Aya 8B - Especialista cultural em 23 idiomas. Texto mais humano e fluido."
    },

    # --- PERFORMANCE (Rápidos ou Densos) ---
    "solar:10.7b": {
        "size_gb": 6.1,
        "description": "Solar 10.7B - O maior que cabe confortavelmente. Muito corporativo e inteligente."
    },
    "mistral:7b": {
        "size_gb": 4.4,
        "description": "Mistral v0.3 - Rápido, obediente e com janela de contexto eficiente."
    },
    "llama3.2:3b": {
        "size_gb": 2.0,
        "description": "Llama 3.2 3B - Ultra leve. Respostas instantâneas para tarefas simples."
    },
    "phi3.5:3.8b": {
        "size_gb": 2.2,
        "description": "Microsoft Phi 3.5 - Compacto, mas com raciocínio lógico surpreendente."
    },
    "nemotron-mini:4b": {
        "size_gb": 2.7,
        "description": "NVIDIA Nemotron 4B - Otimizado pela NVIDIA para densidade de conhecimento."
    }
}

# APENAS estes modelos serão baixados e disponibilizados
MODELS_TO_DOWNLOAD = [
    "gemma2:9b",                # Google Gemma 2 9B (Mais criativo/pesado)
    "qwen2.5:7b",               # Qwen 2.5 7B (Melhor raciocínio/PT-BR)
    "deepseek-llm:7b",          # DeepSeek LLM 7B (Exatas/Matemática)
]

OLLAMA_PORT = 11434
PROXY_PORT = 8000

# Rate limiting (requisições por minuto por IP)
RATE_LIMIT_PER_MINUTE = 30
rate_limit_tracker = defaultdict(list)

# Logging
logging.basicConfig(
    level=logging.INFO,
    format='[%(levelname)s] %(asctime)s - %(message)s',
    datefmt='%H:%M:%S'
)
logger = logging.getLogger("LLM_Server")

# Estatísticas globais
stats = {
    "requests_total": 0,
    "requests_success": 0,
    "requests_failed": 0,
    "start_time": datetime.now(),
    "models_available": [],        # Modelos confirmados no disco
    "last_request": None,
    "current_loaded_model": None   # Modelo atualmente na VRAM
}

# ==========================================
# 1. INSTALAÇÃO DE DEPENDÊNCIAS
# ==========================================
def install_dependencies():
    """
    Instala todas as dependências necessárias para o servidor.

    Instala:
    - Bibliotecas Python (fastapi, uvicorn, httpx, nest_asyncio)
    - Ollama (backend de inferência LLM)
    - Cloudflared (túnel público)

    Utiliza métodos redundantes para garantir instalação em ambientes Colab.
    """
    logger.info("Instalando dependências...")

    # 1.1 Pacotes Python
    logger.info("Instalando bibliotecas Python...")
    subprocess.check_call([
        sys.executable, "-m", "pip", "install", "-q",
        "fastapi", "uvicorn", "httpx", "nest_asyncio"
    ])

    # 1.2 Ollama - Método otimizado para Colab
    ollama_path = "/usr/local/bin/ollama"
    if not os.path.exists(ollama_path) or not os.access(ollama_path, os.X_OK):
        logger.info("Instalando Ollama...")
        installed = False

        # Método 1: Script oficial com sudo (mais confiável)
        if not installed:
            try:
                logger.info("Tentando instalação via script oficial...")
                result = subprocess.run(
                    "curl -fsSL https://ollama.com/install.sh | sudo sh",
                    shell=True,
                    capture_output=True,
                    text=True,
                    timeout=180
                )

                # Verifica múltiplos locais possíveis
                possible_paths = [
                    ollama_path,
                    "/usr/bin/ollama",
                    "/usr/local/ollama",
                    os.path.expanduser("~/.ollama/bin/ollama")
                ]
                
                found_path = None
                for path in possible_paths:
                    if os.path.exists(path):
                        found_path = path
                        break
                
                if found_path:
                    # Move para o local padrão se necessário
                    if found_path != ollama_path:
                        subprocess.run(["sudo", "cp", found_path, ollama_path], check=False)
                    subprocess.run(["sudo", "chmod", "+x", ollama_path], check=False)
                    logger.info(f"Ollama instalado via script oficial (encontrado em {found_path})")
                    installed = True
                else:
                    logger.warning("Script executou mas binário não encontrado em locais conhecidos")
            except Exception as e:
                logger.warning(f"Método 1 falhou: {e}")

        # Método 2: Download direto do arquivo .tar.zst (versão específica)
        if not installed:
            try:
                logger.info(f"Tentando download direto v{OLLAMA_VERSION}...")
                download_url = f"https://github.com/ollama/ollama/releases/download/v{OLLAMA_VERSION}/ollama-linux-amd64.tar.zst"
                
                # Instala zstd se necessário (para descompactar .tar.zst)
                subprocess.run("which zstd || sudo apt-get install -y zstd", shell=True, capture_output=True)
                
                # Baixa o arquivo compactado
                result = subprocess.run(
                    f"curl -fsSL -o /tmp/ollama.tar.zst '{download_url}'",
                    shell=True,
                    capture_output=True,
                    text=True,
                    timeout=180
                )

                if result.returncode == 0 and os.path.exists("/tmp/ollama.tar.zst"):
                    # Descompacta com zstd e extrai com tar
                    subprocess.run(
                        "zstd -d /tmp/ollama.tar.zst -o /tmp/ollama.tar && tar -xf /tmp/ollama.tar -C /tmp",
                        shell=True,
                        check=True
                    )
                    
                    # Move o binário para o local correto
                    if os.path.exists("/tmp/bin/ollama"):
                        subprocess.run(["sudo", "mv", "/tmp/bin/ollama", ollama_path], check=True)
                    elif os.path.exists("/tmp/ollama"):
                        subprocess.run(["sudo", "mv", "/tmp/ollama", ollama_path], check=True)
                    
                    subprocess.run(["sudo", "chmod", "+x", ollama_path], check=True)
                    
                    # Limpa arquivos temporários
                    subprocess.run("rm -rf /tmp/ollama.tar.zst /tmp/ollama.tar /tmp/bin", shell=True, check=False)
                    
                    logger.info("Ollama instalado via download direto (.tar.zst)")
                    installed = True
                else:
                    logger.warning(f"Download .tar.zst falhou: {result.stderr[:200]}")
            except Exception as e:
                logger.warning(f"Método 2 falhou: {e}")

        # Método 3: Usando wget como fallback (arquivo .tar.zst)
        if not installed:
            try:
                logger.info("Tentando com wget...")
                download_url = f"https://github.com/ollama/ollama/releases/download/v{OLLAMA_VERSION}/ollama-linux-amd64.tar.zst"
                
                # Instala zstd se necessário
                subprocess.run("which zstd || sudo apt-get install -y zstd", shell=True, capture_output=True)
                
                subprocess.run(
                    f"wget -q -O /tmp/ollama.tar.zst '{download_url}'",
                    shell=True,
                    check=True,
                    timeout=180
                )

                if os.path.exists("/tmp/ollama.tar.zst"):
                    # Descompacta com zstd e extrai com tar
                    subprocess.run(
                        "zstd -d /tmp/ollama.tar.zst -o /tmp/ollama.tar && tar -xf /tmp/ollama.tar -C /tmp",
                        shell=True,
                        check=True
                    )
                    
                    # Move o binário para o local correto
                    if os.path.exists("/tmp/bin/ollama"):
                        subprocess.run(["sudo", "mv", "/tmp/bin/ollama", ollama_path], check=True)
                    elif os.path.exists("/tmp/ollama"):
                        subprocess.run(["sudo", "mv", "/tmp/ollama", ollama_path], check=True)
                    
                    subprocess.run(["sudo", "chmod", "+x", ollama_path], check=True)
                    
                    # Limpa arquivos temporários
                    subprocess.run("rm -rf /tmp/ollama.tar.zst /tmp/ollama.tar /tmp/bin", shell=True, check=False)
                    
                    logger.info("Ollama instalado via wget (.tar.zst)")
                    installed = True
            except Exception as e:
                logger.warning(f"Método 3 falhou: {e}")

        if not installed:
            logger.error("Todos os métodos de instalação falharam")
            raise RuntimeError(
                "Não foi possível instalar Ollama automaticamente.\n"
                "Execute manualmente no Colab:\n"
                "!curl -fsSL https://ollama.com/install.sh | sudo sh"
            )
    else:
        logger.info("Ollama já instalado")

    # Verifica instalação final
    try:
        version = subprocess.check_output([ollama_path, "--version"], text=True)
        logger.info(f"Ollama version: {version.strip()}")
    except Exception as e:
        logger.error(f"Ollama instalado mas não executável: {e}")
        subprocess.run(["sudo", "chmod", "+x", ollama_path], check=False)

    # 1.3 Cloudflared
    cloudflared_path = "/usr/local/bin/cloudflared"
    if not os.path.exists(cloudflared_path):
        logger.info("Instalando Cloudflared...")
        try:
            subprocess.run([
                "wget", "-q",
                "https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64",
                "-O", "/tmp/cloudflared"
            ], check=True, timeout=60)
            subprocess.run(["chmod", "+x", "/tmp/cloudflared"], check=True)
            subprocess.run(["sudo", "mv", "/tmp/cloudflared", cloudflared_path], check=True)
            subprocess.run(["sudo", "chmod", "+x", cloudflared_path], check=True)
            logger.info("Cloudflared instalado")
        except Exception as e:
            logger.error(f"Erro ao instalar Cloudflared: {e}")
            raise
    else:
        logger.info("Cloudflared já instalado")

    logger.info("Todas as dependências verificadas")

install_dependencies()

# ==========================================
# 2. IMPORTS E CONFIGURAÇÕES
# ==========================================
import uvicorn
import nest_asyncio
import httpx
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.middleware.cors import CORSMiddleware

nest_asyncio.apply()

# ==========================================
# 3. OLLAMA - FUNÇÕES AUXILIARES
# ==========================================
def start_ollama_server():
    """
    Inicia o servidor Ollama na porta configurada.

    Configura variáveis de ambiente para otimização:
    - OLLAMA_KEEP_ALIVE: Tempo de retenção do modelo na VRAM
    - OLLAMA_NUM_PARALLEL: Requisições paralelas (1 para T4)
    - OLLAMA_MAX_LOADED_MODELS: Máximo de modelos simultâneos na VRAM

    Returns:
        bool: True se servidor iniciou com sucesso, False caso contrário
    """
    os.environ["OLLAMA_HOST"] = f"127.0.0.1:{OLLAMA_PORT}"
    os.environ["OLLAMA_KEEP_ALIVE"] = "5m"
    os.environ["OLLAMA_NUM_PARALLEL"] = "1"
    os.environ["OLLAMA_MAX_LOADED_MODELS"] = "1"

    ollama_cmd = "/usr/local/bin/ollama"

    # Garante que tem permissão de execução
    if os.path.exists(ollama_cmd):
        try:
            subprocess.run(["sudo", "chmod", "+x", ollama_cmd], check=False)
        except OSError:
            pass

    logger.info("Iniciando Ollama...")

    try:
        subprocess.Popen(
            [ollama_cmd, "serve"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
    except PermissionError:
        logger.error("Erro de permissão ao executar Ollama")
        logger.info("Tentando com sudo...")
        subprocess.Popen(
            ["sudo", ollama_cmd, "serve"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )

    # Aguarda inicialização
    for i in range(30):
        try:
            httpx.get(f"http://127.0.0.1:{OLLAMA_PORT}", timeout=2)
            logger.info("Ollama online")
            return True
        except Exception:
            time.sleep(1)

    logger.error("Falha ao iniciar Ollama")
    return False

def check_model_exists_on_disk(model_name: str) -> bool:
    """
    Verifica se o modelo realmente existe no sistema de arquivos do Ollama.

    Consulta a API /api/tags e confirma a presença física do modelo.

    Args:
        model_name: Nome do modelo a verificar (ex: "gemma2:9b")

    Returns:
        bool: True se modelo existe no disco, False caso contrário
    """
    try:
        response = httpx.get(f"http://127.0.0.1:{OLLAMA_PORT}/api/tags", timeout=5)
        if response.status_code == 200:
            models_data = response.json().get("models", [])
            # Verifica se o nome exato ou alguma variação existe
            for m in models_data:
                model_disk_name = m.get("name", "")
                if model_disk_name == model_name or model_name in model_disk_name:
                    return True
        return False
    except Exception as e:
        logger.warning(f"Erro ao verificar modelo {model_name}: {e}")
        return False

def test_model_functionality(model_name: str) -> bool:
    """
    Testa se o modelo está funcional fazendo uma requisição simples.

    Envia uma mensagem de teste para verificar se o modelo responde corretamente.
    Detecta modelos corrompidos ou com problemas de inicialização.

    Args:
        model_name: Nome do modelo a testar

    Returns:
        bool: True se modelo responde corretamente, False caso contrário
    """
    logger.info(f"Testando funcionalidade do modelo {model_name}...")

    try:
        with httpx.Client(timeout=60.0) as client:
            response = client.post(
                f"http://127.0.0.1:{OLLAMA_PORT}/api/generate",
                json={
                    "model": model_name,
                    "prompt": "Hi",
                    "stream": False
                }
            )

            if response.status_code == 200:
                data = response.json()
                # Verifica se recebeu uma resposta válida
                if "response" in data and len(data["response"]) > 0:
                    logger.info(f"Modelo {model_name} respondeu corretamente")
                    return True
                else:
                    logger.error(f"Modelo {model_name} retornou resposta vazia")
                    return False
            else:
                error_text = response.text[:200]
                logger.error(f"Modelo {model_name} retornou erro {response.status_code}: {error_text}")
                return False

    except Exception as e:
        logger.error(f"Erro ao testar modelo {model_name}: {str(e)}")
        return False

def remove_model(model_name: str):
    """
    Remove um modelo corrompido do sistema.

    Args:
        model_name: Nome do modelo a remover
    """
    logger.info(f"Removendo modelo {model_name}...")
    ollama_cmd = "/usr/local/bin/ollama"

    try:
        subprocess.run(
            [ollama_cmd, "rm", model_name],
            capture_output=True,
            text=True,
            timeout=30
        )
        logger.info(f"Modelo {model_name} removido")
    except Exception as e:
        logger.warning(f"Erro ao remover {model_name}: {e}")

def pull_models():
    """
    Baixa modelos definidos em MODELS_TO_DOWNLOAD.

    Executa download sequencial durante o boot do servidor.
    Verifica espaço disponível e confirma instalação real de cada modelo.
    TESTA funcionalidade de cada modelo após download.
    Atualiza stats APENAS com modelos confirmados e funcionais.
    """
    total_size = sum(MODELS_CONFIG.get(m, {}).get("size_gb", 0) for m in MODELS_TO_DOWNLOAD)
    logger.info(f"Baixando {len(MODELS_TO_DOWNLOAD)} modelo(s) (~{total_size:.1f}GB)...")

    if total_size > 14:
        logger.warning(f"Total de {total_size:.1f}GB pode exceder VRAM disponível!")

    ollama_cmd = "/usr/local/bin/ollama"

    for model in MODELS_TO_DOWNLOAD:
        config = MODELS_CONFIG.get(model, {})
        desc = config.get("description", "Sem descrição")
        size = config.get("size_gb", "?")

        logger.info(f"-> {model} (~{size}GB) - {desc}")

        # Verifica se modelo já existe e está funcional
        if check_model_exists_on_disk(model):
            logger.info(f"Modelo {model} já existe no disco, testando funcionalidade...")
            if test_model_functionality(model):
                stats["models_available"].append(model)
                continue
            else:
                logger.warning(f"Modelo {model} existe mas está corrompido, re-baixando...")
                remove_model(model)

        # Faz download
        try:
            result = subprocess.run(
                [ollama_cmd, "pull", model],
                capture_output=True,
                text=True,
                timeout=600
            )

            if result.returncode == 0:
                # Verificação 1: Confirma se o modelo aparece na lista do Ollama
                if check_model_exists_on_disk(model):
                    # Verificação 2: Testa se o modelo está funcional
                    if test_model_functionality(model):
                        logger.info(f"Modelo {model} pronto, verificado e testado")
                        stats["models_available"].append(model)
                    else:
                        logger.error(f"Modelo {model} baixado mas não está funcional")
                        logger.error(f"   Tente re-executar o script ou remover este modelo da lista")
                else:
                    logger.error(f"Modelo {model} não apareceu na lista após download")
            else:
                logger.warning(f"Erro ao baixar {model}: {result.stderr[:100]}")
        except subprocess.TimeoutExpired:
            logger.warning(f"Timeout ao baixar {model} (>10min)")
        except Exception as e:
            logger.warning(f"Erro inesperado com {model}: {str(e)}")

    logger.info(f"Download concluído - {len(stats['models_available'])} modelo(s) disponível(is) e funcional(is)")

async def unload_model(model_name: str):
    """
    Força o descarregamento completo de um modelo da VRAM.

    Envia requisição com keep_alive=0 para liberar memória imediatamente.
    Essencial para evitar Out of Memory (OOM) ao trocar modelos.

    Args:
        model_name: Nome do modelo a descarregar da memória
    """
    if not model_name:
        return

    logger.info(f"Descarregando modelo da VRAM: {model_name}")

    async with httpx.AsyncClient(timeout=10.0) as client:
        try:
            # Envia requisição vazia com keep_alive=0
            await client.post(
                f"http://127.0.0.1:{OLLAMA_PORT}/api/generate",
                json={
                    "model": model_name,
                    "keep_alive": 0
                }
            )
            logger.info(f"Comando de unload enviado para {model_name}")

            # Aguarda um pouco para o Ollama processar
            await asyncio.sleep(2)

        except Exception as e:
            logger.warning(f"Falha ao descarregar {model_name}: {e}")

def check_gpu_memory():
    """
    Verifica uso atual de memória GPU via nvidia-smi.

    Returns:
        dict: Dicionário com used_mb, total_mb e percent, ou None se falhar
    """
    try:
        output = subprocess.check_output(
            'nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader,nounits',
            shell=True
        ).decode().strip()
        used, total = map(int, output.split(','))
        percent = (used / total) * 100
        return {"used_mb": used, "total_mb": total, "percent": percent}
    except Exception:
        return None

# ==========================================
# 4. FASTAPI GATEWAY
# ==========================================
app = FastAPI(docs_url=None, redoc_url=None)

# CORS deve ser configurado primeiro
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"]
)

security = HTTPBearer()

def check_rate_limit(client_ip: str) -> bool:
    """
    Implementa rate limiting simples baseado em IP.

    Limita requisições por minuto conforme RATE_LIMIT_PER_MINUTE.
    Remove automaticamente requisições antigas da janela de tempo.

    Args:
        client_ip: Endereço IP do cliente

    Returns:
        bool: True se dentro do limite, False se excedeu
    """
    now = time.time()
    # Remove requisições antigas (>60s)
    rate_limit_tracker[client_ip] = [
        req_time for req_time in rate_limit_tracker[client_ip]
        if now - req_time < 60
    ]

    if len(rate_limit_tracker[client_ip]) >= RATE_LIMIT_PER_MINUTE:
        return False

    rate_limit_tracker[client_ip].append(now)
    return True

async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """
    Valida token de autenticação Bearer.

    Args:
        credentials: Credenciais HTTP Bearer extraídas do header

    Raises:
        HTTPException: Se token for inválido (401)

    Returns:
        str: Token validado
    """
    if not hmac.compare_digest(credentials.credentials, API_SECRET_KEY):
        logger.warning("Acesso negado: token inválido")
        raise HTTPException(status_code=401, detail="Invalid API Key")
    return credentials.credentials

@app.get("/health")
async def health_check():
    """
    Endpoint público de health check (sem autenticação).

    Retorna status do servidor, uptime, modelos disponíveis e uso de GPU.

    Returns:
        dict: Status do servidor e métricas básicas
    """
    gpu_info = check_gpu_memory()
    uptime = (datetime.now() - stats["start_time"]).total_seconds()

    return {
        "status": "healthy",
        "uptime_seconds": int(uptime),
        "models_available": len(stats["models_available"]),
        "current_model_loaded": stats["current_loaded_model"],
        "gpu_memory": gpu_info,
        "version": VERSION
    }

@app.get("/v1/stats")
async def get_stats(token: str = Depends(verify_token)):
    """
    Retorna estatísticas detalhadas do servidor (requer autenticação).

    Inclui métricas de requisições, modelos, uso de GPU e última atividade.

    Args:
        token: Token Bearer validado

    Returns:
        dict: Estatísticas completas do servidor
    """
    uptime = (datetime.now() - stats["start_time"]).total_seconds()
    success_rate = (stats["requests_success"] / max(stats["requests_total"], 1)) * 100

    return {
        "uptime_seconds": int(uptime),
        "requests": {
            "total": stats["requests_total"],
            "success": stats["requests_success"],
            "failed": stats["requests_failed"],
            "success_rate": f"{success_rate:.1f}%"
        },
        "models": {
            "available": stats["models_available"],
            "current_in_vram": stats["current_loaded_model"]
        },
        "last_request": stats["last_request"],
        "gpu": check_gpu_memory()
    }

@app.post("/v1/models/{model_name}/repair")
async def repair_model(model_name: str, token: str = Depends(verify_token)):
    """
    Re-baixa um modelo específico que esteja corrompido (requer autenticação).

    Remove e baixa novamente o modelo, testando sua funcionalidade.

    Args:
        model_name: Nome do modelo a reparar
        token: Token Bearer validado

    Returns:
        JSONResponse: Status da operação de reparo
    """
    # Verifica se modelo está na lista de downloads permitidos
    if model_name not in MODELS_TO_DOWNLOAD:
        raise HTTPException(
            status_code=400,
            detail=f"Modelo '{model_name}' não está em MODELS_TO_DOWNLOAD"
        )

    logger.info(f"Iniciando reparo do modelo {model_name}...")

    # Remove modelo corrompido
    remove_model(model_name)

    # Remove da lista de disponíveis se estiver lá
    if model_name in stats["models_available"]:
        stats["models_available"].remove(model_name)

    # Re-baixa
    ollama_cmd = "/usr/local/bin/ollama"
    try:
        result = subprocess.run(
            [ollama_cmd, "pull", model_name],
            capture_output=True,
            text=True,
            timeout=600
        )

        if result.returncode == 0:
            # Testa funcionalidade
            if test_model_functionality(model_name):
                stats["models_available"].append(model_name)
                logger.info(f"Modelo {model_name} reparado com sucesso")
                return JSONResponse({
                    "status": "success",
                    "message": f"Modelo {model_name} reparado e testado com sucesso"
                })
            else:
                logger.error(f"Modelo {model_name} re-baixado mas ainda não funcional")
                return JSONResponse({
                    "status": "error",
                    "message": f"Modelo {model_name} foi baixado mas ainda apresenta problemas"
                }, status_code=500)
        else:
            logger.error(f"Erro ao re-baixar {model_name}")
            return JSONResponse({
                "status": "error",
                "message": f"Erro ao baixar modelo: {result.stderr[:200]}"
            }, status_code=500)

    except subprocess.TimeoutExpired:
        return JSONResponse({
            "status": "error",
            "message": "Timeout ao baixar modelo (>10min)"
        }, status_code=500)
    except Exception as e:
        return JSONResponse({
            "status": "error",
            "message": f"Erro inesperado: {str(e)}"
        }, status_code=500)

@app.get("/v1/models")
async def list_models(token: str = Depends(verify_token)):
    """
    Lista APENAS modelos que foram baixados em MODELS_TO_DOWNLOAD (requer autenticação).

    Verifica em tempo real com a API do Ollama e retorna metadados.

    Args:
        token: Token Bearer validado

    Returns:
        JSONResponse: Lista de modelos no formato OpenAI

    Raises:
        HTTPException: Se houver erro ao consultar Ollama (500)
    """
    try:
        # Monta resposta com modelos disponíveis
        models_with_info = []
        for m_name in stats["models_available"]:
            config = MODELS_CONFIG.get(m_name, {})
            models_with_info.append({
                "id": m_name,
                "object": "model",
                "owned_by": "ollama",
                "description": config.get("description", "Modelo LLM"),
                "size_gb": config.get("size_gb", 0)
            })

        return JSONResponse({
            "object": "list",
            "data": models_with_info
        })

    except Exception as e:
        logger.error(f"Erro ao listar modelos: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/chat/completions")
async def chat_completions(request: Request, token: str = Depends(verify_token)):
    """
    Proxy para chat completions compatível com OpenAI (requer autenticação).

    Funcionalidades:
    - Rate limiting por IP
    - Validação de modelo disponível
    - Unload inteligente para evitar OOM
    - Streaming compatível com OpenAI

    Args:
        request: Objeto Request do FastAPI
        token: Token Bearer validado

    Returns:
        StreamingResponse: Stream de resposta do modelo

    Raises:
        HTTPException:
            - 429 se rate limit excedido
            - 400 se JSON inválido ou falta parâmetro model
            - 404 se modelo não está disponível
    """
    client_ip = request.client.host

    # Rate limiting
    if not check_rate_limit(client_ip):
        logger.warning(f"Rate limit excedido para IP: {client_ip}")
        raise HTTPException(status_code=429, detail="Rate limit exceeded")

    try:
        body = await request.json()
    except (ValueError, TypeError):
        stats["requests_total"] += 1
        stats["requests_failed"] += 1
        raise HTTPException(status_code=400, detail="Invalid JSON")

    stats["requests_total"] += 1
    stats["last_request"] = datetime.now().isoformat()

    requested_model = body.get('model')
    if not requested_model:
        stats["requests_failed"] += 1
        raise HTTPException(status_code=400, detail="Model parameter required")

    # --- VALIDAÇÃO DE MODELO DISPONÍVEL ---
    # Verifica se modelo está na lista de disponíveis
    if requested_model not in stats["models_available"]:
        stats["requests_failed"] += 1
        logger.warning(f"Modelo não disponível: {requested_model}")
        logger.warning(f"Modelos disponíveis: {', '.join(stats['models_available'])}")
        raise HTTPException(
            status_code=404,
            detail=f"Modelo '{requested_model}' não está disponível. Modelos disponíveis: {', '.join(stats['models_available'])}"
        )

    # Se for um modelo diferente do atual, descarrega o antigo
    current_model = stats.get("current_loaded_model")
    if current_model and requested_model != current_model:
        logger.info(f"Troca de modelo: {current_model} -> {requested_model}")
        await unload_model(current_model)

    # Atualiza modelo atual
    stats["current_loaded_model"] = requested_model
    # --------------------------------

    # Aplica valores padrão se não especificados na requisição
    if 'temperature' not in body:
        body['temperature'] = DEFAULT_TEMPERATURE
    if 'top_p' not in body:
        body['top_p'] = DEFAULT_TOP_P
    if 'repeat_penalty' not in body:
        body['repeat_penalty'] = DEFAULT_REPEAT_PENALTY

    is_stream = body.get('stream', False)
    logger.info(f"Chat | IP: {client_ip} | Model: {requested_model} | Stream: {is_stream}")

    async def stream_response():
        """Generator para streaming de resposta do Ollama"""
        # Timeout granular: conexão rápida, leitura generosa para modelos grandes
        timeout = httpx.Timeout(connect=30.0, read=600.0, write=30.0, pool=30.0)
        async with httpx.AsyncClient(timeout=timeout) as client:
            try:
                async with client.stream(
                    "POST",
                    f"http://127.0.0.1:{OLLAMA_PORT}/v1/chat/completions",
                    json=body
                ) as resp:
                    if resp.status_code != 200:
                        stats["requests_failed"] += 1
                        error_text = await resp.aread()
                        error_decoded = error_text.decode('utf-8', errors='ignore')
                        logger.error(f"Ollama error {resp.status_code}: {error_decoded[:500]}")
                        yield json.dumps({"error": f"Ollama error {resp.status_code}"}).encode()
                        return

                    async for chunk in resp.aiter_bytes():
                        yield chunk

                    stats["requests_success"] += 1

            except httpx.TimeoutException:
                stats["requests_failed"] += 1
                logger.error("Timeout ao processar requisição no Ollama")
                yield b'{"error": "Request timeout"}'
            except Exception as e:
                stats["requests_failed"] += 1
                logger.error(f"Stream error: {str(e)}")
                yield json.dumps({"error": "Internal server error"}).encode()

    return StreamingResponse(stream_response(), media_type="text/event-stream")

def run_fastapi():
    """
    Executa servidor FastAPI em thread separada.

    Configura Uvicorn com host 0.0.0.0 para aceitar conexões externas.
    Log level em 'warning' para reduzir verbosidade.
    """
    config = uvicorn.Config(
        app,
        host="0.0.0.0",
        port=PROXY_PORT,
        log_level="warning"
    )
    server = uvicorn.Server(config)
    server.run()

# ==========================================
# 5. CLOUDFLARE TUNNEL
# ==========================================
def start_cloudflared():
    """
    Inicia túnel público Cloudflare e exibe informações de acesso.

    Cria túnel temporário (trycloudflare.com) que expõe servidor local.
    Exibe URLs de acesso, exemplos de uso e informações dos modelos.
    """
    logger.info("Iniciando túnel Cloudflare...")

    subprocess.run("pkill cloudflared", shell=True, stderr=subprocess.DEVNULL)

    process = subprocess.Popen(
        ["cloudflared", "tunnel", "--url", f"http://localhost:{PROXY_PORT}"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )

    url = None
    for _ in range(30):
        line = process.stderr.readline()
        if "trycloudflare.com" in line:
            match = re.search(r'https://[a-zA-Z0-9-]+\.trycloudflare\.com', line)
            if match:
                url = match.group(0)
                break
        time.sleep(0.5)

    if url:
        print("\n" + "="*70)
        print("SERVIDOR LLM ONLINE!")
        print("="*70)
        print(f"\nENDPOINTS DISPONIVEIS:")
        print(f"   Base URL:        {url}/v1")
        print(f"   Chat:            {url}/v1/chat/completions")
        print(f"   Models:          {url}/v1/models")
        print(f"   Repair Model:    {url}/v1/models/{{model_name}}/repair")
        print(f"   Health (público):{url}/health")
        print(f"   Stats (auth):    {url}/v1/stats")

        print(f"\nAUTENTICACAO:")
        print(f"   API Key: {API_SECRET_KEY}")
        print(f"   Header:  Authorization: Bearer {API_SECRET_KEY}")

        print(f"\nMODELOS DISPONIVEIS ({len(stats['models_available'])} total):")
        if stats['models_available']:
            for model in stats['models_available']:
                config = MODELS_CONFIG.get(model, {})
                size = config.get('size_gb', '?')
                desc = config.get('description', 'Sem descrição')
                print(f"   {model:<20} ({size}GB) - {desc}")
        else:
            print("   NENHUM MODELO DISPONIVEL - Verifique os logs de download")

        print(f"\nMONITORAMENTO:")
        print(f"   Rate Limit:   {RATE_LIMIT_PER_MINUTE} requisições/minuto por IP")
        gpu = check_gpu_memory()
        if gpu:
            print(f"   GPU VRAM:     {gpu['used_mb']}MB / {gpu['total_mb']}MB ({gpu['percent']:.1f}%)")
        print(f"   Modelo atual: {stats['current_loaded_model'] or 'Nenhum carregado'}")

        # --- EXEMPLO 1: JSON PARA POSTMAN/INSOMNIA ---
        first_model = stats['models_available'][0] if stats['models_available'] else "modelo-nao-disponivel"
        print(f"\n" + "="*70)
        print("EXEMPLO 1: POSTMAN / INSOMNIA")
        print("="*70)
        print(f"\nEndpoint:")
        print(f"   POST {url}/v1/chat/completions")
        print(f"\nHeaders:")
        print(f"   Authorization: Bearer {API_SECRET_KEY}")
        print(f"   Content-Type: application/json")
        print(f"\nBody (JSON):")
        print(f"""{{
  "model": "{first_model}",
  "messages": [
    {{
      "role": "system",
      "content": "Você é um assistente útil e educado."
    }},
    {{
      "role": "user",
      "content": "Olá! Explique brevemente o que é inteligência artificial."
    }}
  ],
  "temperature": 0.7,
  "max_tokens": 500
}}""")

        # --- EXEMPLO 2: CURL PARA TERMINAL ---
        print(f"\n" + "="*70)
        print("EXEMPLO 2: TERMINAL (CURL)")
        print("="*70)
        print(f"\n# Teste simples (sem streaming)")
        print(f"curl -X POST {url}/v1/chat/completions \\")
        print(f"  -H \"Authorization: Bearer {API_SECRET_KEY}\" \\")
        print(f"  -H \"Content-Type: application/json\" \\")
        print(f"  -d '{{")
        print(f"    \"model\": \"{first_model}\",")
        print(f"    \"messages\": [")
        print(f"      {{\"role\": \"user\", \"content\": \"Olá! Como você está?\"}}")
        print(f"    ]")
        print(f"  }}'")

        print(f"\n# Com streaming (resposta em tempo real)")
        print(f"curl -X POST {url}/v1/chat/completions \\")
        print(f"  -H \"Authorization: Bearer {API_SECRET_KEY}\" \\")
        print(f"  -H \"Content-Type: application/json\" \\")
        print(f"  -d '{{")
        print(f"    \"model\": \"{first_model}\",")
        print(f"    \"messages\": [")
        print(f"      {{\"role\": \"user\", \"content\": \"Conte uma história curta\"}}")
        print(f"    ],")
        print(f"    \"stream\": true")
        print(f"  }}'")

        # --- EXEMPLO 3: PYTHON ---
        print(f"\n" + "="*70)
        print("EXEMPLO 3: PYTHON (OpenAI Library)")
        print("="*70)
        print(f"""
from openai import OpenAI

# Configura cliente apontando para seu servidor
client = OpenAI(
    base_url="{url}/v1",
    api_key="{API_SECRET_KEY}"
)

# Faz requisição
response = client.chat.completions.create(
    model="{first_model}",
    messages=[
        {{"role": "system", "content": "Você é um assistente útil."}},
        {{"role": "user", "content": "Explique Python em 3 linhas"}}
    ],
    temperature=0.7
)

print(response.choices[0].message.content)

# Com streaming
for chunk in client.chat.completions.create(
    model="{first_model}",
    messages=[{{"role": "user", "content": "Conte uma piada"}}],
    stream=True
):
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="")
""")

        # --- EXEMPLO 4: JAVASCRIPT ---
        print(f"\n" + "="*70)
        print("EXEMPLO 4: JAVASCRIPT (Fetch API)")
        print("="*70)
        print(f"""
const response = await fetch('{url}/v1/chat/completions', {{
  method: 'POST',
  headers: {{
    'Authorization': 'Bearer {API_SECRET_KEY}',
    'Content-Type': 'application/json'
  }},
  body: JSON.stringify({{
    model: '{first_model}',
    messages: [
      {{ role: 'user', content: 'Olá! Como funciona JavaScript?' }}
    ]
  }})
}});

const data = await response.json();
console.log(data.choices[0].message.content);
""")

        # --- INFORMAÇÕES FINAIS ---
        print(f"\n" + "="*70)
        print("INFORMACOES IMPORTANTES")
        print("="*70)
        print(f"• Apenas modelos em MODELS_TO_DOWNLOAD estão disponíveis")
        print(f"• Não há download automático de modelos não listados")
        print(f"• Sistema faz swap automático de modelos para economizar VRAM")
        print(f"• Use /health para verificar status sem autenticação")
        print(f"• Use /v1/stats para ver métricas detalhadas (requer auth)")
        print(f"• Use /v1/models para listar todos os modelos disponíveis")

        print(f"\nREPARAR MODELO CORROMPIDO:")
        print(f"   Se um modelo apresentar erros, você pode repará-lo com:")
        print(f"\n   curl -X POST {url}/v1/models/gemma2:9b/repair \\")
        print(f"     -H \"Authorization: Bearer {API_SECRET_KEY}\"")
        print(f"\n   Isso irá remover e re-baixar o modelo, testando sua funcionalidade.")

        print(f"\nServidor pronto para uso!")
        print("="*70 + "\n")
    else:
        logger.error("Não foi possível obter URL do túnel Cloudflare")

# ==========================================
# 6. EXECUÇÃO PRINCIPAL
# ==========================================
if __name__ == "__main__":
    # Verifica GPU
    try:
        gpu = subprocess.check_output(
            'nvidia-smi --query-gpu=name --format=csv,noheader',
            shell=True
        ).decode().strip()
        logger.info(f"GPU detectada: {gpu}")
    except Exception:
        logger.warning("GPU não detectada - performance será limitada")

    # Inicia serviços
    if not start_ollama_server():
        logger.error("Falha crítica ao iniciar Ollama")
        sys.exit(1)

    pull_models()

    # Inicia threads
    threading.Thread(target=run_fastapi, daemon=True).start()
    time.sleep(3)
    threading.Thread(target=start_cloudflared, daemon=True).start()

    # Mantém processo vivo
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        logger.info("Encerrando servidor...")