"""
Gemma 4 Supa-Fast Dual-API Server for HuggingFace Spaces (CPU, 16GB RAM)
Uses llama.cpp (GGUF) instead of transformers/PyTorch.
OpenAI: POST /v1/chat/completions
Anthropic: POST /v1/messages
"""

from __future__ import annotations

import os
import sys
import json
import time
import uuid
import asyncio
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import AsyncGenerator, Dict, List, Literal, Optional, Union, Any

import huggingface_hub
from llama_cpp import Llama, CreateChatCompletionStreamResponse
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field

# ─────────────────────────── Configuration ───────────────────────────

MODEL_REPO = os.getenv("MODEL_REPO", "ggml-org/gemma-4-E4B-it-GGUF")
QUANT = os.getenv("QUANT", "Q5_K_M")          # Q4_K_M is slightly faster; Q5_K_M is higher quality
MAX_CTX = int(os.getenv("MAX_CTX", "8192"))
N_THREADS = int(os.getenv("N_THREADS", str(os.cpu_count() or 2)))
N_BATCH = int(os.getenv("N_BATCH", "512"))
PORT = int(os.getenv("PORT", "7860"))

GGUF_FILENAME = f"gemma-4-E4B-it-{QUANT}.gguf"

# ─────────────────────────── Global State ────────────────────────────

@dataclass
class ModelState:
    llm: Optional[Llama] = None
    ready: bool = False
    load_error: Optional[str] = None
    load_time_sec: float = 0.0


STATE = ModelState()


# ─────────────────────────── Model Loading ───────────────────────────

def load_model() -> None:
    """Download GGUF (if needed) and load into llama.cpp."""
    global STATE
    start = time.time()
    print(f"[INIT] Repo: {MODEL_REPO}")
    print(f"[INIT] Quant: {QUANT}")
    print(f"[INIT] Threads: {N_THREADS} | Ctx: {MAX_CTX} | Batch: {N_BATCH}")

    try:
        print("[INIT] Downloading / verifying GGUF...")
        model_path = huggingface_hub.hf_hub_download(
            repo_id=MODEL_REPO,
            filename=GGUF_FILENAME,
            local_files_only=False,
        )
        print(f"[INIT] GGUF path: {model_path}")

        print("[INIT] Loading model into llama.cpp (mmap + mlock)...")
        llm = Llama(
            model_path=model_path,
            n_ctx=MAX_CTX,
            n_threads=N_THREADS,
            n_batch=N_BATCH,
            use_mmap=True,      # memory-mapped = instant load, OS pages in on demand
            use_mlock=True,     # prevent swapping to disk
            verbose=False,
        )

        elapsed = time.time() - start
        STATE.llm = llm
        STATE.ready = True
        STATE.load_time_sec = elapsed

        print(f"[INIT] Ready in {elapsed:.1f}s")
        print(f"[INIT] Model size: ~{os.path.getsize(model_path) / 1e9:.1f} GB on disk")

    except Exception as e:
        STATE.load_error = str(e)
        STATE.ready = False
        print(f"[ERROR] {e}")
        import traceback
        traceback.print_exc()


# ─────────────────────────── Pydantic Schemas ────────────────────────

class OAIChatMessage(BaseModel):
    role: Literal["system", "user", "assistant", "tool"] = "user"
    content: str = ""
    name: Optional[str] = None


class OAIChatRequest(BaseModel):
    model: str = "gemma-4"
    messages: List[OAIChatMessage]
    temperature: Optional[float] = Field(default=0.7, ge=0.0, le=2.0)
    top_p: Optional[float] = Field(default=0.9, ge=0.0, le=1.0)
    max_tokens: Optional[int] = Field(default=2048, ge=1, le=8192)
    stream: bool = False
    stop: Optional[Union[str, List[str]]] = None


class OAIChoice(BaseModel):
    index: int = 0
    message: dict = Field(default_factory=dict)
    finish_reason: Optional[str] = None


class OAIChunkChoice(BaseModel):
    index: int = 0
    delta: dict = Field(default_factory=dict)
    finish_reason: Optional[str] = None


class OAIUsage(BaseModel):
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class OAIChatResponse(BaseModel):
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = "gemma-4"
    choices: List[OAIChoice]
    usage: OAIUsage


class OAIChatStreamChunk(BaseModel):
    id: str = ""
    object: str = "chat.completion.chunk"
    created: int = 0
    model: str = "gemma-4"
    choices: List[OAIChunkChoice]


# Anthropic schemas
class AnthropicMessageParam(BaseModel):
    role: Literal["user", "assistant"] = "user"
    content: str = ""


class AnthropicChatRequest(BaseModel):
    model: str = "claude-sonnet-4"
    messages: List[AnthropicMessageParam]
    max_tokens: int = Field(default=4096, ge=1, le=8192)
    system: Optional[str] = None
    temperature: Optional[float] = Field(default=0.7, ge=0.0, le=1.0)
    top_p: Optional[float] = Field(default=0.9, ge=0.0, le=1.0)
    stream: bool = False
    stop_sequences: Optional[List[str]] = None


class AnthropicUsage(BaseModel):
    input_tokens: int = 0
    output_tokens: int = 0


class AnthropicMessageResponse(BaseModel):
    id: str = Field(default_factory=lambda: f"msg_{uuid.uuid4().hex[:24]}")
    type: str = "message"
    role: str = "assistant"
    model: str = "gemma-4"
    content: List[dict]
    stop_reason: Optional[str] = "end_turn"
    usage: AnthropicUsage


# ──────────────────────── Format Conversion ──────────────────────────

def _to_llama_messages(
    messages: List[OAIChatMessage],
    system: Optional[str] = None,
) -> List[Dict[str, str]]:
    """Convert to llama.cpp chat format."""
    out: List[Dict[str, str]] = []
    if system:
        out.append({"role": "system", "content": system})
    for m in messages:
        out.append({"role": m.role, "content": m.content})
    return out


# ─────────────────────────── FastAPI App ─────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("[LIFESPAN] Starting up...")
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, load_model)
    yield
    print("[LIFESPAN] Shutting down...")


app = FastAPI(
    title="Gemma 4 Supa-Fast Server",
    description="llama.cpp GGUF backend — OpenAI + Anthropic endpoints",
    version="2.0.0",
    lifespan=lifespan,
)

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


@app.get("/")
def root():
    return {
        "status": "ok" if STATE.ready else "loading",
        "model": MODEL_REPO,
        "quant": QUANT,
        "version": "2.0.0",
        "endpoints": [
            "POST /v1/chat/completions",
            "POST /v1/messages",
            "GET  /health",
        ],
    }


@app.get("/health")
def health():
    return {
        "status": "healthy" if STATE.ready else "loading",
        "model_loaded": STATE.ready,
        "load_time_sec": STATE.load_time_sec,
        "error": STATE.load_error,
    }


@app.get("/v1/models")
def list_models():
    return {
        "object": "list",
        "data": [
            {
                "id": "gemma-4",
                "object": "model",
                "created": int(time.time()),
                "owned_by": "google",
            }
        ],
    }


# ── OpenAI-compatible endpoint ──

@app.post("/v1/chat/completions")
async def openai_chat_completions(request: OAIChatRequest):
    if not STATE.ready:
        raise HTTPException(
            status_code=503,
            detail=f"Model not loaded. Error: {STATE.load_error}",
        )

    messages = _to_llama_messages(request.messages)
    stop = request.stop or []
    if isinstance(stop, str):
        stop = [stop]

    gen_kwargs = dict(
        messages=messages,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p,
        stop=stop,
        stream=request.stream,
    )

    if request.stream:
        return StreamingResponse(
            _openai_stream(gen_kwargs),
            media_type="text/event-stream",
        )

    # Non-streaming
    t0 = time.time()
    output = STATE.llm.create_chat_completion(**gen_kwargs)
    elapsed = time.time() - t0

    choice = output["choices"][0]
    usage = output.get("usage", {})

    return OAIChatResponse(
        choices=[
            OAIChoice(
                index=0,
                message=choice["message"],
                finish_reason=choice.get("finish_reason", "stop"),
            )
        ],
        usage=OAIUsage(
            prompt_tokens=usage.get("prompt_tokens", 0),
            completion_tokens=usage.get("completion_tokens", 0),
            total_tokens=usage.get("total_tokens", 0),
        ),
    )


async def _openai_stream(gen_kwargs: dict):
    """SSE stream for OpenAI format."""
    req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    created = int(time.time())

    # Role chunk
    first = OAIChatStreamChunk(
        id=req_id,
        created=created,
        choices=[OAIChunkChoice(index=0, delta={"role": "assistant"}, finish_reason=None)],
    )
    yield f"data: {first.model_dump_json()}\n\n"

    content_buf = ""
    prompt_tokens = 0
    completion_tokens = 0

    # llama.cpp streaming generator
    stream = STATE.llm.create_chat_completion(**gen_kwargs)

    for chunk in stream:
        delta = chunk["choices"][0]["delta"]
        finish = chunk["choices"][0].get("finish_reason")

        if delta.get("content"):
            content_buf += delta["content"]
            completion_tokens += 1
            out_chunk = OAIChatStreamChunk(
                id=req_id,
                created=created,
                choices=[OAIChunkChoice(index=0, delta={"content": delta["content"]}, finish_reason=None)],
            )
            yield f"data: {out_chunk.model_dump_json()}\n\n"

        if finish:
            final = OAIChatStreamChunk(
                id=req_id,
                created=created,
                choices=[OAIChunkChoice(index=0, delta={}, finish_reason=finish)],
            )
            yield f"data: {final.model_dump_json()}\n\n"
            break

    usage = {
        "id": req_id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": "gemma-4",
        "choices": [],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
    }
    yield f"data: {json.dumps(usage)}\n\n"
    yield "data: [DONE]\n\n"


# ── Anthropic-compatible endpoint ──

@app.post("/v1/messages")
async def anthropic_messages(request: AnthropicChatRequest):
    if not STATE.ready:
        raise HTTPException(status_code=503, detail=f"Model not loaded. {STATE.load_error}")

    # Flatten Anthropic format → OpenAI format → llama.cpp format
    messages: List[Dict[str, str]] = []
    if request.system:
        messages.append({"role": "system", "content": request.system})
    for m in request.messages:
        messages.append({"role": m.role, "content": m.content})

    stop = request.stop_sequences or []
    gen_kwargs = dict(
        messages=messages,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p,
        stop=stop,
        stream=request.stream,
    )

    if request.stream:
        return StreamingResponse(
            _anthropic_stream(gen_kwargs),
            media_type="text/event-stream",
            headers={"x-api-version": "2023-06-01"},
        )

    output = STATE.llm.create_chat_completion(**gen_kwargs)
    text = output["choices"][0]["message"]["content"]
    usage = output.get("usage", {})

    return AnthropicMessageResponse(
        content=[{"type": "text", "text": text}],
        stop_reason="end_turn",
        usage=AnthropicUsage(
            input_tokens=usage.get("prompt_tokens", 0),
            output_tokens=usage.get("completion_tokens", 0),
        ),
    )


async def _anthropic_stream(gen_kwargs: dict):
    msg_id = f"msg_{uuid.uuid4().hex[:24]}"

    yield f"event: message_start\ndata: {json.dumps({'type': 'message_start', 'message': {'id': msg_id, 'type': 'message', 'role': 'assistant', 'model': 'gemma-4', 'content': [], 'stop_reason': None, 'stop_sequence': None, 'usage': {'input_tokens': 0, 'output_tokens': 0}}})}\n\n"
    yield f"event: content_block_start\ndata: {json.dumps({'type': 'content_block_start', 'index': 0, 'content_block': {'type': 'text', 'text': ''}})}\n\n"

    stream = STATE.llm.create_chat_completion(**gen_kwargs)
    output_tokens = 0

    for chunk in stream:
        delta = chunk["choices"][0]["delta"]
        if delta.get("content"):
            output_tokens += 1
            payload = {
                "type": "content_block_delta",
                "index": 0,
                "delta": {"type": "text_delta", "text": delta["content"]},
            }
            yield f"event: content_block_delta\ndata: {json.dumps(payload)}\n\n"

    yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
    yield f"event: message_delta\ndata: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': 'end_turn', 'stop_sequence': None}, 'usage': {'output_tokens': output_tokens}})}\n\n"
    yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"


# ─────────────────────────── Entry Point ─────────────────────────────

if __name__ == "__main__":
    import uvicorn
    print("=" * 60)
    print("  Gemma 4 Supa-Fast Server (llama.cpp GGUF)")
    print("  OpenAI: POST /v1/chat/completions")
    print("  Anthropic: POST /v1/messages")
    print(f"  Repo: {MODEL_REPO} | Quant: {QUANT}")
    print(f"  Port: {PORT}")
    print("=" * 60)
    uvicorn.run("app:app", host="0.0.0.0", port=PORT, workers=1)