"""
Qwen2.5-Coder-7B on Hugging Face ZeroGPU
-----------------------------------------
- Serves an OpenAI-compatible /v1/chat/completions endpoint (+ /v1/models)
- Protects the API with a Bearer API key (via HF Space Secret, not hardcoded)
- Rate-limits to N requests/minute per API key (default 5)
- Ships a small Gradio admin panel (also gated by the same API key) to
  swap the active model without redeploying

Architecture note: ZeroGPU only schedules GPU time for functions that run
*inside the Gradio process* (functions decorated with @spaces.GPU). So the
FastAPI layer below does NOT do inference itself -- it validates the
request, then calls into a plain Python function that is decorated with
@spaces.GPU. That's what actually grabs the GPU for the few seconds it
needs, then releases it. This is required by ZeroGPU, not a stylistic
choice.
"""

import os
import time
import threading
import uuid
from collections import defaultdict, deque
from typing import List, Optional

import gradio as gr
import spaces
import torch
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

# ---------------------------------------------------------------------------
# Config (read from environment / HF Space Secrets -- never hardcode these)
# ---------------------------------------------------------------------------

# Set this in your Space's Settings -> Variables and secrets -> New secret
# Name: API_KEY   Value: <a long random string you generate yourself>
API_KEY = os.environ.get("API_KEY")

# How many requests per minute a caller is allowed to make (per API key).
RATE_LIMIT_PER_MINUTE = int(os.environ.get("RATE_LIMIT_PER_MINUTE", "5"))

# Default model. Can be changed at runtime from the admin panel (in-memory
# only -- resets to this default if the Space restarts/sleeps).
DEFAULT_MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-Coder-7B-Instruct")

# Max seconds a single generation is allowed to hold the GPU. ZeroGPU kills
# the call if it runs longer than this, so keep some headroom.
GPU_DURATION_SECONDS = int(os.environ.get("GPU_DURATION_SECONDS", "90"))

if not API_KEY:
    # Fail loudly at boot rather than silently serving an unauthenticated API.
    raise RuntimeError(
        "API_KEY environment variable is not set. "
        "Go to your Space's Settings -> Variables and secrets and add a "
        "secret named API_KEY before this Space will start."
    )

# ---------------------------------------------------------------------------
# Model state (mutable so the admin panel can hot-swap it)
# ---------------------------------------------------------------------------

class ModelState:
    """Holds whatever model/tokenizer is currently loaded, plus a lock so a
    hot-swap from the admin panel can't race with an in-flight request."""

    def __init__(self, model_id: str):
        self.lock = threading.Lock()
        self.model_id = None
        self.model = None
        self.tokenizer = None
        self.load(model_id)

    def load(self, model_id: str):
        with self.lock:
            print(f"[model] loading {model_id} ...")
            tokenizer = AutoTokenizer.from_pretrained(model_id)
            model = AutoModelForCausalLM.from_pretrained(
                model_id,
                torch_dtype=torch.bfloat16,
                device_map=None,  # ZeroGPU: do NOT move to cuda here, see note below
            )
            # NOTE: we intentionally do not call .to("cuda") outside of a
            # @spaces.GPU-decorated function. The `spaces` package patches
            # torch so CUDA calls made inside a @spaces.GPU function get
            # routed to the GPU that's allocated for that call. Moving the
            # model to cuda at *load* time (module scope) is the officially
            # supported pattern -- see the generate() function below, where
            # .to("cuda") happens lazily on first GPU-decorated call.
            self.model_id = model_id
            self.tokenizer = tokenizer
            self.model = model
            print(f"[model] loaded {model_id}")

    def swap(self, new_model_id: str) -> str:
        if new_model_id == self.model_id:
            return f"'{new_model_id}' is already the active model."
        try:
            self.load(new_model_id)
            return f"Switched active model to '{new_model_id}'."
        except Exception as e:
            return f"Failed to load '{new_model_id}': {e}"


state = ModelState(DEFAULT_MODEL_ID)

# ---------------------------------------------------------------------------
# Rate limiting: simple in-memory sliding window, keyed by API key.
# Good enough for a single-instance Space. Not distributed -- fine here
# since a Space is one process.
# ---------------------------------------------------------------------------

_rate_lock = threading.Lock()
_request_log: dict[str, deque] = defaultdict(deque)


def check_rate_limit(key: str):
    now = time.monotonic()
    window = 60.0
    with _rate_lock:
        dq = _request_log[key]
        while dq and now - dq[0] > window:
            dq.popleft()
        if len(dq) >= RATE_LIMIT_PER_MINUTE:
            retry_after = max(0, window - (now - dq[0]))
            raise HTTPException(
                status_code=429,
                detail=f"Rate limit exceeded: {RATE_LIMIT_PER_MINUTE} requests/minute. "
                       f"Retry after {retry_after:.1f}s.",
                headers={"Retry-After": str(int(retry_after) + 1)},
            )
        dq.append(now)


def require_api_key(request: Request):
    auth = request.headers.get("authorization", "")
    if not auth.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing bearer token.")
    token = auth.removeprefix("Bearer ").strip()
    if token != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key.")
    check_rate_limit(token)
    return token


# ---------------------------------------------------------------------------
# GPU-bound generation. This is the ONLY function that touches CUDA, and
# it's the only one ZeroGPU actually schedules a GPU for.
# ---------------------------------------------------------------------------

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def generate(messages: List[dict], max_new_tokens: int, temperature: float, top_p: float) -> str:
    model = state.model
    tokenizer = state.tokenizer

    if not model.parameters().__next__().is_cuda:
        model.to("cuda")

    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=temperature > 0,
            temperature=max(temperature, 1e-5),
            top_p=top_p,
            pad_token_id=tokenizer.eos_token_id,
        )

    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
    text = tokenizer.decode(new_tokens, skip_special_tokens=True)

    # Must return CPU/plain Python data -- returning CUDA tensors across the
    # ZeroGPU process boundary breaks (see ZeroGPU pickling constraints).
    return text


# ---------------------------------------------------------------------------
# OpenAI-compatible schema (trimmed to what's actually used)
# ---------------------------------------------------------------------------

class ChatMessage(BaseModel):
    role: str
    content: str


class ChatCompletionRequest(BaseModel):
    model: Optional[str] = None
    messages: List[ChatMessage]
    max_tokens: Optional[int] = 512
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    stream: Optional[bool] = False  # streaming not implemented, see note below


# ---------------------------------------------------------------------------
# FastAPI app, mounted into the Gradio app below
# ---------------------------------------------------------------------------

api = FastAPI(title="Qwen2.5-Coder ZeroGPU API")


@api.get("/v1/models")
def list_models(_=None):
    return {
        "object": "list",
        "data": [{"id": state.model_id, "object": "model", "owned_by": "you"}],
    }


@api.post("/v1/chat/completions")
def chat_completions(payload: ChatCompletionRequest, request: Request):
    require_api_key(request)

    if payload.stream:
        raise HTTPException(
            status_code=400,
            detail="stream=true is not implemented on this endpoint yet. "
                   "Set stream=false in your OpenAI SDK call.",
        )

    messages = [m.model_dump() for m in payload.messages]

    try:
        text = generate(
            messages=messages,
            max_new_tokens=payload.max_tokens or 512,
            temperature=payload.temperature if payload.temperature is not None else 0.7,
            top_p=payload.top_p or 0.9,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Generation failed: {e}")

    now = int(time.time())
    return JSONResponse({
        "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
        "object": "chat.completion",
        "created": now,
        "model": state.model_id,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": text},
            "finish_reason": "stop",
        }],
        # Token usage isn't tracked precisely here; fill with -1 as a signal
        # that these are not real counts, rather than a plausible-looking
        # fake number an SDK might use for billing math.
        "usage": {"prompt_tokens": -1, "completion_tokens": -1, "total_tokens": -1},
    })


# ---------------------------------------------------------------------------
# Gradio admin panel -- gated behind the SAME API key.
# Gradio itself becomes the public page (Space landing page), so we gate
# access at the component level: nothing useful renders until the correct
# key is submitted.
# ---------------------------------------------------------------------------

CANDIDATE_MODELS = [
    "Qwen/Qwen2.5-Coder-7B-Instruct",
    "Qwen/Qwen2.5-Coder-1.5B-Instruct",
    "Qwen/Qwen2.5-Coder-3B-Instruct",
    "Qwen/Qwen2.5-Coder-14B-Instruct",
]


def try_unlock(key: str):
    if key == API_KEY:
        return (
            gr.update(visible=True),                      # admin_panel
            gr.update(visible=False),                      # login_panel
            gr.update(value=f"Current model: {state.model_id}"),
        )
    return (
        gr.update(visible=False),
        gr.update(visible=True),
        gr.update(),
    )


def do_swap(model_choice: str, custom_model_id: str):
    target = custom_model_id.strip() or model_choice
    msg = state.swap(target)
    return msg, f"Current model: {state.model_id}"


with gr.Blocks(title="Qwen2.5-Coder ZeroGPU API") as demo:
    gr.Markdown(
        "# Qwen2.5-Coder — ZeroGPU API\n"
        "This Space exposes an OpenAI-compatible API at `/v1/chat/completions`.\n"
        "It requires a Bearer API key (set as the `API_KEY` Space secret).\n\n"
        "The panel below is for the Space owner only, to switch the active model."
    )

    with gr.Group(visible=True) as login_panel:
        gr.Markdown("### Admin login")
        key_box = gr.Textbox(label="API key", type="password")
        login_btn = gr.Button("Unlock")

    with gr.Group(visible=False) as admin_panel:
        gr.Markdown("### Admin panel")
        status_box = gr.Markdown()
        model_dropdown = gr.Dropdown(
            choices=CANDIDATE_MODELS, label="Switch to a known model", value=None
        )
        custom_model_box = gr.Textbox(
            label="...or enter any Hugging Face model ID (overrides dropdown)",
            placeholder="e.g. Qwen/Qwen2.5-Coder-32B-Instruct",
        )
        swap_btn = gr.Button("Load this model", variant="primary")
        swap_result = gr.Markdown()

    login_btn.click(
        try_unlock,
        inputs=[key_box],
        outputs=[admin_panel, login_panel, status_box],
    )
    swap_btn.click(
        do_swap,
        inputs=[model_dropdown, custom_model_box],
        outputs=[swap_result, status_box],
    )

# Mount the Gradio admin UI onto the FastAPI app, under /admin. The
# OpenAI-compatible routes (/v1/chat/completions, /v1/models) stay directly
# on `api` at the root, since that's the path your OpenAI SDK base_url
# needs to hit. Gradio's own login/model-swap UI lives at /admin instead of
# "/", so it doesn't shadow the API routes.
app = gr.mount_gradio_app(api, demo, path="/admin")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))