import torch
from transformer_lens import HookedTransformer
from model_config import get_config

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

_models: dict[str, HookedTransformer] = {}


def load_model(model_id: str = "llama-3.2-1b") -> HookedTransformer:
    global _models
    if model_id not in _models:
        cfg = get_config(model_id)
        print(f"[load_model] loading {cfg['hf_name']}...", flush=True)
        m = HookedTransformer.from_pretrained(cfg["hf_name"])
        m.eval()
        m.to(DEVICE)
        _models[model_id] = m
        print(f"[load_model] {model_id} ready.", flush=True)
    return _models[model_id]


def _format_prompt(model: HookedTransformer, prompt: str) -> str:
    """Use chat template if available, otherwise use prompt directly."""
    if hasattr(model.tokenizer, "chat_template") and model.tokenizer.chat_template:
        messages = [{"role": "user", "content": prompt}]
        return model.tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
    return prompt


def _get_embeds(tokens, model_id: str = "llama-3.2-1b"):
    m = load_model(model_id)
    with torch.no_grad():
        return m.W_E[tokens]


def _corrupt_all(tokens, noise_scale, model_id: str = "llama-3.2-1b"):
    embeds = _get_embeds(tokens, model_id)
    return embeds + torch.randn_like(embeds) * noise_scale


def _corrupt_single_pos(tokens, pos, noise_scale, model_id: str = "llama-3.2-1b"):
    embeds = _get_embeds(tokens, model_id)
    noise = torch.zeros_like(embeds)
    noise[0, pos] = torch.randn(embeds.shape[-1], device=embeds.device) * noise_scale
    return embeds + noise


def run_trace(
    prompt: str,
    target_token: str,
    model_id: str = "llama-3.2-1b",
    noise_scale: float = 3.0,
    n_noise_runs: int = 10,
):
    m = load_model(model_id)
    tokens = m.to_tokens(prompt)
    target_id = m.to_tokens(target_token, prepend_bos=False)[0, 0].item()

    with torch.no_grad():
        clean_logits, clean_cache = m.run_with_cache(tokens)
    baseline_prob = torch.softmax(clean_logits[0, -1], dim=-1)[target_id].item()

    corrupted_prob = 0.0
    for _ in range(n_noise_runs):
        ne = _corrupt_all(tokens, noise_scale, model_id)
        with torch.no_grad():
            cl = m.run_with_hooks(tokens, fwd_hooks=[("hook_embed", lambda v, hook=None: ne.to(v.device))])
        corrupted_prob += torch.softmax(cl[0, -1], dim=-1)[target_id].item()
    corrupted_prob /= n_noise_runs
    corruption_effect = baseline_prob - corrupted_prob

    results = []
    for layer in range(m.cfg.n_layers):
        patch_probs = []
        for _ in range(n_noise_runs):
            ne = _corrupt_all(tokens, noise_scale, model_id)
            def make_hooks(ne_, li):
                return [
                    ("hook_embed", lambda v, hook=None, n=ne_: n.to(v.device)),
                    (f"blocks.{li}.hook_resid_post", lambda v, hook=None, li=li: (
                        v.__setitem__(
                            (slice(None), -1, slice(None)),
                            clean_cache[f"blocks.{li}.hook_resid_post"][:, -1, :].to(v.device)
                        ) or v
                    )),
                ]
            with torch.no_grad():
                pl = m.run_with_hooks(tokens, fwd_hooks=make_hooks(ne, layer))
            patch_probs.append(torch.softmax(pl[0, -1], dim=-1)[target_id].item())

        restored = sum(patch_probs) / len(patch_probs)
        drop = round(max(restored - corrupted_prob, 0.0) / max(corruption_effect, 1e-6), 4)
        results.append({
            "layer":          layer,
            "baseline_prob":  round(baseline_prob, 4),
            "corrupted_prob": round(corrupted_prob, 4),
            "corruption_effect": round(corruption_effect, 4),
            "drop":           drop,
            "attn_drop":      0.0,
            "mlp_drop":       0.0,
            "patched_prob":   round(restored, 4),
            "attn_patched_prob": 0.0,
            "mlp_patched_prob":  0.0,
        })

    torch.cuda.synchronize()
    return sorted(results, key=lambda x: x["drop"], reverse=True)


def run_prompt_attribution(
    prompt, response, prompt_tokens, response_tokens,
    sig_prompt_tis, sig_response_tis,
    model_id: str = "llama-3.2-1b",
    noise_scale: float = 3.0,
    n_noise_runs: int = 5,
):
    m = load_model(model_id)
    full_ctx = f"{prompt}\n{response}"
    tokens = m.to_tokens(full_ctx)
    n_tokens = tokens.shape[1]

    with torch.no_grad():
        clean_logits = m(tokens)

    target_ids, baseline_probs = {}, {}
    for rti in sig_response_tis:
        tok_str = f" {response_tokens[rti].strip()}"
        try:
            tid = m.to_tokens(tok_str, prepend_bos=False)[0, 0].item()
        except Exception:
            continue
        target_ids[rti] = tid
        baseline_probs[rti] = torch.softmax(clean_logits[0, -1], dim=-1)[tid].item()

    if not target_ids:
        return {"attributions": []}

    prefix_len = m.to_tokens(prompt + "\n", prepend_bos=True).shape[1]

    def get_token_pos(word_idx):
        words = prompt.split()
        if word_idx >= len(words):
            return None
        char_pos = sum(len(w) + 1 for w in words[:word_idx])
        try:
            return m.to_tokens(prompt[:char_pos], prepend_bos=True).shape[1] - 1
        except Exception:
            return None

    scores = {rti: {} for rti in target_ids}
    for pti in sig_prompt_tis:
        tok_pos = get_token_pos(pti)
        if tok_pos is None or tok_pos >= n_tokens:
            continue

        corrupted_probs = {rti: 0.0 for rti in target_ids}
        for _ in range(n_noise_runs):
            ne = _corrupt_single_pos(tokens, tok_pos, noise_scale, model_id)
            with torch.no_grad():
                logits = m.run_with_hooks(
                    tokens,
                    fwd_hooks=[("hook_embed", lambda v, hook=None, n=ne: n.to(v.device))]
                )
            probs = torch.softmax(logits[0, -1], dim=-1)
            for rti, tid in target_ids.items():
                corrupted_probs[rti] += probs[tid].item()

        for rti in target_ids:
            corrupted_probs[rti] /= n_noise_runs
            scores[rti][pti] = max(round(baseline_probs[rti] - corrupted_probs[rti], 4), 0.0)

    attributions = []
    for rti, tid in target_ids.items():
        raw = scores[rti]
        mx = max(raw.values(), default=1e-6) or 1e-6
        prompt_scores = sorted([
            {
                "prompt_ti":    pti,
                "prompt_token": prompt_tokens[pti],
                "score":        round(raw.get(pti, 0.0) / mx, 4),
                "raw_score":    raw.get(pti, 0.0),
            }
            for pti in sig_prompt_tis
        ], key=lambda x: x["score"], reverse=True)
        attributions.append({
            "response_ti":    rti,
            "response_token": response_tokens[rti],
            "baseline_prob":  round(baseline_probs[rti], 4),
            "prompt_scores":  prompt_scores,
        })

    torch.cuda.synchronize()
    return {"attributions": attributions}


def run_logit_lens(
    prompt: str,
    model_id: str = "llama-3.2-1b",
    top_k: int = 5,
):
    m = load_model(model_id)
    tokens = m.to_tokens(prompt)
    with torch.no_grad():
        _, cache = m.run_with_cache(tokens)

    results = []
    for layer in range(m.cfg.n_layers):
        resid = cache[f"blocks.{layer}.hook_resid_post"][0, -1, :]
        resid_normed = m.ln_final(resid.unsqueeze(0).unsqueeze(0))[0, 0, :]
        logits = resid_normed @ m.unembed.W_U + m.unembed.b_U
        probs = torch.softmax(logits, dim=-1)
        topk = probs.topk(top_k)
        results.append({
            "layer": layer,
            "top_tokens": [
                {"token": m.to_string([tid.item()]), "prob": round(p.item(), 4)}
                for tid, p in zip(topk.indices, topk.values)
            ],
        })

    torch.cuda.synchronize()
    return results


def run_chat(
    prompt: str,
    model_id: str = "llama-3.2-1b",
    max_new_tokens: int = 200,
    temperature: float = 0.7,
) -> str:
    m = load_model(model_id)
    formatted = _format_prompt(m, prompt)
    input_ids = m.tokenizer(formatted, return_tensors="pt").input_ids.to(DEVICE)
    generated = []

    with torch.no_grad():
        cur = input_ids
        for _ in range(max_new_tokens):
            logits = m(cur)
            next_logits = logits[0, -1] / max(temperature, 1e-6)
            probs = torch.softmax(next_logits, dim=-1)
            next_id = int(torch.multinomial(probs, 1).item())
            if next_id == m.tokenizer.eos_token_id:
                break
            generated.append(next_id)
            cur = torch.cat([cur, torch.tensor([[next_id]], device=DEVICE)], dim=1)

    return m.tokenizer.decode(generated, skip_special_tokens=True).strip()


def stream_chat(
    prompt: str,
    model_id: str = "llama-3.2-1b",
    max_new_tokens: int = 200,
    temperature: float = 0.7,
):
    m = load_model(model_id)
    formatted = _format_prompt(m, prompt)
    input_ids = m.tokenizer(formatted, return_tensors="pt").input_ids.to(DEVICE)

    with torch.no_grad():
        cur = input_ids
        for _ in range(max_new_tokens):
            logits = m(cur)
            next_logits = logits[0, -1] / max(temperature, 1e-6)
            probs = torch.softmax(next_logits, dim=-1)
            next_id = int(torch.multinomial(probs, 1).item())
            if next_id == m.tokenizer.eos_token_id:
                break
            tok_str = m.tokenizer.decode([next_id], skip_special_tokens=True)
            yield tok_str
            cur = torch.cat([cur, torch.tensor([[next_id]], device=DEVICE)], dim=1)