#!/usr/bin/env python3
"""
infer_bytefast60m.py

Correct standalone inference utility for the custom FastDeepHybridLM defined
by bytefalcon_fast60m.py.

This does NOT instantiate Falcon-H1 or any Hugging Face AutoModel class.
It imports the exact training architecture and calls its load_model_bundle(),
which reconstructs Fast60MConfig + FastDeepHybridLM and strictly loads model.pt.

Expected checkpoint:
    step-XXXXXXXX/
        config.json
        model.pt
        tokenizer.json
        tokenizer_config.json
        ...

Rewrite training format:
    instruction

    "source text"

    "target output"<eos>

For inference, rewrite mode supplies the opening output quote and lets the
model generate the target text, closing quote, and EOS.
"""

from __future__ import annotations

import argparse
import contextlib
import importlib.util
import json
import os
import re
import sys
import time
from pathlib import Path
from types import ModuleType
from typing import Any


CONTEXT_LENGTH = 4096

# Match the training runtime setup before importing the architecture module.
os.environ.setdefault("USE_HUB_KERNELS", "NO")
os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("USE_ROCM_CK_GEMM", "1")
os.environ.pop("PYTORCH_HIP_ALLOC_CONF", None)


def is_checkpoint(path: Path) -> bool:
    return (
        path.is_dir()
        and (path / "config.json").is_file()
        and (path / "model.pt").is_file()
    )


def checkpoint_rank(path: Path) -> tuple[int, float, str]:
    matches = re.findall(r"\d+", path.name)
    step = int(matches[-1]) if matches else -1
    try:
        modified = path.stat().st_mtime
    except OSError:
        modified = 0.0
    return step, modified, path.name


def resolve_checkpoint(
    path: Path,
    *,
    extra_bases: list[Path] | None = None,
) -> Path:
    """
    Accept an exact checkpoint, a run directory, or its checkpoints directory.

    Relative paths are searched from:
      1. the current working directory;
      2. the inference script directory;
      3. any supplied extra bases, such as the architecture script directory.

    Preference within each candidate:
      exact directory -> final/ -> newest immediate checkpoint ->
      newest run/checkpoints checkpoint -> initial/
    """
    raw_path = path.expanduser()
    bases = [
        Path.cwd(),
        Path(__file__).resolve().parent,
    ]
    if extra_bases:
        bases.extend(base.expanduser().resolve() for base in extra_bases)

    candidate_roots: list[Path] = []
    if raw_path.is_absolute():
        candidate_roots.append(raw_path.resolve())
    else:
        candidate_roots.extend(
            (base / raw_path).resolve()
            for base in bases
        )

    candidate_roots = list(dict.fromkeys(candidate_roots))
    inspected: list[dict[str, Any]] = []

    def inspect_root(root: Path) -> Path | None:
        inspected.append(
            {
                "candidate_root": str(root),
                "exists": root.exists(),
                "is_directory": root.is_dir(),
            }
        )

        if is_checkpoint(root):
            return root

        final = root / "final"
        if is_checkpoint(final):
            return final

        direct = (
            sorted(
                (
                    child
                    for child in root.iterdir()
                    if child.is_dir() and is_checkpoint(child)
                ),
                key=checkpoint_rank,
            )
            if root.is_dir()
            else []
        )
        if direct:
            return direct[-1]

        checkpoint_root = root / "checkpoints"
        nested = (
            sorted(
                (
                    child
                    for child in checkpoint_root.iterdir()
                    if child.is_dir() and is_checkpoint(child)
                ),
                key=checkpoint_rank,
            )
            if checkpoint_root.is_dir()
            else []
        )
        if nested:
            return nested[-1]

        initial = root / "initial"
        if is_checkpoint(initial):
            print(
                "WARNING: using initial/; this is an untrained model.",
                file=sys.stderr,
            )
            return initial

        for directory in (root, checkpoint_root):
            if not directory.is_dir():
                continue
            for child in sorted(directory.iterdir()):
                if not child.is_dir():
                    continue
                inspected.append(
                    {
                        "path": str(child),
                        "has_config": (child / "config.json").is_file(),
                        "has_model_pt": (child / "model.pt").is_file(),
                        "files": sorted(
                            item.name
                            for item in child.iterdir()
                            if item.is_file()
                        )[:50],
                    }
                )

        return None

    for candidate_root in candidate_roots:
        resolved = inspect_root(candidate_root)
        if resolved is not None:
            print(
                f"Resolved model path from {candidate_root}",
                file=sys.stderr,
            )
            return resolved

    raise FileNotFoundError(
        "Could not find a FastDeepHybridLM checkpoint containing both "
        "config.json and model.pt.\n"
        "The supplied --model path was searched relative to the working "
        "directory, inference-script directory, and architecture-script "
        "directory.\n"
        + json.dumps(inspected, indent=2)
        + "\n\nCurrent working directory: "
        + str(Path.cwd())
        + "\nInference script directory: "
        + str(Path(__file__).resolve().parent)
    )


def find_architecture_script(explicit: Path | None) -> Path:
    if explicit is not None:
        path = explicit.expanduser().resolve()
        if not path.is_file():
            raise FileNotFoundError(
                f"Architecture script does not exist: {path}"
            )
        return path

    here = Path(__file__).resolve().parent
    cwd = Path.cwd()
    candidates = [
        cwd / "bytefalcon_fast60m.py",
        cwd / "bytefalcon.py",
        here / "bytefalcon_fast60m.py",
        here / "bytefalcon.py",
    ]

    for candidate in candidates:
        if not candidate.is_file():
            continue
        source = candidate.read_text(
            encoding="utf-8",
            errors="replace",
        )
        required = (
            "class Fast60MConfig",
            "def create_model_classes",
            "def load_model_bundle",
        )
        if all(marker in source for marker in required):
            return candidate.resolve()

    raise FileNotFoundError(
        "Could not locate the custom architecture script. Pass it explicitly:\n"
        "  --architecture-script /path/to/bytefalcon_fast60m.py"
    )


def load_architecture_module(path: Path) -> ModuleType:
    module_name = "_bytefast60m_architecture"
    specification = importlib.util.spec_from_file_location(
        module_name,
        path,
    )
    if specification is None or specification.loader is None:
        raise RuntimeError(
            f"Could not create an import specification for {path}"
        )

    module = importlib.util.module_from_spec(specification)
    # Dataclasses and some runtime machinery expect the module to be present.
    sys.modules[module_name] = module
    specification.loader.exec_module(module)

    required = (
        "Fast60MConfig",
        "create_model_classes",
        "import_training_stack",
        "load_model_bundle",
        "load_tokenizer",
    )
    missing = [
        name for name in required if not hasattr(module, name)
    ]
    if missing:
        raise RuntimeError(
            f"{path} is not the FastDeepHybridLM training script; "
            f"missing definitions: {missing}"
        )

    return module


def resolve_tokenizer(
    checkpoint: Path,
    explicit: Path | None,
) -> Path:
    candidates: list[Path] = []

    if explicit is not None:
        candidates.append(explicit.expanduser().resolve())

    candidates.extend(
        [
            checkpoint,
            checkpoint / "tokenizer",
        ]
    )

    project = Path(__file__).resolve().parent
    candidates.extend(
        [
            project / "artifacts" / "byte-tokenizer",
            Path.cwd() / "artifacts" / "byte-tokenizer",
        ]
    )

    for parent in list(checkpoint.parents)[:5]:
        candidates.extend(
            [
                parent / "artifacts" / "byte-tokenizer",
                parent / "byte-tokenizer",
            ]
        )

    candidates = list(dict.fromkeys(candidates))
    for candidate in candidates:
        if (
            candidate.is_dir()
            and (
                (candidate / "tokenizer.json").is_file()
                or (candidate / "tokenizer.model").is_file()
            )
        ):
            return candidate

    raise FileNotFoundError(
        "Tokenizer not found. Pass --tokenizer explicitly. Checked:\n"
        + "\n".join(f"  - {path}" for path in candidates)
    )


def rewrite_prompt(instruction: str, source_text: str) -> str:
    instruction = instruction.strip()
    source = '"' + source_text + '"'
    if instruction:
        return instruction + "\n\n" + source + '\n\n"'
    return source + '\n\n"'


def control_token_id_map(
    tokenizer: Any,
    architecture: ModuleType,
) -> dict[str, int]:
    control_tokens = getattr(
        architecture,
        "CONTROL_TOKENS",
        [
            "<pad>",
            "<bos>",
            "<eos>",
            "<unk>",
            "<instruction>",
            "<text>",
            "<output>",
            "<record>",
            "<byte_start>",
            "<byte_end>",
        ],
    )
    result = {}
    for token in control_tokens:
        token_id = tokenizer.convert_tokens_to_ids(token)
        if token_id is None:
            continue
        token_id = int(token_id)
        if token_id >= 0:
            result[token] = token_id
    return result


def blocked_generation_ids(
    tokenizer: Any,
    architecture: ModuleType,
) -> list[int]:
    mapping = control_token_id_map(tokenizer, architecture)
    return sorted(
        token_id
        for token, token_id in mapping.items()
        if token != "<eos>"
    )


def apply_repetition_penalty(
    torch: Any,
    logits: Any,
    input_ids: Any,
    penalty: float,
) -> Any:
    if penalty == 1.0:
        return logits

    used = torch.unique(input_ids)
    selected = logits[:, used]
    logits[:, used] = torch.where(
        selected < 0,
        selected * penalty,
        selected / penalty,
    )
    return logits


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

    logits = logits / max(temperature, 1e-5)

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

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

    if top_p < 1.0:
        sorted_probabilities, sorted_indices = torch.sort(
            probabilities,
            descending=True,
            dim=-1,
        )
        cumulative = sorted_probabilities.cumsum(dim=-1)
        remove = cumulative > top_p
        remove[:, 1:] = remove[:, :-1].clone()
        remove[:, 0] = False
        sorted_probabilities = sorted_probabilities.masked_fill(
            remove,
            0.0,
        )
        denominator = sorted_probabilities.sum(
            dim=-1,
            keepdim=True,
        ).clamp_min(1e-12)
        sorted_probabilities = (
            sorted_probabilities / denominator
        )
        sampled = torch.multinomial(
            sorted_probabilities,
            num_samples=1,
        )
        return sorted_indices.gather(-1, sampled)

    return torch.multinomial(probabilities, num_samples=1)


def clean_completion(text: str, rewrite_mode: bool) -> str:
    for marker in ("<eos>", "<record>", "<pad>"):
        position = text.find(marker)
        if position >= 0:
            text = text[:position]

    if rewrite_mode:
        text = text.rstrip()
        if text.endswith('"'):
            text = text[:-1]

    return text


def generate(
    *,
    architecture: ModuleType,
    checkpoint: Path,
    tokenizer_path: Path,
    prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_k: int,
    top_p: float,
    repetition_penalty: float,
    seed: int,
    compile_model: bool,
    compile_mode: str,
    stream: bool,
    rewrite_mode: bool,
    allow_control_tokens: bool,
    show_top_tokens: int,
) -> tuple[str, dict[str, Any]]:
    (
        _np,
        torch,
        nn,
        F,
        _DataLoader,
        _Dataset,
    ) = architecture.import_training_stack()

    if not torch.cuda.is_available():
        raise RuntimeError(
            "ROCm PyTorch did not expose the AMD GPU through torch.cuda."
        )

    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    device = torch.device("cuda")
    tokenizer = architecture.load_tokenizer(tokenizer_path)
    model = architecture.load_model_bundle(
        checkpoint,
        torch,
        nn,
        F,
    )
    model.to(device)
    model.eval()

    blocked_ids = (
        []
        if allow_control_tokens
        else blocked_generation_ids(tokenizer, architecture)
    )
    blocked_tensor = (
        torch.tensor(
            blocked_ids,
            device=device,
            dtype=torch.long,
        )
        if blocked_ids
        else None
    )

    active_model = model
    if compile_model:
        active_model = torch.compile(
            model,
            mode=compile_mode,
            fullgraph=False,
            dynamic=False,
        )

    encoded = tokenizer(
        prompt,
        add_special_tokens=False,
        return_tensors="pt",
        return_token_type_ids=False,
    )
    input_ids = encoded.input_ids.to(device)

    prompt_tokens = int(input_ids.shape[1])
    maximum_context = int(
        model.config.max_position_embeddings
    )
    if prompt_tokens >= maximum_context:
        raise ValueError(
            f"Prompt has {prompt_tokens} tokens and exceeds the "
            f"{maximum_context}-token context."
        )

    max_new_tokens = min(
        max_new_tokens,
        maximum_context - prompt_tokens,
    )

    eos_id = int(tokenizer.eos_token_id)
    generated_ids: list[int] = []

    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    started = time.perf_counter()

    # This architecture has no KV/conv recurrent inference cache. It therefore
    # recomputes the active prefix each step, matching the original CLI.
    with torch.inference_mode():
        for _ in range(max_new_tokens):
            model_input = input_ids[
                :, -maximum_context:
            ]

            with torch.autocast(
                device_type="cuda",
                dtype=torch.bfloat16,
                enabled=True,
            ):
                logits = active_model(
                    input_ids=model_input,
                    return_last_logits=True,
                ).logits[:, -1, :]

            if not bool(torch.isfinite(logits).all().item()):
                print(
                    "Non-finite BF16 logits; retrying this token in FP32.",
                    file=sys.stderr,
                )
                with torch.autocast(
                    device_type="cuda",
                    enabled=False,
                ):
                    logits = model(
                        input_ids=model_input,
                        return_last_logits=True,
                    ).logits[:, -1, :].float()

            if not bool(torch.isfinite(logits).all().item()):
                logits = torch.nan_to_num(
                    logits,
                    nan=-float("inf"),
                    posinf=1e4,
                    neginf=-1e4,
                )

            if show_top_tokens > 0:
                top_values, top_indices = torch.topk(
                    logits,
                    min(show_top_tokens, logits.shape[-1]),
                    dim=-1,
                )
                report = [
                    {
                        "id": int(token_id),
                        "token": tokenizer.decode(
                            [int(token_id)],
                            skip_special_tokens=False,
                            clean_up_tokenization_spaces=False,
                        ),
                        "logit": float(value),
                    }
                    for token_id, value in zip(
                        top_indices[0].tolist(),
                        top_values[0].float().tolist(),
                    )
                ]
                print(
                    "raw top tokens: "
                    + json.dumps(report, ensure_ascii=False),
                    file=sys.stderr,
                )

            if blocked_tensor is not None:
                logits.index_fill_(
                    1,
                    blocked_tensor,
                    -float("inf"),
                )

            logits = apply_repetition_penalty(
                torch,
                logits,
                model_input,
                repetition_penalty,
            )

            if not bool(torch.isfinite(logits).any().item()):
                next_token = torch.tensor(
                    [[int(tokenizer.eos_token_id)]],
                    device=device,
                    dtype=torch.long,
                )
            else:
                next_token = sample_next_token(
                    torch,
                    logits,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                )

            token_id = int(next_token.item())
            if token_id in blocked_ids:
                raise RuntimeError(
                    "A reserved control token escaped masking: "
                    f"id={token_id}, token={tokenizer.decode([token_id], skip_special_tokens=False)!r}"
                )
            generated_ids.append(token_id)
            input_ids = torch.cat(
                (input_ids, next_token),
                dim=-1,
            )

            if stream:
                piece = tokenizer.decode(
                    [token_id],
                    skip_special_tokens=False,
                    clean_up_tokenization_spaces=False,
                )
                print(piece, end="", flush=True)

            if token_id == eos_id:
                break

    torch.cuda.synchronize()
    elapsed = time.perf_counter() - started

    raw_completion = tokenizer.decode(
        generated_ids,
        skip_special_tokens=False,
        clean_up_tokenization_spaces=False,
    )
    completion = clean_completion(
        raw_completion,
        rewrite_mode,
    )

    if stream:
        print()

    metrics = {
        "architecture": model.config.architecture,
        "model_type": model.config.model_type,
        "checkpoint": str(checkpoint),
        "tokenizer": str(tokenizer_path),
        "parameters": sum(
            parameter.numel()
            for parameter in model.parameters()
        ),
        "prompt_tokens": prompt_tokens,
        "generated_tokens": len(generated_ids),
        "elapsed_seconds": elapsed,
        "tokens_per_second": (
            len(generated_ids) / elapsed
            if elapsed > 0
            else None
        ),
        "peak_vram_gib": (
            torch.cuda.max_memory_allocated() / (1024**3)
        ),
        "compiled": compile_model,
        "blocked_control_token_ids": blocked_ids,
        "note": (
            "Generation recomputes the active prefix because this custom "
            "architecture does not implement an incremental inference cache."
        ),
    }

    return completion, metrics


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Inference for the custom byte-deep-hybrid FastDeepHybridLM."
        ),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--model",
        type=Path,
        default=Path("runs/bytefast-60m"),
        help=(
            "Exact checkpoint, run directory, or checkpoints directory. "
            "Relative paths are searched from the shell, script, and "
            "architecture-script directories."
        ),
    )
    parser.add_argument(
        "--architecture-script",
        type=Path,
        help=(
            "Path to bytefalcon_fast60m.py. Automatically discovered "
            "when omitted."
        ),
    )
    parser.add_argument(
        "--tokenizer",
        type=Path,
        help=(
            "Tokenizer directory. The checkpoint tokenizer is preferred."
        ),
    )

    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument(
        "--prompt",
        help="Raw language-model prompt.",
    )
    input_group.add_argument(
        "--text",
        help="Source text for rewrite mode.",
    )
    parser.add_argument(
        "--instruction",
        default="Rewrite this clearly and naturally.",
        help="Instruction used with --text.",
    )

    parser.add_argument("--max-new-tokens", type=int, default=128)
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--top-p", type=float, default=0.95)
    parser.add_argument("--top-k", type=int, default=50)
    parser.add_argument(
        "--repetition-penalty",
        type=float,
        default=1.1,
    )
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--stream",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    parser.add_argument(
        "--metrics",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    parser.add_argument("--compile", action="store_true")
    parser.add_argument(
        "--compile-mode",
        choices=[
            "default",
            "reduce-overhead",
            "max-autotune",
        ],
        default="reduce-overhead",
    )
    parser.add_argument(
        "--show-prompt",
        action="store_true",
    )
    parser.add_argument(
        "--allow-control-tokens",
        action="store_true",
        help="Allow structural tokens such as <pad>; disabled by default.",
    )
    parser.add_argument(
        "--show-top-tokens",
        type=int,
        default=0,
        help="Print the raw top-N logits before control-token masking.",
    )

    return parser


def validate_args(args: argparse.Namespace) -> None:
    if args.max_new_tokens <= 0:
        raise ValueError("--max-new-tokens must be positive.")
    if args.temperature < 0:
        raise ValueError("--temperature cannot be negative.")
    if not 0 < args.top_p <= 1:
        raise ValueError("--top-p must be in (0, 1].")
    if args.top_k < 0:
        raise ValueError("--top-k cannot be negative.")
    if args.repetition_penalty <= 0:
        raise ValueError(
            "--repetition-penalty must be positive."
        )
    if args.show_top_tokens < 0:
        raise ValueError("--show-top-tokens must be non-negative.")


def main() -> int:
    args = build_parser().parse_args()
    validate_args(args)

    architecture_path = find_architecture_script(
        args.architecture_script
    )
    checkpoint = resolve_checkpoint(
        args.model,
        extra_bases=[architecture_path.parent],
    )
    architecture = load_architecture_module(
        architecture_path
    )
    tokenizer_path = resolve_tokenizer(
        checkpoint,
        args.tokenizer,
    )

    rewrite_mode = args.text is not None
    prompt = (
        rewrite_prompt(args.instruction, args.text)
        if rewrite_mode
        else args.prompt
    )
    assert prompt is not None

    print(
        json.dumps(
            {
                "checkpoint": str(checkpoint),
                "architecture_script": str(architecture_path),
                "tokenizer": str(tokenizer_path),
                "rewrite_mode": rewrite_mode,
            },
            indent=2,
        ),
        file=sys.stderr,
    )

    if args.show_prompt:
        print(
            "----- PROMPT -----\n"
            + prompt
            + "\n----- END PROMPT -----",
            file=sys.stderr,
        )

    completion, metrics = generate(
        architecture=architecture,
        checkpoint=checkpoint,
        tokenizer_path=tokenizer_path,
        prompt=prompt,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
        repetition_penalty=args.repetition_penalty,
        seed=args.seed,
        compile_model=args.compile,
        compile_mode=args.compile_mode,
        stream=args.stream,
        rewrite_mode=rewrite_mode,
        allow_control_tokens=args.allow_control_tokens,
        show_top_tokens=args.show_top_tokens,
    )

    if not args.stream:
        print(completion)

    if args.metrics:
        print(
            "\n" + json.dumps(metrics, indent=2),
            file=sys.stderr,
        )

    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        print("\nInterrupted.", file=sys.stderr)
        raise SystemExit(130)