"""Model loading and affect-conditioned generation for CODELOCK."""

from __future__ import annotations

from typing import List, Optional
from contextlib import nullcontext

import torch

DEFAULT_MODEL = "Qwen/Qwen3-8B"

from .persona import CODELOCK_SYSTEM


def load(
    model_id: str = DEFAULT_MODEL,
    four_bit: bool = False,
    device_map: str = "auto",
    revision=None,
):
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(
        model_id, revision=revision, trust_remote_code=False
    )
    kwargs = {
        "device_map": device_map,
        "torch_dtype": torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        "revision": revision,
        "trust_remote_code": False,
        "attn_implementation": "sdpa",
    }
    if four_bit:
        from transformers import BitsAndBytesConfig

        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
    model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
    model.eval()
    return model, tokenizer


def format_prompt(tokenizer, user, history=None, system=CODELOCK_SYSTEM):
    messages = (
        [{"role": "system", "content": system}]
        + list(history or [])
        + [{"role": "user", "content": user}]
    )
    kwargs = {"tokenize": False, "add_generation_prompt": True}
    try:
        prompt = tokenizer.apply_chat_template(
            messages, enable_thinking=False, **kwargs
        )
    except TypeError:
        prompt = tokenizer.apply_chat_template(messages, **kwargs)
    return prompt


@torch.no_grad()
def generate(
    model,
    tokenizer,
    user: str,
    history: Optional[List[dict]] = None,
    system: str = CODELOCK_SYSTEM,
    max_new_tokens: int = 200,
    temperature: float = 0.7,
    top_p: float = 0.8,
    top_k: int = 20,
    repetition_penalty: float = 1.15,
    seed: Optional[int] = None,
    do_sample: bool = True,
) -> str:
    prompt = format_prompt(tokenizer, user, history, system)
    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(
        model.get_input_embeddings().weight.device
    )
    owner = getattr(model, "_codelock_steerer", None)
    sampling = (
        {"temperature": temperature, "top_p": top_p, "top_k": top_k}
        if do_sample
        else {}
    )
    devices = sorted(
        {p.device.index for p in model.parameters() if p.device.type == "cuda"}
    )
    with torch.random.fork_rng(devices=devices), (
        owner.frozen() if owner is not None else nullcontext()
    ):
        if seed is not None:
            torch.manual_seed(seed)
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            **sampling,
            repetition_penalty=repetition_penalty,
            pad_token_id=(
                tokenizer.pad_token_id
                if tokenizer.pad_token_id is not None
                else tokenizer.eos_token_id
            ),
        )
    return tokenizer.decode(
        out[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True
    ).strip()