"""LLM clients: HFTransformersClient (transformers/AWQ, T4, greedy) plus
NullClient/CallableClient for dev and an optional VLLMClient."""

from __future__ import annotations

import re
from pathlib import Path
from typing import Callable, List, Optional, Sequence, Tuple

MAX_NEW_TOKENS = 1024
DEFAULT_MODEL_DIR = "."  # submission: weights ship at the repo root


class LLMClient:
    """Interface: generate(prompts) -> one completion per prompt."""

    available: bool = False
    # can this backend score candidate next-tokens (needed for the match_letters
    # assignment solver)? Only the real transformers backend can.
    can_score: bool = False
    deadline: Optional[float] = None  # monotonic wall-clock abort (armed by caller)

    def generate(self, prompts: Sequence[str], max_new_tokens: int = MAX_NEW_TOKENS,
                 system: Optional[str] = None, sample: bool = False,
                 temperature: float = 0.5) -> List[str]:
        raise NotImplementedError


class NullClient(LLMClient):
    available = False

    def generate(self, prompts, max_new_tokens=MAX_NEW_TOKENS, system=None,
                 sample=False, temperature=0.5):
        return ["" for _ in prompts]


class CallableClient(LLMClient):
    available = True

    def __init__(self, fn: Callable[[str], str]):
        self.fn = fn

    def generate(self, prompts, max_new_tokens=MAX_NEW_TOKENS, system=None,
                 sample=False, temperature=0.5):
        return [self.fn(p) for p in prompts]


def pack_by_tokens(lengths: Sequence[int], token_budget: int, max_batch: int
                   ) -> List[List[int]]:
    """Group prompt indices into batches whose TOTAL token count stays under
    `token_budget`. On the sandbox's transformers version, prefill computes
    float32 logits over every prompt position (batch x seq x 152k vocab), so
    total-tokens-per-batch — not batch count — is what bounds T4 memory.
    An oversized single prompt still gets its own batch (handled by the OOM
    retry path)."""
    batches: List[List[int]] = []
    cur: List[int] = []
    cur_tokens = 0
    for i, n in enumerate(lengths):
        if cur and (cur_tokens + n > token_budget or len(cur) >= max_batch):
            batches.append(cur)
            cur, cur_tokens = [], 0
        cur.append(i)
        cur_tokens += n
    if cur:
        batches.append(cur)
    return batches


class HFTransformersClient(LLMClient):
    """transformers backend, mirroring the official notebook's loading code
    (fp16, device_map=auto, greedy) plus token-budget batching and OOM
    recovery. Import is lazy."""

    available = True
    can_score = True  # supports score_next_logprobs (match_letters assignment)

    # total prompt tokens per generation batch. On the sandbox's transformers,
    # prefill computes fp32 logits over every prompt position (batch x seq x
    # 152k vocab), ~0.9 MB/token; total-tokens-per-batch — not batch count — is
    # the T4 memory bound. 3500 is the safe fallback used when we cannot
    # measure free VRAM; __init__ raises it to fit whatever headroom the loaded
    # model actually leaves (≈6500 for a 7B-AWQ, ≈3500 for a 14B-AWQ).
    TOKEN_BUDGET = 3500
    MB_PER_TOKEN = 0.9      # fp32 prefill logits at Qwen's 152k vocab
    VRAM_RESERVE_GB = 1.8   # KV cache + activations + fragmentation slack

    def __init__(self, model_dir: str = DEFAULT_MODEL_DIR, batch_size: int = 8):
        import inspect
        import time as _time

        import torch
        from transformers import (AutoModelForCausalLM, AutoTokenizer,
                                  StoppingCriteria, StoppingCriteriaList)

        self.torch = torch
        self.tok = AutoTokenizer.from_pretrained(model_dir)
        # Pin to GPU 0: device_map="auto" can silently offload layers to CPU on a
        # tight T4 (~100x slower). Fall back to "auto" if the pinned load fails.
        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                model_dir, torch_dtype=torch.float16,
                device_map={"": 0} if torch.cuda.is_available() else "auto",
            ).eval()
        except Exception:
            self.model = AutoModelForCausalLM.from_pretrained(
                model_dir, torch_dtype=torch.float16, device_map="auto",
            ).eval()
        if self.tok.pad_token_id is None:
            self.tok.pad_token = self.tok.eos_token
        self.batch_size = batch_size
        # newer transformers can skip full-sequence prefill logits entirely
        fwd_params = inspect.signature(self.model.forward).parameters
        self._logits_kwarg = next(
            (k for k in ("logits_to_keep", "num_logits_to_keep") if k in fwd_params),
            None)
        self._tune_token_budget()
        self.last_truncated: List[bool] = []

        # per-token wall-clock abort: a batch started near the deadline can't
        # overrun and get the process killed (the budget is otherwise only
        # checked between batches). set self.deadline (monotonic ts) to arm it.
        self.deadline: Optional[float] = None

        class _Deadline(StoppingCriteria):
            def __call__(self, input_ids, scores, **kw):
                return _time.monotonic() > deadline_holder[0]

        deadline_holder = [float("inf")]
        self._deadline_holder = deadline_holder
        self._deadline_crit = StoppingCriteriaList([_Deadline()])

    def _tune_token_budget(self) -> None:
        """Size the per-batch token budget to the VRAM the loaded weights
        actually leave free. Falls back to the safe class default if the GPU
        can't be queried (CPU dev, older CUDA)."""
        torch = self.torch
        try:
            if not torch.cuda.is_available():
                return
            free_bytes, _ = torch.cuda.mem_get_info()
            free_gb = free_bytes / 1e9
            budget = int((free_gb - self.VRAM_RESERVE_GB) * 1000 / self.MB_PER_TOKEN)
            # the logits_to_keep fast path removes the big allocation entirely,
            # but stay conservative regardless; clamp to a sane window
            self.TOKEN_BUDGET = max(2500, min(7000, budget))
        except Exception:
            pass

    def _chat_texts(self, prompts: Sequence[str], system: Optional[str]) -> List[str]:
        texts = []
        for p in prompts:
            messages = ([{"role": "system", "content": system}] if system else []) \
                + [{"role": "user", "content": p}]
            texts.append(self.tok.apply_chat_template(
                messages, add_generation_prompt=True, tokenize=False))
        return texts

    def _generate_chunk(self, chunk: List[str], max_new_tokens: int,
                        sample: bool = False, temperature: float = 0.5
                        ) -> Tuple[List[str], List[bool]]:
        torch = self.torch
        # pad only when batching more than one prompt (padding can perturb a
        # quantized model's greedy outputs)
        enc = self.tok(chunk, return_tensors="pt", padding=len(chunk) > 1,
                       add_special_tokens=False).to(self.model.device)
        kwargs = {self._logits_kwarg: 1} if self._logits_kwarg else {}
        # repetition_penalty=1.0 explicitly: the shipped generation_config sets
        # 1.05, which is applied even under greedy and biases against the
        # repeated characters common in these answers.
        kwargs["repetition_penalty"] = 1.0
        if sample:
            kwargs.update(do_sample=True, temperature=temperature, top_p=0.95)
        else:
            kwargs["do_sample"] = False
        if self.deadline is not None:
            self._deadline_holder[0] = self.deadline
            kwargs["stopping_criteria"] = self._deadline_crit
        with torch.no_grad():
            gen = self.model.generate(
                **enc, max_new_tokens=max_new_tokens,
                pad_token_id=self.tok.pad_token_id, **kwargs,
            )
        new_tokens = gen[:, enc["input_ids"].shape[1]:]
        eos = self.tok.eos_token_id
        texts, truncated = [], []
        for row in new_tokens:
            texts.append(_guard(self.tok.decode(row, skip_special_tokens=True)).strip())
            # no EOS in the generated span => generation was cut at the cap
            truncated.append(eos is None or int((row == eos).sum()) == 0)
        return texts, truncated

    def generate(self, prompts, max_new_tokens=MAX_NEW_TOKENS, system=None,
                 sample=False, temperature=0.5):
        if not prompts:
            self.last_truncated = []
            return []
        torch = self.torch
        texts = self._chat_texts(prompts, system)
        lengths = [len(self.tok(t, add_special_tokens=False)["input_ids"])
                   for t in texts]
        out: List[str] = [""] * len(texts)
        trunc: List[bool] = [False] * len(texts)
        prev_side = self.tok.padding_side
        self.tok.padding_side = "left"  # sequences must end at the gen position
        try:
            for batch in pack_by_tokens(lengths, self.TOKEN_BUDGET, self.batch_size):
                chunk = [texts[i] for i in batch]
                try:
                    results, tflags = self._generate_chunk(
                        chunk, max_new_tokens, sample, temperature)
                except torch.cuda.OutOfMemoryError:
                    # halve pressure: clear cache, retry one prompt at a time;
                    # a prompt that OOMs alone yields "" (symbolic answer stands)
                    torch.cuda.empty_cache()
                    results, tflags = [], []
                    for t in chunk:
                        try:
                            r1, f1 = self._generate_chunk(
                                [t], max_new_tokens, sample, temperature)
                            results.append(r1[0])
                            tflags.append(f1[0])
                        except torch.cuda.OutOfMemoryError:
                            torch.cuda.empty_cache()
                            results.append("")
                            tflags.append(False)
                for i, r, f in zip(batch, results, tflags):
                    out[i] = r
                    trunc[i] = f
        finally:
            self.tok.padding_side = prev_side
        self.last_truncated = trunc
        return out

    def score_next_logprobs(self, prompts, cand_token_ids, system=None,
                            batch_size=4):
        """For each prompt, return the max next-token log-prob over each
        candidate group. `cand_token_ids` is a list of token-id groups (one per
        option), shared across prompts. Returns List[List[float]]
        (prompt x option). One forward pass per batch; used by the match_letters
        assignment solver. Missing/OOM/timed-out prompts get all-zero rows so
        the caller can fall back."""
        import time
        if not prompts:
            return []
        torch = self.torch
        n_opt = len(cand_token_ids)
        texts = self._chat_texts(prompts, system)
        out: List[List[float]] = []
        prev_side = self.tok.padding_side
        self.tok.padding_side = "left"
        bs = batch_size
        i = 0
        try:
            while i < len(texts):
                if self.deadline is not None and time.monotonic() > self.deadline:
                    out.extend([[0.0] * n_opt for _ in range(len(texts) - i)])
                    break
                chunk = texts[i:i + bs]
                try:
                    enc = self.tok(chunk, return_tensors="pt", padding=True,
                                   add_special_tokens=False, truncation=True,
                                   max_length=6144).to(self.model.device)
                    fwd = {self._logits_kwarg: 1} if self._logits_kwarg else {}
                    with torch.no_grad():
                        logits = self.model(**enc, **fwd).logits[:, -1, :].float()
                    logprobs = torch.log_softmax(logits, dim=-1)
                    for b in range(len(chunk)):
                        row = [max((logprobs[b, t].item() for t in group),
                                   default=-1e9) if group else -1e9
                               for group in cand_token_ids]
                        out.append(row)
                    i += bs
                except torch.cuda.OutOfMemoryError:
                    torch.cuda.empty_cache()
                    if bs == 1:
                        out.append([0.0] * n_opt)
                        i += 1
                    else:
                        bs = max(1, bs // 2)
        finally:
            self.tok.padding_side = prev_side
        return out


class VLLMClient(LLMClient):
    """Optional vLLM backend for throughput experiments. Never required."""

    available = True

    def __init__(self, model_dir: str = DEFAULT_MODEL_DIR, max_model_len: int = 4096):
        from vllm import LLM

        self.llm = LLM(model=model_dir, dtype="half", max_model_len=max_model_len,
                       gpu_memory_utilization=0.90)

    def generate(self, prompts, max_new_tokens=MAX_NEW_TOKENS, system=None):
        from vllm import SamplingParams

        tok = self.llm.get_tokenizer()
        texts = []
        for p in prompts:
            messages = ([{"role": "system", "content": system}] if system else []) \
                + [{"role": "user", "content": p}]
            texts.append(tok.apply_chat_template(messages, add_generation_prompt=True,
                                                 tokenize=False))
        params = SamplingParams(temperature=0.0, max_tokens=max_new_tokens)
        outs = self.llm.generate(texts, params)
        return [_guard(o.outputs[0].text if o.outputs else "").strip() for o in outs]


def _guard(text: str, max_repeat: int = 4) -> str:
    """Loop-collapse guard: truncate at the point where a line repeats more
    than `max_repeat` times consecutively."""
    lines = text.splitlines()
    out, streak = [], 0
    for i, l in enumerate(lines):
        if i > 0 and l.strip() and l == lines[i - 1]:
            streak += 1
            if streak >= max_repeat:
                break
        else:
            streak = 0
        out.append(l)
    return "\n".join(out)


def load_client(model_id: Optional[str] = None) -> LLMClient:
    """Best available client for `model_id` (script.py's MODEL_ID):
    - a local path ("." in the submission, weights/base in dev) is loaded
      when its config.json exists;
    - a Hub name (contains "/" and is not a local dir) is passed straight to
      transformers — the Colab-testing path, mirroring the workshop notebook;
    - anything unloadable degrades to NullClient (symbolic-only pipeline)."""
    if model_id and "/" in model_id and not Path(model_id).exists():
        try:
            return HFTransformersClient(model_dir=model_id)
        except Exception:
            return NullClient()
    for d in ([model_id] if model_id else []) + [DEFAULT_MODEL_DIR, "weights/base"]:
        if d and Path(d, "config.json").exists():
            try:
                return HFTransformersClient(model_dir=d)
            except Exception:
                continue
    return NullClient()