"""Transformers inference with autoregressive decoding and batched field decisions."""

import copy
import json
import os
import threading
import time
from functools import wraps
from typing import Any, Dict, Generator

import torch
from torch.nn.attention.flex_attention import create_block_mask
from transformers import AutoModelForCausalLM, AutoTokenizer

from core.prompt_builder import build_naive_json_prompt
from core.schema import StructuredSchema

MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
_model = None
_tokenizer = None
_gpu_lock = threading.Lock()
_load_lock = threading.Lock()


def gpu_locked(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        with _gpu_lock, torch.inference_mode():
            return fn(*args, **kwargs)
    return wrapper


def gpu_locked_gen(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        # Inference mode is thread-local. Starlette can resume a sync iterator on
        # another worker, so enter it for each next(), never across a yield.
        iterator = fn(*args, **kwargs)
        _gpu_lock.acquire()
        try:
            while True:
                with torch.inference_mode():
                    try:
                        event = next(iterator)
                    except StopIteration:
                        return
                yield event
        finally:
            iterator.close()
            _gpu_lock.release()
    return wrapper


def get_engine():
    with _load_lock, torch.inference_mode():
        return _load_engine()


def _load_engine():
    global _model, _tokenizer
    if _model is None:
        device = torch.device(os.environ.get("DEVICE", "cuda" if torch.cuda.is_available() else "cpu"))
        dtype = torch.float32
        if device.type == "cuda":
            with torch.cuda.device(device):
                dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        print(f"Loading {MODEL_ID} on {device} ({dtype})...")
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=dtype).to(device).eval()
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token = tokenizer.eos_token
        warmup = tokenizer("Warmup", return_tensors="pt").to(device)
        model(**warmup, use_cache=True, logits_to_keep=1)
        _synchronize(model)
        _model, _tokenizer = model, tokenizer
        print("Engine ready.")
    return _model, _tokenizer


def _synchronize(model):
    if model.device.type == "cuda":
        torch.cuda.synchronize(model.device)


def _naive_result(text, schema, token_count, passes, elapsed_ms):
    parsed = None
    error = None
    try:
        parsed = json.loads(text.strip())
    except ValueError as exc:
        error = str(exc)
    valid = error is None
    missing, invalid = [], []
    if isinstance(parsed, dict):
        for name, field in schema.fields.items():
            if name not in parsed:
                missing.append(name)
            elif (field.field_type == "boolean" and not isinstance(parsed[name], bool)) or (
                field.field_type != "boolean" and parsed[name] not in field.choices
            ):
                invalid.append(f"{name}={parsed[name]}")
    return {
        "mode": "naive_autoregressive", "elapsed_ms": round(elapsed_ms, 2),
        "total_tokens": token_count,
        "tokens_per_second": round(token_count * 1000 / max(elapsed_ms, 1e-6), 1),
        "sequential_forward_passes": passes, "is_valid_json": valid,
        "schema_match": isinstance(parsed, dict) and set(parsed) == set(schema.fields) and not invalid,
        "raw_text": text, "parsed_json": parsed, "parse_error": error,
        "missing_keys": missing, "invalid_enums": invalid,
        "has_calibrated_probabilities": False,
    }


@gpu_locked_gen
def stream_naive_generation(context: str, schema: StructuredSchema, max_tokens: int = 700,
                            temperature: float = 0.2) -> Generator[Dict[str, Any], None, None]:
    if max_tokens < 1 or temperature < 0:
        raise ValueError("max_tokens must be positive and temperature nonnegative")
    model, tokenizer = get_engine()
    inputs = tokenizer(build_naive_json_prompt(context, schema), add_special_tokens=False,
                       return_tensors="pt").to(model.device)
    t0 = time.perf_counter()
    cache = None
    generated = []
    current_text = "{\n  "
    eos = model.generation_config.eos_token_id
    stop_tokens = set(eos if isinstance(eos, list) else [eos]) | {tokenizer.eos_token_id}
    passes = 0
    emitted = False
    for _ in range(max_tokens):
        output = model(**inputs, past_key_values=cache, use_cache=True, logits_to_keep=1)
        cache = output.past_key_values
        passes += 1
        logits = output.logits[0, -1].float()
        token = int(torch.argmax(logits) if temperature == 0 else
                    torch.multinomial(torch.softmax(logits / temperature, dim=-1), 1)[0])
        if token in stop_tokens:
            break
        generated.append(token)
        text = "{\n  " + tokenizer.decode(generated, skip_special_tokens=True,
                                           clean_up_tokenization_spaces=False)
        # Wait for complete UTF-8 characters before emitting a text delta.
        if not text.endswith("\ufffd"):
            delta = text[len(current_text):] if emitted else text
            emitted = True
            current_text = text
            yield {"type": "token", "token": delta, "accumulated": text,
                   "token_count": len(generated),
                   "elapsed_ms": round((time.perf_counter() - t0) * 1000, 1)}
        try:
            if isinstance(json.loads(text), dict):
                break
        except ValueError:
            pass
        inputs = {"input_ids": torch.tensor([[token]], device=model.device),
                  "attention_mask": torch.cat([inputs["attention_mask"],
                      torch.ones((1, 1), dtype=torch.long, device=model.device)], dim=1)}
    _synchronize(model)
    yield {"type": "done", "result": _naive_result(
        current_text, schema, len(generated), passes, (time.perf_counter() - t0) * 1000)}


def run_naive_generation(context: str, schema: StructuredSchema, max_tokens: int = 700,
                         temperature: float = 0.2) -> Dict[str, Any]:
    for event in stream_naive_generation(context, schema, max_tokens, temperature):
        if event["type"] == "done":
            return event["result"]


def _candidate_scores(model, cache, first_logits, candidates, pad_id):
    """Score complete allowed continuations, including their closing delimiter."""
    first_log_probs = first_logits.float().log_softmax(-1)
    scores = first_log_probs[[tokens[0] for tokens in candidates]].clone()
    passes = 0
    # Bound memory when a field has hundreds of multi-token choices.
    for start in range(0, len(candidates), 16):
        sequences = candidates[start:start + 16]
        width = max(map(len, sequences)) - 1
        if width == 0:
            continue
        batch = len(sequences)
        ids = torch.full((batch, width), pad_id, dtype=torch.long, device=model.device)
        mask = torch.zeros_like(ids)
        for row, tokens in enumerate(sequences):
            ids[row, :len(tokens) - 1] = torch.tensor(tokens[:-1], device=model.device)
            mask[row, :len(tokens) - 1] = 1
        branch = copy.deepcopy(cache)
        branch.batch_repeat_interleave(batch)
        attention = torch.cat([torch.ones((batch, cache.get_seq_length()),
                              dtype=torch.long, device=model.device), mask], dim=1)
        output = model(input_ids=ids, attention_mask=attention,
                       past_key_values=branch, use_cache=True)
        passes += 1
        log_probs = output.logits.float().log_softmax(-1)
        for row, tokens in enumerate(sequences):
            if len(tokens) > 1:
                scores[start + row] += log_probs[row, torch.arange(len(tokens) - 1,
                    device=model.device), tokens[1:]].sum()
        del output, branch, log_probs
    return scores, passes


@gpu_locked
def run_parallel_generation(context: str, schema: StructuredSchema,
                            temperature: float = 1.0) -> Dict[str, Any]:
    if temperature < 0:
        raise ValueError("temperature must be nonnegative")
    model, tokenizer = get_engine()
    t0 = time.perf_counter()
    meta = schema.compile_parallel_metadata(tokenizer)
    prompt = (f"<|im_start|>system\nClassify JSON attributes:\n{schema.to_parallel_schema_str()}"
              f"<|im_end|>\n<|im_start|>user\n{context}<|im_end|>\n"
              f"<|im_start|>assistant\n{{\n")
    inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
    t_pre = time.perf_counter()
    output = model(**inputs, use_cache=True, logits_to_keep=1)
    cache = output.past_key_values
    prefix_length = cache.get_seq_length()
    _synchronize(model)
    prefill_ms = (time.perf_counter() - t_pre) * 1000
    del output
    cache.batch_repeat_interleave(len(schema))
    suffixes = meta["suffixes_batch"].to(model.device)
    suffix_mask = meta["suffix_attention_mask"].to(model.device)
    attention = torch.cat([inputs["attention_mask"].repeat(len(schema), 1), suffix_mask], dim=1)
    t_suffix = time.perf_counter()
    output = model(input_ids=suffixes, attention_mask=attention, past_key_values=cache, use_cache=True)
    _synchronize(model)
    suffix_ms = (time.perf_counter() - t_suffix) * 1000
    parsed, telemetry = {}, {}
    passes = 2  # Prefix prefill plus batched field suffixes.
    for i, (name, field) in enumerate(meta["field_items"]):
        length = meta["suffix_lengths"][i]
        logits = output.logits[i, length - 1]
        candidates = meta["candidate_sequences"][i]
        if meta["has_collisions"][i]:
            branch = copy.deepcopy(cache)
            branch.batch_select_indices(torch.tensor([i], device=model.device))
            # Remove right padding before evaluating candidate continuations.
            padding = branch.get_seq_length() - prefix_length - length
            if padding:
                branch.crop(-padding)
            scores, extra_passes = _candidate_scores(model, branch, logits, candidates, tokenizer.pad_token_id)
            passes += extra_passes
            del branch
        else:
            scores = logits[[tokens[0] for tokens in candidates]].float()
        probabilities = torch.softmax(scores / max(temperature, 1e-4), dim=-1).tolist()
        winner = max(range(len(probabilities)), key=probabilities.__getitem__)
        choice = field.choices[winner]
        value = choice == "true" if field.field_type == "boolean" else choice
        confidence = round(probabilities[winner], 4)
        parsed[name] = {"value": value, "prob": confidence}
        top = sorted([{"choice": c, "probability": round(p, 4)}
                      for c, p in zip(field.choices, probabilities)],
                     key=lambda item: item["probability"], reverse=True)
        telemetry[name] = {"value": value, "type": field.field_type, "confidence": confidence,
                           "cardinality": field.cardinality, "top_choices": top[:5]}
    _synchronize(model)
    return {"mode": "parallel_constrained_calibrated", "elapsed_ms": round((time.perf_counter() - t0) * 1000, 2),
            "prefill_ms": round(prefill_ms, 2), "suffix_eval_ms": round(suffix_ms, 2),
            "total_tokens_generated": 0, "sequential_forward_passes": passes,
            "is_valid_json": True, "schema_match": True, "parsed_json": parsed,
            "field_telemetry": telemetry, "has_calibrated_probabilities": True, "num_fields": len(schema)}


@torch.compile(fullgraph=True)
def _tree_attention_mask(starts, ends, prefix_length):
    """Compile the mask predicate into block metadata without a dense token mask."""
    size = starts.shape[0]
    limits = torch.tensor([size - 1, prefix_length, prefix_length + size], device=starts.device)

    def visible(batch, head, query, key):
        # FlexAttention can evaluate padded block indices. Clamp all lookups;
        # the validity checks below keep those positions masked out.
        q = torch.minimum(query.clamp_min(0), limits[0])
        k = torch.minimum((key - limits[1]).clamp_min(0), limits[0])
        ancestor = (starts[k] <= starts[q]) & (starts[q] < ends[k])
        return ((query <= limits[0]) & (key < limits[2])
                & ((key < limits[1]) | ancestor))

    return create_block_mask(visible, B=None, H=None, Q_LEN=size,
                             KV_LEN=prefix_length + size, device=starts.device)


@gpu_locked
def run_parallel_generation_tree(context: str, schema: StructuredSchema,
                                 temperature: float = 1.0) -> Dict[str, Any]:
    """Score all complete candidates in one tree pass over a single prefix cache.

    Uses a FlexAttention BlockMask with DFS intervals for ancestor visibility.
    Scores are mean token log-probabilities over the full candidate, including
    shared value prefixes and closing delimiters. The model's original
    attention backend is restored afterward.
    """
    if temperature < 0:
        raise ValueError("temperature must be nonnegative")
    model, tokenizer = get_engine()
    if "sliding_attention" in getattr(model.config, "layer_types", []):
        raise ValueError("Tree generation requires full attention, not a sliding-window cache")
    t0 = time.perf_counter()
    meta = schema.compile_tree_metadata(tokenizer)
    prompt = (f"<|im_start|>system\nClassify JSON attributes:\n{schema.to_parallel_schema_str()}"
              f"<|im_end|>\n<|im_start|>user\n{context}<|im_end|>\n"
              f"<|im_start|>assistant\n{{\n")
    inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
    t_pre = time.perf_counter()
    output = model(**inputs, use_cache=True, logits_to_keep=1)
    cache = output.past_key_values
    prefix_length = cache.get_seq_length()
    _synchronize(model)
    prefill_ms = (time.perf_counter() - t_pre) * 1000
    del output

    t_suffix = time.perf_counter()
    ids = meta["input_ids"].to(model.device)
    positions = meta["depths"].to(model.device) + prefix_length
    size = ids.shape[1]
    attention = _tree_attention_mask(meta["dfs_starts"].to(model.device),
                                     meta["dfs_ends"].to(model.device), prefix_length)
    original_backend = model.config._attn_implementation
    try:
        model.set_attn_implementation("flex_attention")
        # The short-query GQA decode kernel can pack more queries than fit in
        # a BlockMask tile. Use the general FlexAttention kernel for the tree.
        output = model(input_ids=ids, attention_mask=attention, position_ids=positions,
                       past_key_values=cache, use_cache=True,
                       logits_to_keep=meta["logit_indices"].to(model.device),
                       kernel_options={"FORCE_USE_FLEX_ATTENTION": True})
    finally:
        model.set_attn_implementation(original_backend)
    _synchronize(model)
    suffix_ms = (time.perf_counter() - t_suffix) * 1000
    # Normalize only the parent logits needed to score candidate edges, without
    # allocating another full [nodes, vocabulary] log-probability tensor.
    logits = output.logits[0].float()
    rows = meta["score_rows"].to(model.device)
    tokens = meta["score_tokens"].to(model.device)
    edge_scores = logits[rows, tokens] - logits.logsumexp(-1)[rows]
    scores = logits.new_zeros(sum(meta["candidate_counts"]))
    scores.scatter_add_(0, meta["score_candidates"].to(model.device), edge_scores)
    scores /= meta["candidate_lengths"].to(model.device)
    del output, cache, logits, attention

    parsed, telemetry = {}, {}
    for (name, field), field_scores in zip(meta["field_items"], scores.split(meta["candidate_counts"])):
        probabilities = torch.softmax(field_scores / max(temperature, 1e-4), dim=-1).tolist()
        winner = max(range(len(probabilities)), key=probabilities.__getitem__)
        choice = field.choices[winner]
        value = choice == "true" if field.field_type == "boolean" else choice
        confidence = round(probabilities[winner], 4)
        parsed[name] = {"value": value, "prob": confidence}
        top = sorted([{"choice": c, "probability": round(p, 4)}
                      for c, p in zip(field.choices, probabilities)],
                     key=lambda item: item["probability"], reverse=True)
        telemetry[name] = {"value": value, "type": field.field_type, "confidence": confidence,
                           "cardinality": field.cardinality, "top_choices": top[:5]}
    _synchronize(model)
    return {"mode": "parallel_constrained_tree", "elapsed_ms": round((time.perf_counter() - t0) * 1000, 2),
            "prefill_ms": round(prefill_ms, 2), "suffix_eval_ms": round(suffix_ms, 2),
            "total_tokens_generated": 0, "sequential_forward_passes": 2,
            "is_valid_json": True, "schema_match": True, "parsed_json": parsed,
            "field_telemetry": telemetry, "has_calibrated_probabilities": True, "num_fields": len(schema),
            "tree_nodes": size, "prefix_tokens": prefix_length,
            "scoring": "mean_token_log_probability"}


run_rlcd_generation = run_parallel_generation