from __future__ import annotations

import json
from dataclasses import asdict
from pathlib import Path

import torch

from .checkpoints import load_token_checkpoint
from .config import PipelineConfig, VerbalizationTokenSetConfig
from .modeling import ModelBundle


def apply_token_pair(bundle: ModelBundle, token_dir: Path) -> None:
    input_emb = bundle.model.get_input_embeddings()
    ai_ckpt = load_token_checkpoint(token_dir / "ai_token.pt")
    human_ckpt = load_token_checkpoint(token_dir / "human_token.pt")
    ai_id = bundle.tokenizer.convert_tokens_to_ids(bundle.config.model.ai_token)
    human_id = bundle.tokenizer.convert_tokens_to_ids(bundle.config.model.human_token)
    with torch.no_grad():
        input_emb.weight[ai_id].copy_(ai_ckpt.embedding.to(input_emb.weight.device, dtype=input_emb.weight.dtype))
        input_emb.weight[human_id].copy_(
            human_ckpt.embedding.to(input_emb.weight.device, dtype=input_emb.weight.dtype)
        )


@torch.inference_mode()
def sample_text(
    bundle: ModelBundle,
    prompt: str,
    *,
    max_new_tokens: int,
    do_sample: bool,
    temperature: float,
    top_p: float,
) -> str:
    input_ids = bundle.tokenizer(prompt, return_tensors="pt", add_special_tokens=False)["input_ids"].to(
        bundle.model.device
    )
    output = bundle.model.generate(
        input_ids=input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=do_sample,
        temperature=temperature,
        top_p=top_p,
        pad_token_id=bundle.tokenizer.eos_token_id,
    )
    return bundle.tokenizer.decode(output[0, input_ids.shape[1] :], skip_special_tokens=True).strip()


def _chat_prompt(bundle: ModelBundle, content: str) -> str:
    return bundle.tokenizer.apply_chat_template(
        [{"role": "user", "content": content}],
        tokenize=False,
        add_generation_prompt=True,
    )


@torch.inference_mode()
def verbalize_token(bundle: ModelBundle, token: str, *, max_new_tokens: int, do_sample: bool) -> str:
    prompt = _chat_prompt(bundle, f"Describe what qualities of text would be implied by {token}.")
    prompt += f"{token} text typically describes text that"
    return sample_text(
        bundle,
        prompt,
        max_new_tokens=max_new_tokens,
        do_sample=do_sample,
        temperature=1.0,
        top_p=0.95 if do_sample else 1.0,
    )


@torch.inference_mode()
def verbalize_difference(bundle: ModelBundle, *, max_new_tokens: int, do_sample: bool) -> str:
    ai_token = bundle.config.model.ai_token
    human_token = bundle.config.model.human_token
    prompt = _chat_prompt(
        bundle,
        (
            f"What is the difference between {ai_token} text and {human_token} text? "
            f"Refer to {ai_token} text as A-type text and {human_token} text as B-type text. "
            "Do not discuss the literal token strings; describe the passage qualities they accompany."
        ),
    )
    prompt += (
        f"{ai_token}, which I will refer to as A-type text, and {human_token}, "
        "B-type text, have many similarities and differences."
    )
    return sample_text(
        bundle,
        prompt,
        max_new_tokens=max_new_tokens,
        do_sample=do_sample,
        temperature=1.0,
        top_p=0.95 if do_sample else 1.0,
    )


def verbalize_token_set(
    bundle: ModelBundle,
    token_set: VerbalizationTokenSetConfig,
    config: PipelineConfig,
) -> dict:
    apply_token_pair(bundle, token_set.token_dir)
    payload = {
        "name": token_set.name,
        "token_dir": str(token_set.token_dir),
        "model_name": config.model.model_name,
        "ai_token": config.model.ai_token,
        "human_token": config.model.human_token,
        "greedy": {
            "ai": verbalize_token(
                bundle,
                config.model.ai_token,
                max_new_tokens=config.verbalization.max_new_tokens,
                do_sample=False,
            ),
            "human": verbalize_token(
                bundle,
                config.model.human_token,
                max_new_tokens=config.verbalization.max_new_tokens,
                do_sample=False,
            ),
            "difference": verbalize_difference(
                bundle,
                max_new_tokens=config.verbalization.max_new_tokens,
                do_sample=False,
            ),
        },
        "samples": [],
    }
    for _ in range(config.verbalization.n_samples):
        payload["samples"].append(
            {
                "ai": verbalize_token(
                    bundle,
                    config.model.ai_token,
                    max_new_tokens=config.verbalization.max_new_tokens,
                    do_sample=True,
                ),
                "human": verbalize_token(
                    bundle,
                    config.model.human_token,
                    max_new_tokens=config.verbalization.max_new_tokens,
                    do_sample=True,
                ),
                "difference": verbalize_difference(
                    bundle,
                    max_new_tokens=config.verbalization.max_new_tokens,
                    do_sample=True,
                ),
            }
        )
    return payload


def save_verbalizations(output_dir: Path, results: list[dict]) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "verbalizations.json").write_text(json.dumps(results, indent=2), encoding="utf-8")
    for result in results:
        (output_dir / f"{result['name']}.json").write_text(json.dumps(result, indent=2), encoding="utf-8")