import argparse
import ast
import json
import os
from typing import List

import matplotlib.pyplot as plt
import numpy as np
import torch
from matplotlib.colors import LogNorm
from peft import PeftModel
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

from lib.tokenize_data_mask import apply_chat_with_tokenize_with_mark


CUSTOM_MASK_IDENTIFIER = {
    "data": ["<data>", "</data>"],
    "inst": ["<inst>", "</inst>"],
}


def load_message_list(path: str) -> List[dict]:
    with open(path, "r", encoding="utf-8") as f:
        raw_text = f.read().strip()

    if not raw_text:
        raise ValueError(f"Input file is empty: {path}")

    candidates = []

    try:
        parsed = json.loads(raw_text)
        candidates.append(parsed)
    except json.JSONDecodeError:
        pass

    try:
        parsed = ast.literal_eval(raw_text)
        candidates.append(parsed)
    except Exception:
        pass

    jsonl_items = []
    jsonl_ok = True
    for line in raw_text.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            jsonl_items.append(json.loads(line))
        except json.JSONDecodeError:
            jsonl_ok = False
            break
    if jsonl_ok and jsonl_items:
        candidates.append(jsonl_items)

    for candidate in candidates:
        if isinstance(candidate, dict) and isinstance(candidate.get("messages"), list):
            candidate = candidate["messages"]
        if isinstance(candidate, list) and all(isinstance(item, dict) for item in candidate):
            return candidate

    raise ValueError(
        "Input must be a message list: JSON array of {role, content}, JSON object with `messages`, "
        "or JSONL with one message object per line."
    )


def load_model_and_tokenizer(model_path: str, lora_path: str):
    if not os.path.isdir(model_path):
        raise FileNotFoundError(f"Model path not found or not a directory: {model_path}")
    if lora_path and not os.path.isdir(lora_path):
        raise FileNotFoundError(f"LoRA path not found or not a directory: {lora_path}")

    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 = "right"

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        config=cfg,
        device_map="auto",
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation="eager",
    )
    if lora_path:
        model = PeftModel.from_pretrained(model, lora_path, device_map="auto")
        model = model.merge_and_unload()

    model.eval()
    return model, tok


def load_head_order(path: str) -> List[str]:
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)

    if not isinstance(data, list):
        raise ValueError(f"Head-order JSON must contain a list: {path}")

    ordered_heads = []
    for item in data:
        if isinstance(item, list) and item:
            ordered_heads.append(str(item[0]))
        else:
            ordered_heads.append(str(item))
    return ordered_heads


@torch.inference_mode()
def collect_prompt_and_attention(model, tok, messages: List[dict]):
    (
        input_ids,
        instruction_mask,
        _data_mask,
        segment_type,
        _is_normal_token,
        custom_mask,
        _rendered_prompt,
    ) = apply_chat_with_tokenize_with_mark(
        messages,
        tok,
        custom_mask_identifier=CUSTOM_MASK_IDENTIFIER,
        add_generation_prompt=True,
    )

    input_ids_tensor = torch.tensor([input_ids], dtype=torch.long, device=model.device)
    attention_mask = torch.ones_like(input_ids_tensor)
    out = model(
        input_ids=input_ids_tensor,
        attention_mask=attention_mask,
        output_attentions=True,
    )

    n_layers = len(out.attentions)
    n_heads = out.attentions[0].shape[1]
    token_len = len(input_ids)

    attn_rows = []
    head_names = []
    for layer_idx in range(n_layers):
        layer_attn = out.attentions[layer_idx][0]
        for head_idx in range(n_heads):
            head_names.append(f"L{layer_idx}_H{head_idx}")
            attn_rows.append(layer_attn[head_idx].to(torch.float32).cpu().numpy()[-1, :])

    attn_matrix = np.array(attn_rows, dtype=np.float32).T
    if attn_matrix.shape[0] != token_len:
        raise ValueError(
            f"Attention/token mismatch: attn rows={attn_matrix.shape[0]} tokens={token_len}"
        )

    user_mask = [seg == "usr" for seg in segment_type]
    inst_mask = [cust == "inst" for cust in custom_mask]
    instr_mask = [bool(v) for v in instruction_mask]

    tokens = [
        tok.decode([token_id], skip_special_tokens=False).replace("\n", "\\n")
        for token_id in input_ids
    ]

    return tokens, inst_mask, instr_mask, user_mask, head_names, attn_matrix


def reorder_heads(head_names: List[str], attn_matrix: np.ndarray, requested_order: List[str]):
    index_by_head = {name: idx for idx, name in enumerate(head_names)}
    ordered_indices = []
    seen = set()

    for head_name in requested_order:
        idx = index_by_head.get(head_name)
        if idx is None or idx in seen:
            continue
        ordered_indices.append(idx)
        seen.add(idx)

    for idx, head_name in enumerate(head_names):
        if idx in seen:
            continue
        ordered_indices.append(idx)

    reordered_head_names = [head_names[idx] for idx in ordered_indices]
    reordered_attn_matrix = attn_matrix[:, ordered_indices]
    return reordered_head_names, reordered_attn_matrix


def render_image(
    tokens: List[str],
    inst_mask: List[bool],
    instr_mask: List[bool],
    user_mask: List[bool],
    head_names: List[str],
    attn_matrix: np.ndarray,
    output_path: str,
    title: str,
):
    token_len = len(tokens)
    num_heads = len(head_names)

    text_panel_width = 4.0
    width_per_head = 0.25
    colorbar_pad = 1.5

    heatmap_width = max(2.0, num_heads * width_per_head)
    total_width = text_panel_width + heatmap_width + colorbar_pad
    total_height = max(8, token_len * 0.18)

    fig = plt.figure(figsize=(total_width, total_height))
    gs = fig.add_gridspec(1, 2, width_ratios=[text_panel_width, heatmap_width], wspace=0.05)

    ax_text = fig.add_subplot(gs[0, 0])
    ax_heat = fig.add_subplot(gs[0, 1])

    ax_text.set_axis_off()
    ax_text.set_xlim(0, 1)
    ax_text.set_ylim(token_len - 0.5, -0.5)
    ax_text.text(
        0.0,
        -1.0,
        "idx  inst  instr  user  token",
        fontsize=9,
        fontfamily="monospace",
        fontweight="bold",
    )

    for i, (token, im, inrm, um) in enumerate(zip(tokens, inst_mask, instr_mask, user_mask)):
        bg_color = "white"
        if um:
            bg_color = "#e6f2ff"
        elif inrm:
            bg_color = "#fff2e6"
        elif im:
            bg_color = "#f2ffe6"

        ax_text.text(
            0.0,
            i,
            f"{i:03d}  {int(im):4d}  {int(inrm):5d}  {int(um):4d}  {token}",
            fontsize=9,
            fontfamily="monospace",
            va="center",
            bbox=dict(facecolor=bg_color, edgecolor="none", pad=1),
        )

    positive_values = attn_matrix[attn_matrix > 0]
    vmin = float(max(1e-6, positive_values.min())) if positive_values.size else 1e-6
    vmax = float(max(1.0, attn_matrix.max()))
    norm = LogNorm(vmin=vmin, vmax=vmax)

    im = ax_heat.imshow(
        attn_matrix,
        aspect="auto",
        interpolation="nearest",
        cmap="viridis",
        norm=norm,
    )
    ax_heat.set_yticks(range(token_len))
    ax_heat.set_yticklabels([f"{i:03d}" for i in range(token_len)], fontsize=7)

    max_ticks = 40
    step = max(1, num_heads // max_ticks)
    indices = range(0, num_heads, step)
    labels = [head_names[i] for i in indices]

    ax_heat.set_xticks(list(indices))
    ax_heat.set_xticklabels(labels, rotation=90, ha="center", fontsize=8)
    ax_heat.set_xlabel(f"Top {num_heads} Heads")
    ax_heat.set_ylabel("Tokens")

    cbar = fig.colorbar(im, ax=ax_heat, fraction=0.046, pad=0.04)
    cbar.set_label("Attention Weight (log scale)")

    fig.suptitle(title, y=0.99, fontsize=12)
    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    fig.savefig(output_path, bbox_inches="tight", dpi=150)
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(
        description="Visualize last-token attention for one arbitrary message list."
    )
    parser.add_argument("--model-path", required=True)
    parser.add_argument("--lora-path", default="")
    parser.add_argument("--input-path", required=True, help="Message list file")
    parser.add_argument("--output-path", default="head_heatmap_any.png")
    parser.add_argument("--head-order-path", default="")
    parser.add_argument("--title", default="")
    args = parser.parse_args()

    messages = load_message_list(args.input_path)
    model, tok = load_model_and_tokenizer(args.model_path, args.lora_path)
    tokens, inst_mask, instr_mask, user_mask, head_names, attn_matrix = collect_prompt_and_attention(
        model, tok, messages
    )
    if args.head_order_path:
        requested_order = load_head_order(args.head_order_path)
        head_names, attn_matrix = reorder_heads(head_names, attn_matrix, requested_order)

    title = args.title.strip()
    if not title:
        model_name = os.path.basename(os.path.normpath(args.model_path))
        lora_name = os.path.basename(os.path.normpath(args.lora_path)) if args.lora_path else "base"
        title = f"Model: {model_name} | LoRA: {lora_name}"

    render_image(
        tokens=tokens,
        inst_mask=inst_mask,
        instr_mask=instr_mask,
        user_mask=user_mask,
        head_names=head_names,
        attn_matrix=attn_matrix,
        output_path=args.output_path,
        title=title,
    )
    print(f"Saved image to: {args.output_path}")


if __name__ == "__main__":
    main()