"""
OMNI-NEXUS · Worker de cómputo elástico
=======================================
Envuelve llama.cpp en una API compatible con OpenAI para que Odysseus lo
vea como un endpoint local más, sin modificar nada del proyecto.

Endpoints:
    GET  /health              → estado (lo usa el healthcheck de HF)
    GET  /v1/models           → catálogo (lo usa Odysseus al añadir el endpoint)
    POST /v1/chat/completions → inferencia, con y sin streaming
    GET  /warm                → despierta el Space sin gastar inferencia

Nota sobre /warm: los Spaces gratuitos se suspenden tras un rato inactivos y
tardan ~60-90 s en volver. El router del gateway llama a /warm para levantar
el siguiente worker mientras el actual trabaja, así el usuario no percibe el
arranque en frío.
"""

from __future__ import annotations

import json
import os
import time
import uuid
from typing import Any, AsyncIterator, Dict, List, Optional

import httpx
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

LLAMA_URL = f"http://127.0.0.1:{os.getenv('LLAMA_PORT', '8080')}"
WORKER_ID = os.getenv("WORKER_ID", "A")
N_CTX = int(os.getenv("N_CTX", "4096"))
TIMEOUT = float(os.getenv("LLAMA_TIMEOUT", "180.0"))
MODEL_NAME = os.getenv("MODEL_NAME", "omni-nexus/worker")

# Plantilla de chat. BitNet b1.58-2B-4T se entrenó con el formato de Llama 3,
# no con ChatML. Usar la plantilla equivocada degrada mucho la calidad sin dar
# ningún error visible — el modelo simplemente responde peor. Por eso es una
# variable explícita y no una suposición.
CHAT_TEMPLATE = os.getenv("CHAT_TEMPLATE", "llama3").lower()

app = FastAPI(title=f"Omni-Nexus Worker {WORKER_ID}", version="1.0.0")

# CORS abierto: el endpoint se consume desde el navegador de Odysseus.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

_stats = {"requests": 0, "tokens": 0, "started_at": time.time()}


# ---------------------------------------------------------------------------
# Modelos de entrada (subconjunto del esquema OpenAI que Odysseus usa)
# ---------------------------------------------------------------------------
class Message(BaseModel):
    role: str
    content: Any = ""


class ChatRequest(BaseModel):
    model: Optional[str] = None
    messages: List[Message]
    temperature: float = 0.7
    top_p: float = 0.95
    max_tokens: Optional[int] = Field(default=1024)
    stream: bool = False
    stop: Optional[List[str]] = None


# ---------------------------------------------------------------------------
def _flatten(content: Any) -> str:
    """El contenido puede venir como string o como lista de bloques."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                parts.append(block.get("text", ""))
            elif isinstance(block, str):
                parts.append(block)
        return "\n".join(parts)
    return str(content or "")


def _prompt_llama3(messages: List[Message]) -> str:
    """Formato Llama 3 — el correcto para BitNet b1.58-2B-4T."""
    out = ["<|begin_of_text|>"]
    for m in messages:
        role = m.role if m.role in ("system", "user", "assistant") else "user"
        out.append(
            f"<|start_header_id|>{role}<|end_header_id|>\n\n"
            f"{_flatten(m.content)}<|eot_id|>"
        )
    out.append("<|start_header_id|>assistant<|end_header_id|>\n\n")
    return "".join(out)


def _prompt_chatml(messages: List[Message]) -> str:
    """Formato ChatML — Qwen 2.5 y buena parte de los instruct actuales."""
    out = []
    for m in messages:
        role = m.role if m.role in ("system", "user", "assistant") else "user"
        out.append(f"<|im_start|>{role}\n{_flatten(m.content)}<|im_end|>")
    out.append("<|im_start|>assistant\n")
    return "\n".join(out)


def to_prompt(messages: List[Message]) -> str:
    if CHAT_TEMPLATE == "chatml":
        return _prompt_chatml(messages)
    return _prompt_llama3(messages)


# Se incluyen los stops de ambas plantillas: sobra un par de tokens en la
# lista y evita que un cambio de modelo deje basura al final de la respuesta.
STOP_TOKENS = [
    "<|eot_id|>", "<|start_header_id|>", "<|end_header_id|>",
    "<|im_end|>", "<|im_start|>", "</s>", "<|end_of_text|>",
]


# ---------------------------------------------------------------------------
@app.get("/health")
async def health():
    llama_ok = False
    try:
        async with httpx.AsyncClient(timeout=5.0) as c:
            r = await c.get(f"{LLAMA_URL}/health")
            llama_ok = r.status_code == 200
    except Exception:
        pass
    return {
        "status": "ok" if llama_ok else "starting",
        "worker": WORKER_ID,
        "llama": llama_ok,
        "ctx": N_CTX,
        "template": CHAT_TEMPLATE,
        "kind": os.getenv("WORKER_KIND", "chat"),
        "uptime_s": round(time.time() - _stats["started_at"]),
        "requests": _stats["requests"],
    }


@app.get("/warm")
async def warm():
    """Despierta el Space sin gastar inferencia. Responde apenas el proceso vive."""
    return {"worker": WORKER_ID, "awake": True, "ts": time.time()}


@app.get("/v1/models")
async def models():
    return {
        "object": "list",
        "data": [
            {
                "id": MODEL_NAME,
                "object": "model",
                "created": int(_stats["started_at"]),
                "owned_by": f"omni-nexus-worker-{WORKER_ID}",
                "context_length": N_CTX,
            }
        ],
    }


@app.post("/v1/chat/completions")
async def chat(req: ChatRequest):
    prompt = to_prompt(req.messages)
    payload = {
        "prompt": prompt,
        "temperature": req.temperature,
        "top_p": req.top_p,
        "n_predict": req.max_tokens or 1024,
        "stop": STOP_TOKENS + (req.stop or []),
        "stream": req.stream,
        "cache_prompt": True,  # reutiliza el KV cache entre turnos: gran ahorro
    }
    _stats["requests"] += 1

    if req.stream:
        return StreamingResponse(
            _stream(payload, req.model), media_type="text/event-stream"
        )

    try:
        async with httpx.AsyncClient(timeout=TIMEOUT) as c:
            r = await c.post(f"{LLAMA_URL}/completion", json=payload)
            r.raise_for_status()
            data = r.json()
    except httpx.TimeoutException:
        raise HTTPException(504, detail=f"Worker {WORKER_ID}: timeout de inferencia")
    except Exception as exc:
        raise HTTPException(502, detail=f"Worker {WORKER_ID}: {exc}")

    text = (data.get("content") or "").strip()
    for tok in STOP_TOKENS:
        text = text.replace(tok, "")

    n_out = data.get("tokens_predicted", 0)
    n_in = data.get("tokens_evaluated", 0)
    _stats["tokens"] += n_out

    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": req.model or MODEL_NAME,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": text.strip()},
                "finish_reason": "stop",
            }
        ],
        "usage": {
            "prompt_tokens": n_in,
            "completion_tokens": n_out,
            "total_tokens": n_in + n_out,
        },
        "_worker": WORKER_ID,
    }


async def _stream(payload: Dict[str, Any], model: Optional[str]) -> AsyncIterator[str]:
    cid = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    created = int(time.time())

    def frame(delta: Dict[str, Any], finish: Optional[str] = None) -> str:
        return "data: " + json.dumps(
            {
                "id": cid,
                "object": "chat.completion.chunk",
                "created": created,
                "model": model or MODEL_NAME,
                "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
            }
        ) + "\n\n"

    yield frame({"role": "assistant", "content": ""})

    try:
        async with httpx.AsyncClient(timeout=TIMEOUT) as c:
            async with c.stream("POST", f"{LLAMA_URL}/completion", json=payload) as r:
                async for line in r.aiter_lines():
                    if not line.startswith("data: "):
                        continue
                    try:
                        chunk = json.loads(line[6:])
                    except json.JSONDecodeError:
                        continue
                    piece = chunk.get("content", "")
                    if piece:
                        _stats["tokens"] += 1
                        yield frame({"content": piece})
                    if chunk.get("stop"):
                        break
    except Exception as exc:
        yield frame({"content": f"\n[worker {WORKER_ID}: {exc}]"})

    yield frame({}, finish="stop")
    yield "data: [DONE]\n\n"


@app.get("/")
async def root():
    return {
        "name": f"Omni-Nexus Worker {WORKER_ID}",
        "usage": "Añádelo en Odysseus → Settings → Add Local Models con la URL /v1",
        "endpoints": ["/health", "/warm", "/v1/models", "/v1/chat/completions"],
    }