import math
from typing import List

import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer


def select_heads(head_list: List[str], topk: int) -> List[str]:
    if topk <= 0:
        return []
    return head_list[: min(topk, len(head_list))]


def build_head_map(head_names: List[str]) -> dict:
    head_map = {}
    for head_name in head_names:
        if not head_name.startswith("L"):
            raise ValueError(f"Invalid head name: {head_name}")
        try:
            layer_part, head_part = head_name.split("_", 1)
            layer_idx = int(layer_part[1:])
            head_idx = int(head_part[1:])
        except Exception as exc:
            raise ValueError(f"Invalid head name: {head_name}") from exc
        head_map.setdefault(layer_idx, []).append(head_idx)
    return head_map


def load_model(model_path: str):
    cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
        tok.pad_token_id = tok.eos_token_id
    tok.padding_side = "left"
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        config=cfg,
        device_map="auto",
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation="eager",
    )
    model.eval()
    return model, tok


def mask_attention(attn_head: torch.Tensor, data_positions: List[int]) -> torch.Tensor:
    if not data_positions:
        return attn_head
    masked = attn_head.clone()
    masked[:, data_positions] = 0.0
    return masked


def build_head_summaries(
    attentions,
    selected_heads: List[str],
    data_positions_batch: List[List[int]],
    n_layers: int,
    n_heads: int,
):
    summaries = []
    for batch_idx, data_positions in enumerate(data_positions_batch):
        head_info = {}
        for head_name in selected_heads:
            if not head_name.startswith("L"):
                raise ValueError(f"Invalid head name: {head_name}")
            try:
                layer_part, head_part = head_name.split("_", 1)
                layer_idx = int(layer_part[1:])
                head_idx = int(head_part[1:])
            except Exception as exc:
                raise ValueError(f"Invalid head name: {head_name}") from exc
            if layer_idx < 0 or layer_idx >= n_layers or head_idx < 0 or head_idx >= n_heads:
                raise ValueError(f"Head out of range for model: {head_name}")
            attn_head = attentions[layer_idx][batch_idx, head_idx]
            masked = mask_attention(attn_head, data_positions)
            pre_sum = attn_head[:, data_positions].sum().float().item() if data_positions else 0.0
            post_sum = masked[:, data_positions].sum().float().item() if data_positions else 0.0
            head_info[head_name] = {
                "pre_data_attention_sum": pre_sum,
                "post_data_attention_sum": post_sum,
            }
        summaries.append(head_info)
    return summaries


def total_attn_to_data_batch(attn, data_positions_by_sample: List[List[int]]) -> List[float]:
    totals = []
    for b, data_positions in enumerate(data_positions_by_sample):
        if not data_positions:
            totals.append(0.0)
            continue
        total = 0.0
        for layer_attn in attn:
            total += layer_attn[b, :, -1, data_positions].sum().float().item()
        totals.append(total)
    return totals


def debug_print_attention_totals(
    data_indices,
    unmasked_attn,
    masked_attn,
    data_positions_batch: List[List[int]],
    debug: bool = False,
):
    if not debug:
        return
    unmasked_totals = total_attn_to_data_batch(unmasked_attn, data_positions_batch)
    masked_totals = total_attn_to_data_batch(masked_attn, data_positions_batch)
    for idx, (u, m) in enumerate(zip(unmasked_totals, masked_totals)):
        sample_id = data_indices[idx] if idx < len(data_indices) else idx
        print(f"[DEBUG] idx={sample_id} unmasked_attn_sum={u:.6f}")
        print(f"[DEBUG] idx={sample_id} masked_attn_sum={m:.6f}")


def debug_print_head_summaries(
    data_indices,
    head_summaries: List[dict],
    debug: bool = False,
):
    if not debug:
        return
    for idx, head_info in enumerate(head_summaries):
        sample_id = data_indices[idx] if idx < len(data_indices) else idx
        print(f"[DEBUG] idx={sample_id} head_summaries={len(head_info)}")


def _make_head_mask_hook(
    layer_idx: int,
    head_indices: List[int],
    data_positions_by_sample: List[List[int]],
    num_heads: int,
    debug: bool = False,
):
    head_indices = list(sorted(set(head_indices)))
    seen = {"printed": False}

    def _hook(_module, args, kwargs):
        if not data_positions_by_sample:
            return None

        attention_mask = None
        if kwargs is not None:
            attention_mask = kwargs.get("attention_mask", None)
        if attention_mask is None and len(args) >= 2:
            attention_mask = args[1]

        if attention_mask is None:
            return None

        bsz, mask_heads, q_len, k_len = attention_mask.shape

        if mask_heads == 1 and num_heads > 1:
            attention_mask = attention_mask.expand(bsz, num_heads, q_len, k_len).clone()

        if attention_mask.dtype == torch.bool:
            val_to_fill = True
        else:
            val_to_fill = torch.finfo(attention_mask.dtype).min

        for b in range(bsz):
            masked_positions = [p for p in data_positions_by_sample[b] if p < k_len]
            if not masked_positions:
                continue
            masked_pos_tensor = torch.tensor(masked_positions, device=attention_mask.device)
            if debug and layer_idx == 0 and not seen["printed"]:
                target_head = head_indices[0] if head_indices else 0
                sample_pos = masked_positions[:3]
                if sample_pos:
                    sample_vals = attention_mask[b, target_head, -1, sample_pos].detach().cpu().tolist()
                    print(f"[DEBUG] L{layer_idx}: head={target_head} sample_before={sample_vals}")

            for h in head_indices:
                attention_mask[b, h].index_fill_(1, masked_pos_tensor, val_to_fill)

            if debug and layer_idx == 0 and not seen["printed"]:
                target_head = head_indices[0] if head_indices else 0
                sample_pos = masked_positions[:3]
                if sample_pos:
                    sample_vals = attention_mask[b, target_head, -1, sample_pos].detach().cpu().tolist()
                    print(f"[DEBUG] L{layer_idx}: head={target_head} sample_after={sample_vals}")
                seen["printed"] = True

        if kwargs is not None and "attention_mask" in kwargs:
            kwargs["attention_mask"] = attention_mask
            return args, kwargs

        new_args = list(args)
        if len(new_args) >= 2:
            new_args[1] = attention_mask
            return tuple(new_args), kwargs

        return None

    return _hook


def _install_mask_hooks(
    model,
    head_map: dict,
    data_positions_by_sample: List[List[int]],
    num_heads: int,
    debug: bool = False,
):
    handles = []
    layers = getattr(getattr(model, "model", None), "layers", None)
    if layers is None:
        raise ValueError("Unsupported model layout: missing model.model.layers")
    for layer_idx, head_indices in head_map.items():
        if layer_idx < 0 or layer_idx >= len(layers):
            raise ValueError(f"Layer index out of range: L{layer_idx}")
        attn = getattr(layers[layer_idx], "self_attn", None)
        if attn is None:
            raise ValueError(f"Layer L{layer_idx} missing self_attn")
        hook = _make_head_mask_hook(layer_idx, head_indices, data_positions_by_sample, num_heads, debug=debug)
        handles.append(attn.register_forward_pre_hook(hook, with_kwargs=True))
    return handles


class MaskedCausalLM:
    def __init__(self, model, head_map: dict, num_heads: int, debug: bool = False):
        self._model = model
        self._head_map = head_map
        self._num_heads = num_heads
        self._debug = debug

    def __getattr__(self, name):
        return getattr(self._model, name)

    def _validate_data_positions(self, data_positions_batch, batch_size):
        if data_positions_batch is None:
            raise ValueError("data_positions_batch is required for masked inference.")
        if batch_size is not None and len(data_positions_batch) != batch_size:
            raise ValueError(
                f"data_positions_batch size {len(data_positions_batch)} does not match batch size {batch_size}."
            )

    def _with_masking(self, data_positions_batch, fn):
        if not self._head_map:
            return fn()
        handles = _install_mask_hooks(
            self._model,
            self._head_map,
            data_positions_batch,
            self._num_heads,
            debug=self._debug,
        )
        try:
            return fn()
        finally:
            for handle in handles:
                handle.remove()

    def __call__(self, *args, **kwargs):
        data_positions_batch = kwargs.pop("data_positions_batch", None)
        batch_size = None
        input_ids = kwargs.get("input_ids", None)
        if input_ids is None and args:
            input_ids = args[0]
        if input_ids is not None and hasattr(input_ids, "shape"):
            batch_size = input_ids.shape[0]
        self._validate_data_positions(data_positions_batch, batch_size)
        return self._with_masking(data_positions_batch, lambda: self._model(*args, **kwargs))

    def generate(self, *args, **kwargs):
        data_positions_batch = kwargs.pop("data_positions_batch", None)
        batch_size = None
        input_ids = kwargs.get("input_ids", None)
        if input_ids is None and args:
            input_ids = args[0]
        if input_ids is not None and hasattr(input_ids, "shape"):
            batch_size = input_ids.shape[0]
        self._validate_data_positions(data_positions_batch, batch_size)
        return self._with_masking(data_positions_batch, lambda: self._model.generate(*args, **kwargs))


def build_masked_model(model, head_list: List[str], topk: str, debug: bool = False):
    if not isinstance(head_list, list) or not head_list:
        raise ValueError("head_list must be a non-empty list like ['L1H5', 'L15H23'].")
    if len(head_list) <= 0:
        topk_count = 0
    elif topk is None:
        topk_count = len(head_list)
    else:
        topk_str = str(topk).strip().lower()
        if topk_str.endswith("p"):
            pct = float(topk_str[:-1])
            if pct <= 0:
                topk_count = 0
            else:
                topk_count = max(1, int(math.ceil(len(head_list) * pct / 100.0)))
        else:
            topk_count = max(0, int(topk_str))
    selected_heads = select_heads(head_list, topk_count)
    num_heads = getattr(model.config, "num_attention_heads", None)
    if num_heads is None:
        raise ValueError("Model config missing num_attention_heads.")
    head_map = build_head_map(selected_heads)
    masked_model = MaskedCausalLM(model, head_map, num_heads, debug=debug)
    return masked_model, selected_heads

def _generate_batch(model, tok, input_ids_batch, attention_mask_batch, max_new_tokens, data_positions_batch=None):
    if not input_ids_batch:
        return []
    input_ids_tensor = torch.tensor(input_ids_batch, dtype=torch.long, device=model.device)
    attention_mask_tensor = torch.tensor(attention_mask_batch, dtype=torch.long, device=model.device)
    if data_positions_batch is None or (type(model) != MaskedCausalLM):
        out = model.generate(
            input_ids=input_ids_tensor,
            attention_mask=attention_mask_tensor,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            eos_token_id=tok.eos_token_id,
            pad_token_id=tok.pad_token_id,
        )
    else:
        out = model.generate(
            input_ids=input_ids_tensor,
            attention_mask=attention_mask_tensor,
            data_positions_batch=data_positions_batch,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            eos_token_id=tok.eos_token_id,
            pad_token_id=tok.pad_token_id,
        )
    prompt_len = len(input_ids_batch[0])
    outputs = []
    for row in out:
        gen_ids = row.tolist()
        outputs.append(tok.decode(gen_ids[prompt_len:], skip_special_tokens=True))
    return outputs