"""Generation on the consolidated GPU backend.

Loads Qwen2.5-1.5B-Instruct + the whosouravsharma/paper-qa-lora adapter once
at startup (module-level state, populated by the FastAPI lifespan hook in
app.py). One model instance serves both the fine-tuned behavior (default)
and the raw base model's behavior via peft's disable_adapter() context, so
only one copy of the weights lives in memory.

A single GPU can only run one generate() call at a time, but that must not
block the whole server (health checks, other in-flight requests) while one
generation is running. So actual generation is serialized through a single
asyncio worker task consuming a queue, with the blocking model.generate()
call itself pushed into a one-worker thread executor — the async endpoint
in app.py just awaits its own future. This replaces the earlier archived
design's plain threading.Lock(), which under FastAPI's sync-route threadpool
could let a /health check queue up behind an in-flight generation instead of
returning immediately.
"""

import asyncio
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
ADAPTER_REPO = "whosouravsharma/paper-qa-lora"


def generator_model_name(use_adapter: bool) -> str:
    return f"{BASE_MODEL}+LoRA({ADAPTER_REPO})" if use_adapter else BASE_MODEL

SYSTEM_PROMPT = (
    "You are a research paper assistant. Answer the question using only the "
    "provided context. Quote the exact sentence(s) you relied on as a "
    "citation. If the context does not contain the answer, say so plainly "
    "instead of guessing."
)

state: dict = {}
_executor = ThreadPoolExecutor(max_workers=1)  # one GPU: exactly one worker
_queue: "asyncio.Queue | None" = None
_worker_task: "asyncio.Task | None" = None


@dataclass
class GenerateRequest:
    context: str
    question: str
    use_adapter: bool = True
    max_new_tokens: int = 256


def load_model() -> None:
    """Called once from app.py's lifespan startup."""
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    # Qwen2.5's config defaults to bfloat16; force fp16 explicitly rather than
    # inheriting that default — this exact mismatch caused real training
    # crashes on T4 (see paper-qa-lora/PLAN.md's debugging notes), so it's
    # forced here too even though this is inference, not training.
    base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.float16).to("cuda")
    model = PeftModel.from_pretrained(base, ADAPTER_REPO)
    model.eval()
    state["tokenizer"] = tokenizer
    state["model"] = model


def _blocking_generate(req: GenerateRequest) -> tuple[str, float]:
    tokenizer = state["tokenizer"]
    model = state["model"]

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Context:\n{req.context}\n\nQuestion: {req.question}"},
    ]
    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
    ).to(model.device)

    start = time.perf_counter()
    with torch.no_grad():
        if req.use_adapter:
            output = model.generate(
                **inputs, max_new_tokens=req.max_new_tokens, do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )
        else:
            with model.disable_adapter():
                output = model.generate(
                    **inputs, max_new_tokens=req.max_new_tokens, do_sample=False,
                    pad_token_id=tokenizer.eos_token_id,
                )
    latency_ms = (time.perf_counter() - start) * 1000

    answer = tokenizer.decode(
        output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
    ).strip()
    return answer, latency_ms


async def _worker_loop() -> None:
    loop = asyncio.get_running_loop()
    while True:
        req, future = await _queue.get()
        try:
            result = await loop.run_in_executor(_executor, _blocking_generate, req)
            if not future.done():
                future.set_result(result)
        except Exception as e:  # noqa: BLE001 - surface any generation failure to the caller
            if not future.done():
                future.set_exception(e)
        finally:
            _queue.task_done()


def start_worker() -> None:
    """Called once from app.py's lifespan startup, after load_model()."""
    global _queue, _worker_task
    _queue = asyncio.Queue()
    _worker_task = asyncio.create_task(_worker_loop())


async def stop_worker() -> None:
    if _worker_task is not None:
        _worker_task.cancel()


async def generate_answer(req: GenerateRequest) -> tuple[str, float]:
    """Enqueues one generation request and awaits its result. Safe to call
    concurrently — the queue serializes actual GPU work while this
    coroutine itself doesn't block the event loop."""
    future = asyncio.get_running_loop().create_future()
    await _queue.put((req, future))
    return await future