from __future__ import annotations

import math
import os
import importlib.util
from dataclasses import dataclass

import torch
import torch.nn.functional as F
from huggingface_hub import login
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

from .checkpoints import load_token_checkpoint
from .config import PipelineConfig


def _load_causal_lm(model_name: str, **kwargs):
    """Load a causal LM, falling back to architecture-specific classes when
    AutoModelForCausalLM doesn't recognize the model type."""
    try:
        return AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
    except (ValueError, ModuleNotFoundError):
        config = AutoConfig.from_pretrained(model_name)
        model_type = getattr(config, "model_type", "")
        if model_type == "gemma3":
            from transformers import Gemma3ForConditionalGeneration
            return Gemma3ForConditionalGeneration.from_pretrained(model_name, **kwargs)
        if model_type == "gemma4":
            from transformers import Gemma4ForConditionalGeneration
            return Gemma4ForConditionalGeneration.from_pretrained(model_name, **kwargs)
        raise


def resolve_hf_token() -> str | None:
    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")


def maybe_login_hf() -> None:
    token = resolve_hf_token()
    if token:
        login(token=token, add_to_git_credential=False)


def resolve_dtype() -> torch.dtype:
    return torch.bfloat16 if torch.cuda.is_available() else torch.float32


def has_accelerate() -> bool:
    return importlib.util.find_spec("accelerate") is not None


@dataclass(slots=True)
class ModelBundle:
    config: PipelineConfig
    tokenizer: AutoTokenizer
    model: AutoModelForCausalLM
    initial_tokenizer_len: int


def initialize_model_bundle(config: PipelineConfig) -> ModelBundle:
    maybe_login_hf()
    hf_token = resolve_hf_token()
    dtype = resolve_dtype()

    tokenizer = AutoTokenizer.from_pretrained(config.model.model_name, use_fast=True, token=hf_token)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    initial_len = len(tokenizer)
    tokenizer.add_special_tokens(
        {"additional_special_tokens": [config.model.ai_token, config.model.human_token]}
    )

    load_kwargs = {
        "token": hf_token,
        "dtype": dtype,
    }
    if torch.cuda.is_available() and has_accelerate():
        load_kwargs["device_map"] = "auto"
    model = _load_causal_lm(config.model.model_name, **load_kwargs)
    if torch.cuda.is_available() and not has_accelerate():
        model = model.to("cuda")
    elif not torch.cuda.is_available():
        model = model.to("cpu")

    try:
        model.resize_token_embeddings(len(tokenizer), mean_resizing=False)
    except TypeError:
        model.resize_token_embeddings(len(tokenizer))

    # Mean-fill Gemma 4's auxiliary per-layer token table. Some Transformers
    # versions resize it internally, while others leave it at the old size.
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Embedding) and module is not model.get_input_embeddings():
            if module.weight.shape[0] == initial_len:
                mean_row = module.weight.data[:initial_len].mean(
                    dim=0, dtype=torch.float32
                ).to(dtype=module.weight.dtype)
                new_emb = torch.nn.Embedding(
                    len(tokenizer), module.weight.shape[1],
                    device=module.weight.device, dtype=module.weight.dtype,
                )
                new_emb.weight.data[:initial_len] = module.weight.data
                new_emb.weight.data[initial_len:] = mean_row
                if not torch.equal(
                    new_emb.weight.data[initial_len], mean_row
                ):
                    raise RuntimeError(
                        f"Failed to mean-fill resized embedding {name}"
                    )
                parent_name, attr_name = name.rsplit(".", 1)
                parent = dict(model.named_modules())[parent_name]
                setattr(parent, attr_name, new_emb)
                print(
                    f"Mean-filled secondary embedding {name}: "
                    f"{initial_len} -> {len(tokenizer)}"
                )
            elif module.weight.shape[0] == len(tokenizer):
                mean_row = module.weight.data[:initial_len].mean(
                    dim=0, dtype=torch.float32
                ).to(dtype=module.weight.dtype)
                module.weight.data[initial_len:] = mean_row
                if not torch.equal(module.weight.data[initial_len], mean_row):
                    raise RuntimeError(
                        f"Failed to mean-fill expanded embedding {name}"
                    )
                print(
                    f"Mean-filled expanded secondary embedding {name}: "
                    f"rows {initial_len}:{len(tokenizer)}"
                )

    model.config.use_cache = False
    if hasattr(model, "gradient_checkpointing_enable"):
        model.gradient_checkpointing_enable()
    model.eval()

    input_emb = model.get_input_embeddings()
    with torch.no_grad():
        mean_in = input_emb.weight[:initial_len].mean(dim=0)
        for token in (config.model.ai_token, config.model.human_token):
            token_id = tokenizer.convert_tokens_to_ids(token)
            input_emb.weight[token_id].copy_(mean_in + torch.randn_like(mean_in) * 1e-5)

    bundle = ModelBundle(
        config=config,
        tokenizer=tokenizer,
        model=model,
        initial_tokenizer_len=initial_len,
    )
    apply_initial_checkpoints(bundle)
    return bundle


def apply_initial_checkpoints(bundle: ModelBundle) -> None:
    input_emb = bundle.model.get_input_embeddings()
    token_dir = bundle.config.output.model_tokens_dir(bundle.config.model.model_name)
    ai_path = bundle.config.init_checkpoints.ai_token_path or token_dir / "ai_token.pt"
    human_path = bundle.config.init_checkpoints.human_token_path or token_dir / "human_token.pt"

    secondary_embs = [
        module
        for module in bundle.model.modules()
        if (
            isinstance(module, torch.nn.Embedding)
            and module is not input_emb
            and module.weight.shape[0] == len(bundle.tokenizer)
        )
    ]

    def install(path, token):
        if not path.exists():
            return
        checkpoint = load_token_checkpoint(path)
        token_id = bundle.tokenizer.convert_tokens_to_ids(token)
        saved_secondary = checkpoint.secondary_embeddings or []
        if len(saved_secondary) != len(secondary_embs):
            raise ValueError(
                f"{path} has {len(saved_secondary)} secondary rows, but "
                f"{bundle.config.model.model_name} exposes "
                f"{len(secondary_embs)} secondary token embeddings."
            )
        input_emb.weight[token_id].copy_(
            checkpoint.embedding.to(
                input_emb.weight.device, dtype=input_emb.weight.dtype
            )
        )
        for embedding, row in zip(secondary_embs, saved_secondary):
            embedding.weight[token_id].copy_(
                row.to(embedding.weight.device, dtype=embedding.weight.dtype)
            )

    with torch.no_grad():
        install(ai_path, bundle.config.model.ai_token)
        install(human_path, bundle.config.model.human_token)


def build_prompt(bundle: ModelBundle, token: str) -> str:
    content = bundle.config.model.prompt_template.format(token=token)
    messages = [{"role": "user", "content": content}]
    return bundle.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)


def encode_response(bundle: ModelBundle, prompt_text: str, response_text: str) -> tuple[torch.Tensor, int]:
    full_text = prompt_text + response_text
    prompt_ids = bundle.tokenizer(prompt_text, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
    full_ids = bundle.tokenizer(
        full_text,
        return_tensors="pt",
        add_special_tokens=False,
        truncation=True,
        max_length=bundle.config.model.max_length,
    )["input_ids"][0]
    return full_ids, int(len(prompt_ids))


def compute_average_logprob(bundle: ModelBundle, input_ids: torch.Tensor, prompt_len: int) -> torch.Tensor:
    token_logps = compute_token_logprobs(bundle, input_ids, prompt_len)
    return token_logps.mean()


def _model_forward(bundle: ModelBundle, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None):
    """Run a forward pass, injecting token_type_ids for models that require it."""
    fwd_kwargs: dict = {"input_ids": input_ids}
    if attention_mask is not None:
        fwd_kwargs["attention_mask"] = attention_mask
    model_type = getattr(bundle.model.config, "model_type", "")
    if model_type in ("gemma3", "gemma4"):
        fwd_kwargs["token_type_ids"] = torch.zeros_like(input_ids)
    return bundle.model(**fwd_kwargs)


def compute_token_logprobs(bundle: ModelBundle, input_ids: torch.Tensor, prompt_len: int) -> torch.Tensor:
    input_ids = input_ids.unsqueeze(0).to(bundle.model.device)
    logits = _model_forward(bundle, input_ids).logits[0]
    shift_logits = logits[prompt_len - 1 : -1]
    shift_labels = input_ids[0, prompt_len:]
    log_probs = F.log_softmax(shift_logits, dim=-1)
    return log_probs[torch.arange(len(shift_labels), device=bundle.model.device), shift_labels]


def compute_sequence_logprob(bundle: ModelBundle, input_ids: torch.Tensor, prompt_len: int) -> torch.Tensor:
    input_ids = input_ids.unsqueeze(0).to(bundle.model.device)
    logits = _model_forward(bundle, input_ids).logits[0]
    shift_logits = logits[prompt_len - 1 : -1]
    shift_labels = input_ids[0, prompt_len:]
    log_probs = F.log_softmax(shift_logits, dim=-1)
    token_logps = log_probs[torch.arange(len(shift_labels), device=bundle.model.device), shift_labels]
    return token_logps.sum()


def cosine_with_floor(
    step: int,
    total_steps: int,
    base_lr: float,
    *,
    min_lr: float,
    warmup_steps: int,
) -> float:
    if step < warmup_steps:
        return base_lr * float(step + 1) / float(max(1, warmup_steps))
    progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr + (base_lr - min_lr) * cosine