import argparse
import os
import pickle

from transformers import AutoTokenizer


def _load_pickle(path: str):
    if not os.path.isfile(path):
        raise FileNotFoundError(f"Pickle not found: {path}")
    with open(path, "rb") as f:
        return pickle.load(f)


def _as_bool_list(values, name: str, expected_len: int):
    if len(values) != expected_len:
        raise ValueError(f"{name} length {len(values)} != token length {expected_len}")
    return [bool(v) for v in values]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-path", required=True)
    parser.add_argument("--prompt-index", type=int, default=0)
    parser.add_argument("--head-name", required=True)
    parser.add_argument(
        "--model-path",
        default="../../models/Llama-3.1-8B-Instruct",
    )
    args = parser.parse_args()

    raw = _load_pickle(args.input_path)
    prompts = raw.get("prompts", [])
    heads = raw.get("heads", {})
    if not prompts:
        raise ValueError("No prompts found in pickle.")
    if args.prompt_index < 0 or args.prompt_index >= len(prompts):
        raise IndexError(
            f"prompt-index {args.prompt_index} out of range (0..{len(prompts) - 1})"
        )
    if args.head_name not in heads:
        head_keys = ", ".join(sorted(heads.keys()))
        raise KeyError(f"Unknown head-name {args.head_name}. Available: {head_keys}")

    prompt = prompts[args.prompt_index]
    token_ids = list(prompt["token_ids"])
    token_len = len(token_ids)
    inst_mask = _as_bool_list(prompt["inst_mask"], "inst_mask", token_len)
    instr_mask = _as_bool_list(prompt["instr_mask"], "instr_mask", token_len)
    user_mask = _as_bool_list(prompt["user_mask"], "user_mask", token_len)
    attn_weight = heads[args.head_name]["attn_weight"][args.prompt_index]
    if len(attn_weight) != token_len:
        raise ValueError(
            f"attn_weight length {len(attn_weight)} != token length {token_len}"
        )

    tok = AutoTokenizer.from_pretrained(args.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"

    tokens = [tok.decode([tid], skip_special_tokens=False) for tid in token_ids]
    print(f"pickle: {args.input_path}")
    print(f"prompt_index: {args.prompt_index}")
    print(f"head_name: {args.head_name}")
    print("idx\tinst\tinstr\tuser\tattn_weight\ttoken")
    for i, (t, im, inrm, um, aw) in enumerate(
        zip(tokens, inst_mask, instr_mask, user_mask, attn_weight)
    ):
        token_str = t.replace("\n", "\\n")
        print(f"{i:03d}\t{int(im)}\t{int(inrm)}\t{int(um)}\t{float(aw):.6f}\t{token_str}")


if __name__ == "__main__":
    main()
