"""Standalone SHARED inference after a bounded visible-response Fable SFT stage."""

import argparse
import hashlib
import importlib.util
import json
import os
import sys
from pathlib import Path

import torch

FORMAT = "native-loom-fable-response-sft-v1"
PAIR_HELPER_SHA = "62f3432b3174d6825b39aca1ac9fbab38f924f8b5e6c04f7e45c17e1a8c69e6e"
PARENT_SHA = "809a3836e2ad6a62c0bc95ae95431b4a18ba50bc839d25998a03bfa62eadd915"
INITIAL_SHA = "18c3eac60d7b51d7211cd943eb646bf1cb50d0dacfb587bd182cd9fd9d88b935"
DATA_SHA = "8dc518ae52704a7b24d6f2d5b0e43fa9f6e4374b6a9284a4a5db9f377a07936c"
PROMPT_TEMPLATE = "### Context\n{prompt}\n\n### Response\n"


def require(condition, message):
    if not condition:
        raise ValueError(message)


def digest(path):
    with Path(path).open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def pair_helper(root):
    """Import only the exact reviewed, bundled pure helper before model construction."""
    path = root / "runtime/loom_melt_pair_portable.py"
    require(
        path.is_file()
        and not path.is_symlink()
        and not path.parent.is_symlink()
        and path.resolve().is_relative_to(root),
        "Reviewed portable helper is missing or changed",
    )
    source = path.read_bytes()
    require(
        hashlib.sha256(source).hexdigest() == PAIR_HELPER_SHA,
        "Reviewed portable helper is missing or changed",
    )
    spec = importlib.util.spec_from_file_location("_loom_fable_pair_reference", path)
    module = importlib.util.module_from_spec(spec)
    previous = sys.dont_write_bytecode
    try:
        sys.dont_write_bytecode = True
        # A valid stale .pyc must never execute before inventory rejection.
        exec(compile(source, str(path), "exec"), module.__dict__)
    finally:
        sys.dont_write_bytecode = previous
    return module


def verify_bundle(directory):
    root = Path(directory).resolve()
    helper = pair_helper(root)
    manifest = helper.read(helper.local_file(root, "manifest.json"))
    require(
        manifest["format"] == FORMAT and manifest["variant"] == "shared",
        "Expected the single SHARED Fable response-SFT export",
    )
    actual = set()
    for path in root.rglob("*"):
        require(
            not path.is_symlink()
            and "__pycache__" not in path.relative_to(root).parts
            and path.suffix not in {".pyc", ".pyo"},
            "Unbound executable cache or symlink",
        )
        if path.is_file():
            actual.add(str(path.relative_to(root)))
    require(actual == set(manifest["files"]) | {"manifest.json"}, "Bundle inventory differs")
    for name, binding in manifest["files"].items():
        require(
            helper.file_row(helper.local_file(root, name)) == binding,
            "Bundled file changed: " + name,
        )
    return root, manifest, helper


def sft_contract(manifest):
    sft = manifest["sft"]
    require(
        manifest["parent_manifest_sha256"] == PARENT_SHA
        and manifest["initial_shared_state_sha256"] == INITIAL_SHA
        and manifest["data_receipt_sha256"] == DATA_SHA
        and manifest["adapted_updates"] == 32
        and manifest["adapted_targets_per_arm"] == 4087238
        and manifest["prompt_template"] == PROMPT_TEMPLATE
        and sft["target_kind"] == "visible_explanatory_response"
        and sft["private_cot_claim"] is False
        and type(sft["epochs"]) is int
        and sft["epochs"] in {1, 2, 3}
        and sft["updates"] == 7 * sft["epochs"]
        and sft["input_tokens"] == 54651 * sft["epochs"]
        and sft["supervised_targets"] == 49199 * sft["epochs"]
        and sft["optimizer_restart"] is True,
        "Saved parent, data, formatting or complete SFT epoch contract differs",
    )


def load_bundle(directory, device="cpu"):
    require(device in {"cpu", "cuda"}, "Unknown inference device")
    require(torch.__version__.split("+")[0] == "2.13.0", "Use pinned Torch 2.13.0")
    require(
        all(
            os.environ.get(name) in {None, "0"}
            for name in ("NVIDIA_TF32_OVERRIDE", "TORCH_ALLOW_TF32_CUBLAS_OVERRIDE")
        ),
        "TF32 overrides conflict with FP32 routing",
    )
    root, manifest, helper = verify_bundle(directory)
    sft_contract(manifest)
    config = helper.read(helper.local_file(root, "model_config.json"))
    # These inherited fields describe the unchanged scratch/MELT parent geometry.
    helper.native_contract(manifest, config)
    require(
        digest(helper.local_file(root, "source-integrity.json")) == helper.INTEGRITY_SHA,
        "Frozen native source receipt changed",
    )
    reviewed = helper.read(root / "source-integrity.json")["child_hashes"]
    require(
        {
            name[7:]: binding["sha256"]
            for name, binding in manifest["files"].items()
            if name.startswith("source/")
        }
        == reviewed,
        "Frozen source inventory differs",
    )
    adapter, exporter = helper.runtime_modules(root)
    contract = sys.modules["research.probes.loom_melt_v2_contract"]
    require(
        contract.matches_compute_contract(manifest["compute_contract"]),
        "Saved BF16 training compute contract changed",
    )
    model = helper.construct_model(root, config, manifest["files"])
    require(
        sum(p.numel() for p in model.parameters()) == 112237056
        and sum(".melt_gate." in name for name, _ in model.named_parameters()) == 18,
        "Native parameter or gate inventory differs",
    )
    helper.load_state(model, config, helper.local_file(root, "model.pt"), exporter)
    packed = adapter.pack_experts(model)
    require(
        packed["roundtrip_bit_equal"]
        and packed["parameters"] == 112237056
        and packed["parameter_tensors"] == 92
        and helper.state_digest(model) == manifest["packed_state_sha256"],
        "Saved SFT state differs from its validated checkpoint",
    )
    if device == "cuda":
        require(
            torch.version.cuda == "13.0" and torch.cuda.is_bf16_supported(),
            "CUDA inference requires pinned CUDA 13.0 and BF16",
        )
    model.to(device).eval()
    from tokenizers import Tokenizer

    tokenizer_path = helper.local_file(root, "tokenizer/tokenizer.json")
    require(
        digest(tokenizer_path)
        == "9ca9acddb6525a194ec8ac7a87f24fbba7232a9a15ffa1af0c1224fcd888e47c",
        "Unchanged native tokenizer required",
    )
    tokenizer = Tokenizer.from_file(str(tokenizer_path))
    engine = adapter.CudaMeltEngine(
        model,
        "shared",
        50,
        50,
        precision="fp32" if device == "cpu" else "bf16",
        cpu_reference=device == "cpu",
    )
    return engine, tokenizer, config, helper


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--bundle", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--max-new-tokens", type=int, default=64)
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    args = parser.parse_args()
    torch.set_num_threads(1)
    engine, tokenizer, config, helper = load_bundle(args.bundle, args.device)
    ids = tokenizer.encode(PROMPT_TEMPLATE.format(prompt=args.prompt), add_special_tokens=False).ids
    generated = helper.greedy(
        engine,
        ids,
        args.max_new_tokens,
        config["model"]["max_seq_len"],
        config["tokenizer"]["eos_token_id"],
    )
    print(
        json.dumps(
            {
                "text": tokenizer.decode(generated, skip_special_tokens=True),
                "new_tokens": len(generated),
                "variant": "shared",
                "target_kind": "visible_explanatory_response",
            }
        )
    )


if __name__ == "__main__":
    main()