import os
from typing import List, Tuple

import torch
import pandas as pd
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig


# ============================================================
# Settings
# ============================================================

MODEL_ID = os.getenv("MODEL_ID", "llm-jp/llm-jp-4-8b-instruct")
LOAD_IN_4BIT = os.getenv("LOAD_IN_4BIT", "true").lower() == "true"
USE_CHAT_TEMPLATE = os.getenv("USE_CHAT_TEMPLATE", "true").lower() == "true"
USE_FAST_TOKENIZER = os.getenv("USE_FAST_TOKENIZER", "true").lower() == "true"

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


# ============================================================
# Model / tokenizer loading
# ============================================================

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    use_fast=USE_FAST_TOKENIZER,
)

if LOAD_IN_4BIT:
    if device != "cuda":
        raise RuntimeError(
            "LOAD_IN_4BIT=true requires a CUDA GPU Space. "
            "Please switch the Space hardware to T4/L4 or set LOAD_IN_4BIT=false."
        )

    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=quantization_config,
        device_map="auto",
        torch_dtype=torch.bfloat16,
    )
else:
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    )
    model.to(device)

if tokenizer.pad_token_id is None:
    if tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token

if model.config.pad_token_id is None and tokenizer.pad_token_id is not None:
    model.config.pad_token_id = tokenizer.pad_token_id

model.eval()


# ============================================================
# Utility functions
# ============================================================


def build_model_prompt_text(user_prompt: str) -> str:
    """
    Convert the user's raw prompt into the actual model prompt.
    For instruct/chat models, use the tokenizer chat template.
    """
    if USE_CHAT_TEMPLATE and getattr(tokenizer, "chat_template", None):
        messages = [
            {"role": "user", "content": user_prompt}
        ]
        return tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )

    return user_prompt


def encode_model_prompt(user_prompt: str, max_prompt_tokens: int) -> List[int]:
    model_prompt_text = build_model_prompt_text(user_prompt)

    enc = tokenizer(
        model_prompt_text,
        add_special_tokens=False,
        truncation=True,
        max_length=int(max_prompt_tokens),
    )

    return enc["input_ids"]
    

def visible_token_text(token_id: int) -> str:
    text = tokenizer.decode([token_id], skip_special_tokens=False)
    text = text.replace("\n", "\\n").replace("\t", "\\t").replace(" ", "␠")
    return text if text != "" else "<empty>"


def raw_token_text(token_id: int) -> str:
    token = tokenizer.convert_ids_to_tokens([token_id])[0]
    return token.replace("\n", "\\n").replace("\t", "\\t").replace(" ", "␠")


def encode_prompt(prompt: str, max_prompt_tokens: int) -> List[int]:
    enc = tokenizer(
        prompt,
        add_special_tokens=False,
        truncation=True,
        max_length=int(max_prompt_tokens),
    )
    return enc["input_ids"]


def remove_special_tokens(token_ids: List[int]) -> List[int]:
    special_ids = set(tokenizer.all_special_ids)
    return [tid for tid in token_ids if tid not in special_ids]


# ============================================================
# Generation
# ============================================================

def generate_answer_from_prompt_ids(
    prompt_ids: List[int],
    max_new_tokens: int,
) -> Tuple[str, List[int]]:
    input_ids = torch.tensor(
        [prompt_ids],
        dtype=torch.long,
        device=device,
    )

    attention_mask = torch.ones_like(input_ids, device=device)

    generation_kwargs = {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "max_new_tokens": int(max_new_tokens),
        "do_sample": False,
        "use_cache": True,
    }

    if tokenizer.pad_token_id is not None:
        generation_kwargs["pad_token_id"] = tokenizer.pad_token_id

    if tokenizer.eos_token_id is not None:
        generation_kwargs["eos_token_id"] = tokenizer.eos_token_id

    with torch.no_grad():
        generated = model.generate(**generation_kwargs)

    generated_ids = generated[0].detach().cpu().tolist()
    answer_ids = generated_ids[len(prompt_ids):]
    answer_ids = remove_special_tokens(answer_ids)

    answer_text = tokenizer.decode(
        answer_ids,
        skip_special_tokens=True,
    )

    return answer_text, answer_ids


# ============================================================
# Logprob scoring
# ============================================================

def score_answer_logprobs_from_ids(
    prompt_ids: List[int],
    answer_ids: List[int],
) -> torch.Tensor:
    if len(prompt_ids) == 0:
        raise ValueError("prompt_ids is empty.")

    if len(answer_ids) == 0:
        raise ValueError("answer_ids is empty.")

    input_ids = torch.tensor(
        [prompt_ids + answer_ids],
        dtype=torch.long,
        device=device,
    )

    attention_mask = torch.ones_like(input_ids, device=device)

    with torch.no_grad():
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )

    logits = outputs.logits

    prompt_len = len(prompt_ids)
    answer_len = len(answer_ids)

    # causal LMでは logits[:, k, :] が input_ids[:, k + 1] を予測する。
    pred_positions = torch.arange(
        prompt_len - 1,
        prompt_len + answer_len - 1,
        device=device,
    )

    target_positions = torch.arange(
        prompt_len,
        prompt_len + answer_len,
        device=device,
    )

    selected_logits = logits[:, pred_positions, :]
    target_ids = input_ids[:, target_positions]

    log_probs = torch.log_softmax(selected_logits, dim=-1)

    target_log_probs = log_probs.gather(
        dim=-1,
        index=target_ids.unsqueeze(-1),
    ).squeeze(-1)

    return target_log_probs.squeeze(0)


# ============================================================
# Attribution analysis
# ============================================================

def analyze_prompt_token_influence(
    prompt_ids: List[int],
    answer_ids: List[int],
    low_influence_threshold: float,
) -> pd.DataFrame:
    base_logps = score_answer_logprobs_from_ids(
        prompt_ids=prompt_ids,
        answer_ids=answer_ids,
    )

    base_sum_logprob = float(base_logps.sum().item())
    base_mean_logprob = float(base_logps.mean().item())

    rows = []

    for pos, token_id in enumerate(prompt_ids):
        ablated_prompt_ids = prompt_ids[:pos] + prompt_ids[pos + 1:]

        if len(ablated_prompt_ids) == 0:
            rows.append(
                {
                    "position": pos,
                    "token": visible_token_text(token_id),
                    "raw_token": raw_token_text(token_id),
                    "token_id": token_id,
                    "ablated_prompt": "<empty prompt: skipped>",
                    "base_mean_logprob": base_mean_logprob,
                    "base_sum_logprob": base_sum_logprob,
                    "ablated_mean_logprob": None,
                    "ablated_sum_logprob": None,
                    "mean_delta_logprob": None,
                    "sum_delta_logprob": None,
                    "mean_abs_delta_logprob": None,
                    "judgment": "skipped",
                }
            )
            continue

        ablated_logps = score_answer_logprobs_from_ids(
            prompt_ids=ablated_prompt_ids,
            answer_ids=answer_ids,
        )

        delta = base_logps - ablated_logps

        mean_delta = float(delta.mean().item())
        sum_delta = float(delta.sum().item())
        mean_abs_delta = float(delta.abs().mean().item())

        ablated_mean = float(ablated_logps.mean().item())
        ablated_sum = float(ablated_logps.sum().item())

        if abs(mean_delta) <= low_influence_threshold:
            judgment = "low influence"
        elif mean_delta > 0:
            judgment = "supports answer"
        else:
            judgment = "suppresses answer"

        ablated_prompt_text = tokenizer.decode(
            ablated_prompt_ids,
            skip_special_tokens=False,
        )

        rows.append(
            {
                "position": pos,
                "token": visible_token_text(token_id),
                "raw_token": raw_token_text(token_id),
                "token_id": token_id,
                "ablated_prompt": ablated_prompt_text,
                "base_mean_logprob": base_mean_logprob,
                "base_sum_logprob": base_sum_logprob,
                "ablated_mean_logprob": ablated_mean,
                "ablated_sum_logprob": ablated_sum,
                "mean_delta_logprob": mean_delta,
                "sum_delta_logprob": sum_delta,
                "mean_abs_delta_logprob": mean_abs_delta,
                "judgment": judgment,
            }
        )

    df = pd.DataFrame(rows)

    if "mean_delta_logprob" in df.columns:
        df["_sort_key"] = df["mean_delta_logprob"].apply(
            lambda x: float("-inf") if pd.isna(x) else float(x)
        )
        df = df.sort_values(
            by="_sort_key",
            ascending=False,
        ).drop(columns=["_sort_key"])

    return df.reset_index(drop=True)


# ============================================================
# Gradio handler
# ============================================================

def make_summary(
    prompt_ids: List[int],
    answer_ids: List[int],
    used_prompt_text: str,
    base_mean: float,
    base_sum: float,
) -> str:
    lines = [
        "### Result",
        "",
        f"- Model: `{MODEL_ID}`",
        f"- Device: `{device}`",
        f"- Prompt tokens used: `{len(prompt_ids)}`",
        f"- Generated answer tokens scored: `{len(answer_ids)}`",
        f"- Base mean logprob of generated answer: `{base_mean:.6f}`",
        f"- Base sum logprob of generated answer: `{base_sum:.6f}`",
        "",
        "### Used prompt after tokenization / truncation",
        "",
        "```text",
        used_prompt_text,
        "```",
        "",
        "### Score definition",
        "",
        "`mean_delta_logprob` は次を表します。",
        "",
        "```text",
        "平均 [",
        "  log P(生成済み回答トークン | 元プロンプト)",
        "  -",
        "  log P(生成済み回答トークン | その prompt token を削除したプロンプト)",
        "]",
        "```",
        "",
        "読み方:",
        "",
        "- 正の値が大きい: その prompt token は生成済み回答を出しやすくしていた",
        "- 0 に近い: 削除しても生成済み回答の出やすさがあまり変わらない",
        "- 負の値: その固定回答に対しては、むしろ抑制的だった可能性がある",
    ]

    return "\n".join(lines)


def generate_and_analyze(
    prompt: str,
    max_prompt_tokens: int,
    max_new_tokens: int,
    low_influence_threshold: float,
):
    if prompt is None or prompt.strip() == "":
        return "プロンプトが空です。", "", pd.DataFrame()

    try:
        prompt_ids = encode_model_prompt(
            user_prompt=prompt,
            max_prompt_tokens=int(max_prompt_tokens),
        )

        if len(prompt_ids) == 0:
            return "有効な prompt token がありません。", "", pd.DataFrame()

        used_prompt_text = tokenizer.decode(
            prompt_ids,
            skip_special_tokens=False,
        )

        answer_text, answer_ids = generate_answer_from_prompt_ids(
            prompt_ids=prompt_ids,
            max_new_tokens=int(max_new_tokens),
        )

        if len(answer_ids) == 0:
            lines = [
                "### Result",
                "",
                f"- Model: `{MODEL_ID}`",
                f"- Device: `{device}`",
                f"- Prompt tokens used: `{len(prompt_ids)}`",
                "- Generated answer tokens: `0`",
                "",
                "生成された回答トークンが空でした。`max_new_tokens` を増やすか、別のモデルを指定してください。",
            ]
            return "\n".join(lines), answer_text, pd.DataFrame()

        df = analyze_prompt_token_influence(
            prompt_ids=prompt_ids,
            answer_ids=answer_ids,
            low_influence_threshold=float(low_influence_threshold),
        )

        base_mean = float(df["base_mean_logprob"].dropna().iloc[0])
        base_sum = float(df["base_sum_logprob"].dropna().iloc[0])

        summary = make_summary(
            prompt_ids=prompt_ids,
            answer_ids=answer_ids,
            used_prompt_text=used_prompt_text,
            base_mean=base_mean,
            base_sum=base_sum,
        )

        return summary, answer_text, df

    except Exception as e:
        return f"エラーが発生しました: `{type(e).__name__}: {e}`", "", pd.DataFrame()


# ============================================================
# UI
# ============================================================

with gr.Blocks() as demo:
    gr.Markdown(
        "# LLM Prompt Token Influence Probe\n\n"
        "プロンプトから回答を一度生成し、その回答を固定したうえで、"
        "プロンプト内の各 token が回答全体にどれだけ影響したかを推定します。\n\n"
        "評価単位は、ユーザー指定ではなく、モデルの tokenizer が実際に分割した prompt token です。"
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                lines=6,
                value="自己紹介をしてください",
            )

            max_prompt_tokens = gr.Slider(
                minimum=8,
                maximum=512,
                value=128,
                step=1,
                label="Max prompt tokens",
            )

            max_new_tokens = gr.Slider(
                minimum=1,
                maximum=256,
                value=64,
                step=1,
                label="Max new tokens",
            )

            low_influence_threshold = gr.Number(
                value=0.02,
                label="Low influence threshold",
            )

            run_button = gr.Button("Generate & Analyze")

        with gr.Column(scale=2):
            summary = gr.Markdown(label="Summary")

            answer = gr.Textbox(
                label="Generated answer",
                lines=6,
                interactive=False,
            )

            table = gr.Dataframe(
                label="Prompt token influence",
                wrap=True,
                interactive=False,
            )

    run_button.click(
        fn=generate_and_analyze,
        inputs=[
            prompt,
            max_prompt_tokens,
            max_new_tokens,
            low_influence_threshold,
        ],
        outputs=[
            summary,
            answer,
            table,
        ],
    )


if __name__ == "__main__":
    demo.queue().launch()