from __future__ import annotations

import argparse
import io
import sys
from pathlib import Path
from typing import Any

import torch

sys.path.insert(0, str(Path(__file__).parent.parent))

from model_cpu_gpt2 import CPUGPT, gpt2_small_config, smoke_config


def _load_ckpt(path: str) -> dict:
    if path.startswith("s3://"):
        import boto3

        bucket, key = path[5:].split("/", 1)
        buf = io.BytesIO()
        boto3.client("s3").download_fileobj(bucket, key, buf)
        buf.seek(0)
        return torch.load(buf, map_location="cpu", weights_only=False)
    return torch.load(path, map_location="cpu", weights_only=False)


def build_our_model(config_name: str, ckpt_path: str | None) -> CPUGPT:
    cfg_map = {
        "gpt2-small": gpt2_small_config,
        "smoke": smoke_config,
    }
    if config_name not in cfg_map:
        raise ValueError(f"Unknown config '{config_name}'. Choices: {list(cfg_map)}")
    cfg = cfg_map[config_name]()
    model = CPUGPT(cfg)
    if ckpt_path:
        ckpt = _load_ckpt(ckpt_path)
        state = ckpt.get("model", ckpt)
        state = {k.replace("_orig_mod.", ""): v for k, v in state.items()}
        model.load_state_dict(state, strict=True)
        print(f"Loaded checkpoint: {ckpt_path}", flush=True)
    else:
        print("No checkpoint — using random weights.", flush=True)
    model.eval()
    return model


def build_hf_gpt2():
    try:
        from transformers import GPT2LMHeadModel, GPT2Tokenizer
    except ImportError:
        print(
            "\nERROR: transformers is not installed (needed for GPT-2 baseline).\n"
            "Install with:  pip install transformers\n",
            file=sys.stderr,
        )
        return None, None

    print("Loading HuggingFace gpt2...", flush=True)
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    model = GPT2LMHeadModel.from_pretrained("gpt2")
    model.eval()
    return model, tokenizer


def generate_our_model(
    model: CPUGPT,
    prompt_text: str,
    n_samples: int,
    temperature: float,
    top_p: float,
    max_new_tokens: int,
) -> list[str]:
    import tiktoken

    enc = tiktoken.get_encoding("gpt2")

    prompt_ids = enc.encode_ordinary(prompt_text)
    max_T = model.cfg.seq_len
    if len(prompt_ids) > max_T:
        prompt_ids = prompt_ids[-max_T:]
    prompt_tensor = torch.tensor(prompt_ids, dtype=torch.long).unsqueeze(0)

    top_k = max(1, int(50 * top_p))

    completions = []
    for _ in range(n_samples):
        with torch.no_grad():
            out = model.generate(
                prompt_tensor,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
            )
        gen_ids = out[0].tolist()
        completions.append(enc.decode(gen_ids))
    return completions


def generate_hf_gpt2(
    hf_model,
    tokenizer,
    prompt_text: str,
    n_samples: int,
    temperature: float,
    top_p: float,
    max_new_tokens: int,
) -> list[str]:
    inputs = tokenizer(prompt_text, return_tensors="pt")
    input_ids = inputs["input_ids"]

    completions = []
    for _ in range(n_samples):
        with torch.no_grad():
            if temperature == 0.0:
                out = hf_model.generate(
                    input_ids,
                    max_new_tokens=max_new_tokens,
                    do_sample=False,
                    pad_token_id=tokenizer.eos_token_id,
                )
            else:
                out = hf_model.generate(
                    input_ids,
                    max_new_tokens=max_new_tokens,
                    do_sample=True,
                    temperature=temperature,
                    top_p=top_p,
                    pad_token_id=tokenizer.eos_token_id,
                )
        gen_ids = out[0, input_ids.shape[1] :].tolist()
        gen_str = tokenizer.decode(gen_ids, skip_special_tokens=True)
        completions.append(gen_str)
    return completions


def load_prompts(yaml_path: str) -> list[dict[str, str]]:
    try:
        import yaml
    except ImportError:
        return _parse_prompts_fallback(yaml_path)
    with open(yaml_path) as f:
        data = yaml.safe_load(f)
    return data["prompts"]


def _parse_prompts_fallback(yaml_path: str) -> list[dict[str, str]]:
    prompts = []
    current: dict[str, str] = {}
    with open(yaml_path) as f:
        for line in f:
            line = line.rstrip()
            if line.startswith("  - category:"):
                if current:
                    prompts.append(current)
                current = {"category": line.split(":", 1)[1].strip()}
            elif line.startswith("    text:"):
                val = line.split(":", 1)[1].strip().strip('"')
                current["text"] = val
    if current:
        prompts.append(current)
    return prompts


def _escape_md(text: str) -> str:
    return text.replace("|", "\\|").replace("\n", " ").strip()


def build_markdown_table(
    prompts: list[dict[str, str]],
    our_completions: list[list[str]],
    hf_completions: list[list[str]] | None,
    n_samples: int,
) -> str:
    lines = []

    if hf_completions:
        lines.append("| # | Prompt | Our Model | GPT-2 Small |")
        lines.append("|---|--------|-----------|-------------|")
    else:
        lines.append("| # | Prompt | Our Model |")
        lines.append("|---|--------|-----------|")

    for pi, prompt in enumerate(prompts):
        prompt_text = _escape_md(prompt["text"])
        for si in range(n_samples):
            our_text = (
                _escape_md(our_completions[pi][si])
                if si < len(our_completions[pi])
                else ""
            )
            row_label = f"{pi + 1}.{si + 1}"
            if hf_completions:
                hf_text = (
                    _escape_md(hf_completions[pi][si])
                    if si < len(hf_completions[pi])
                    else ""
                )
                lines.append(
                    f"| {row_label} | {prompt_text} | {our_text} | {hf_text} |"
                )
            else:
                lines.append(f"| {row_label} | {prompt_text} | {our_text} |")

        if pi < len(prompts) - 1:
            if hf_completions:
                lines.append("| | | | |")
            else:
                lines.append("| | | |")

    return "\n".join(lines)


def run_gen_samples(
    config_name: str,
    prompts: list[dict[str, str]],
    n_samples: int,
    temperature: float,
    top_p: float,
    max_new_tokens: int,
    ckpt_path: str | None = None,
    include_hf_baseline: bool = True,
) -> str:
    our_model = build_our_model(config_name, ckpt_path)

    hf_model, hf_tokenizer = None, None
    if include_hf_baseline:
        hf_model, hf_tokenizer = build_hf_gpt2()

    our_all: list[list[str]] = []
    hf_all: list[list[str]] = []

    for prompt in prompts:
        print(f"  [{prompt['category']}] {prompt['text'][:60]}...", flush=True)
        our_completions = generate_our_model(
            our_model, prompt["text"], n_samples, temperature, top_p, max_new_tokens
        )
        our_all.append(our_completions)

        if hf_model is not None:
            hf_completions = generate_hf_gpt2(
                hf_model,
                hf_tokenizer,
                prompt["text"],
                n_samples,
                temperature,
                top_p,
                max_new_tokens,
            )
            hf_all.append(hf_completions)

    table = build_markdown_table(
        prompts, our_all, hf_all if hf_model else None, n_samples
    )
    return table


def main():
    p = argparse.ArgumentParser(
        description="Generate qualitative samples (our model vs GPT-2)"
    )
    p.add_argument(
        "--ckpt", default=None, help="Checkpoint path (omit for random weights)"
    )
    p.add_argument(
        "--config", default="gpt2-small", help="Model config: gpt2-small | smoke"
    )
    p.add_argument(
        "--prompts",
        default="scripts/eval_prompts.yaml",
        help="Path to eval_prompts.yaml",
    )
    p.add_argument("--n-samples", type=int, default=3, help="Completions per prompt")
    p.add_argument("--temperature", type=float, default=0.8)
    p.add_argument("--top-p", type=float, default=0.9)
    p.add_argument("--max-new-tokens", type=int, default=100)
    p.add_argument(
        "--output",
        required=True,
        help="Output markdown file (e.g. results/samples_v5.md)",
    )
    p.add_argument(
        "--no-baseline", action="store_true", help="Skip HuggingFace GPT-2 baseline"
    )
    args = p.parse_args()

    print(f"\n=== gen_samples.py ===", flush=True)
    print(f"Config:      {args.config}", flush=True)
    print(f"Ckpt:        {args.ckpt or '(random weights)'}", flush=True)
    print(f"Prompts:     {args.prompts}", flush=True)
    print(f"N-samples:   {args.n_samples}", flush=True)
    print(f"Temperature: {args.temperature}", flush=True)
    print(f"Top-p:       {args.top_p}", flush=True)
    print(f"Max tokens:  {args.max_new_tokens}", flush=True)
    print(flush=True)

    prompts_path = Path(args.prompts)
    if not prompts_path.exists():
        print(f"ERROR: prompts file not found: {prompts_path}", file=sys.stderr)
        sys.exit(1)
    prompts = load_prompts(str(prompts_path))
    print(f"Loaded {len(prompts)} prompts from {prompts_path}", flush=True)

    table = run_gen_samples(
        config_name=args.config,
        prompts=prompts,
        n_samples=args.n_samples,
        temperature=args.temperature,
        top_p=args.top_p,
        max_new_tokens=args.max_new_tokens,
        ckpt_path=args.ckpt,
        include_hf_baseline=not args.no_baseline,
    )

    print("\n" + "=" * 60)
    print("Sample outputs")
    print("=" * 60)
    print(table)

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w") as f:
        f.write(f"# Qualitative Generation Samples\n\n")
        f.write(f"Config: `{args.config}`  \n")
        f.write(f"Checkpoint: `{args.ckpt or 'random_weights'}`  \n")
        f.write(
            f"Temperature: {args.temperature}, Top-p: {args.top_p}, "
            f"Max new tokens: {args.max_new_tokens}\n\n"
        )
        f.write(table)
        f.write("\n")
    print(f"\nSaved to: {out_path}", flush=True)


if __name__ == "__main__":
    main()