#!/usr/bin/env python3

from __future__ import annotations

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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Generate text with ObsidianSmall-Base using "
            "full-sequence Multiscreen decoding."
        )
    )

    parser.add_argument(
        "--model-dir",
        type=Path,
        default=Path(__file__).resolve().parent,
    )
    parser.add_argument("--prompt", type=str)
    parser.add_argument("--prompt-file", type=Path)
    parser.add_argument("--interactive", action="store_true")

    parser.add_argument(
        "--backend",
        choices=("auto", "torch", "triton"),
        default="auto",
    )
    parser.add_argument(
        "--device",
        choices=("auto", "cpu", "cuda"),
        default="auto",
    )
    parser.add_argument(
        "--dtype",
        choices=("auto", "float32", "float16", "bfloat16"),
        default="auto",
    )

    parser.add_argument("--max-new-tokens", type=int, default=128)
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top-k", type=int, default=40)
    parser.add_argument("--top-p", type=float, default=1.0)
    parser.add_argument("--seed", type=int, default=1337)

    return parser.parse_args()


def extract_state_dict(payload: Any, torch: Any) -> dict[str, Any]:
    if isinstance(payload, dict):
        for container in ("model", "state_dict", "model_state_dict"):
            candidate = payload.get(container)

            if isinstance(candidate, dict):
                payload = candidate
                break

    if not isinstance(payload, dict):
        raise TypeError(
            f"Unsupported checkpoint type: {type(payload).__name__}"
        )

    state_dict: dict[str, Any] = {}

    for original_name, value in payload.items():
        if not torch.is_tensor(value):
            continue

        name = str(original_name)

        for prefix in ("_orig_mod.", "module.", "model."):
            while name.startswith(prefix):
                name = name[len(prefix):]

        state_dict[name] = value

    if not state_dict:
        raise RuntimeError("No model tensors were found")

    return state_dict


def sample_token(
    logits: Any,
    torch: Any,
    temperature: float,
    top_k: int,
    top_p: float,
) -> Any:
    if temperature <= 0:
        return logits.argmax(dim=-1, keepdim=True)

    logits = logits / temperature

    if top_k > 0:
        k = min(top_k, logits.size(-1))
        threshold = torch.topk(logits, k).values[:, -1:]
        logits = logits.masked_fill(
            logits < threshold,
            float("-inf"),
        )

    if 0.0 < top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(
            logits,
            descending=True,
        )

        sorted_probs = torch.softmax(sorted_logits, dim=-1)
        cumulative_probs = sorted_probs.cumsum(dim=-1)

        remove = cumulative_probs > top_p
        remove[:, 1:] = remove[:, :-1].clone()
        remove[:, 0] = False

        sorted_logits = sorted_logits.masked_fill(
            remove,
            float("-inf"),
        )

        filtered = torch.full_like(
            logits,
            float("-inf"),
        )

        logits = filtered.scatter(
            dim=-1,
            index=sorted_indices,
            src=sorted_logits,
        )

    probabilities = torch.softmax(logits, dim=-1)

    return torch.multinomial(
        probabilities,
        num_samples=1,
    )


def resolve_prompt(args: argparse.Namespace) -> str:
    if args.prompt is not None:
        return args.prompt

    if args.prompt_file is not None:
        return args.prompt_file.read_text(encoding="utf-8")

    raise SystemExit(
        "Provide --prompt, --prompt-file, or --interactive."
    )


def main() -> None:
    args = parse_args()
    model_dir = args.model_dir.resolve()

    os.environ["MULTISCREEN_BACKEND"] = args.backend
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

    sys.path.insert(0, str(model_dir / "runtime"))

    import torch
    from litgpt import Config, GPT, Tokenizer

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

    if args.backend == "triton" and device.type != "cuda":
        raise RuntimeError("The Triton backend requires CUDA")

    if args.dtype == "auto":
        dtype = (
            torch.bfloat16
            if device.type == "cuda"
            else torch.float32
        )
    else:
        dtype = {
            "float32": torch.float32,
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
        }[args.dtype]

    torch.manual_seed(args.seed)

    if device.type == "cuda":
        torch.cuda.manual_seed_all(args.seed)

    config = Config.from_file(
        model_dir / "model_config.yaml"
    )

    model = GPT(config)

    try:
        checkpoint = torch.load(
            model_dir / "lit_model.pth",
            map_location="cpu",
            weights_only=False,
        )
    except TypeError:
        checkpoint = torch.load(
            model_dir / "lit_model.pth",
            map_location="cpu",
        )

    state_dict = extract_state_dict(checkpoint, torch)

    incompatible = model.load_state_dict(
        state_dict,
        strict=False,
    )

    if incompatible.missing_keys or incompatible.unexpected_keys:
        raise RuntimeError(
            "Checkpoint mismatch:\n"
            f"Missing: {incompatible.missing_keys}\n"
            f"Unexpected: {incompatible.unexpected_keys}"
        )

    model = model.to(
        device=device,
        dtype=dtype,
    )
    model.eval()

    tokenizer = Tokenizer(model_dir)

    print(f"Device: {device}")
    print(f"Backend: {args.backend}")
    print(f"Precision: {dtype}")
    print(
        "Decoding: full sequence "
        "(incremental KV cache is not used)"
    )
    print()

    def complete(prompt: str) -> str:
        tokens = tokenizer.encode(
            prompt,
            device=device,
        ).reshape(1, -1).long()

        with torch.inference_mode():
            for _ in range(args.max_new_tokens):
                context = tokens[:, -config.block_size:]

                output = model(context)

                logits = (
                    output[0]
                    if isinstance(output, tuple)
                    else output
                )

                next_token = sample_token(
                    logits[:, -1, :].float(),
                    torch=torch,
                    temperature=args.temperature,
                    top_k=args.top_k,
                    top_p=args.top_p,
                )

                tokens = torch.cat(
                    (tokens, next_token),
                    dim=1,
                )

                if (
                    tokenizer.eos_id is not None
                    and next_token.item() == tokenizer.eos_id
                ):
                    break

        return tokenizer.decode(
            tokens[0].detach().cpu()
        )

    if args.interactive:
        print("Enter a prompt. Use Ctrl+D or /quit to exit.")

        while True:
            try:
                prompt = input("\nPrompt> ")
            except EOFError:
                print()
                break

            if prompt.strip().lower() in {"/quit", "/exit"}:
                break

            if not prompt.strip():
                continue

            print("\n" + complete(prompt))
    else:
        print(complete(resolve_prompt(args)))


if __name__ == "__main__":
    main()