"""Verify small, immutable official source files without weights or GPU imports.

No model code is executed. No credentials are read. Files stay in a temporary
folder unless --source-dir supplies an existing reviewed source snapshot.
"""
from __future__ import annotations

import argparse
import ast
import hashlib
import json
import ssl
import tempfile
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
LOCK = json.loads((ROOT / "model-lock.json").read_text())


def class_node(tree, name):
    return next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == name)


def method(node, name):
    return next(item for item in node.body if isinstance(item, ast.FunctionDef) and item.name == name)


def attribute_calls(node):
    return [item.func.attr for item in ast.walk(node) if isinstance(item, ast.Call) and isinstance(item.func, ast.Attribute)]


def verify(folder: Path) -> dict:
    for filename, digest in LOCK["reviewed_source_sha256"].items():
        actual = hashlib.sha256((folder / filename).read_bytes()).hexdigest()
        if actual != digest:
            raise ValueError("Immutable source fingerprint mismatch: " + filename)
    tree = ast.parse((folder / "modeling_minicpmo.py").read_text())
    model = class_node(tree, "MiniCPMO")
    duplex = class_node(tree, "MiniCPMODuplex")
    assert "from_existing_model" in attribute_calls(method(model, "as_duplex"))
    defaults = next(ast.literal_eval(node.value) for node in duplex.body if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "_default_duplex_params" for target in node.targets))
    assert defaults["sample_rate"] == 16000
    assert defaults["chunk_ms"] == 1000 and defaults["first_chunk_ms"] == 1035
    assert defaults["text_repetition_window_size"] == 512
    assert defaults["generate_audio"] is True  # Adapter explicitly opts out.
    assert defaults["force_listen_count"] == 0
    prefill = method(duplex, "streaming_prefill")
    assert {"audio_waveform", "frame_list", "max_slice_nums", "batch_vision_feed"} <= {arg.arg for arg in prefill.args.args}
    assert "torch.no_grad()" in [ast.unparse(d) for d in prefill.decorator_list]
    assert "process_audio_streaming" in attribute_calls(prefill)
    assert "get_streaming_chunk_size" in attribute_calls(prefill)
    generate = method(duplex, "streaming_generate")
    assert "torch.no_grad()" in [ast.unparse(d) for d in generate.decorator_list]
    generate_args = {arg.arg for arg in generate.args.args}
    assert {"max_new_speak_tokens_per_chunk", "listen_top_k", "listen_prob_scale"} <= generate_args
    assert "force_listen_override" not in generate_args  # Demo API differs.
    generate_source = ast.unparse(generate)
    # Verify the native listen/speak handoff and exclusion of its control token from
    # public words. No invented force_speak argument or foreign control token.
    assert "logits = self.pending_logits" in generate_source
    assert "self.pending_logits = None" in generate_source
    assert "last_id.item() in self.chunk_speak_token_ids" in generate_source
    assert "if j != 0:" in generate_source
    calls = attribute_calls(generate)
    assert "register_unit_end" in calls and "enforce_window" in calls
    assert "finalize_unit" not in calls
    constructor = method(duplex, "from_existing_model")
    assert "convert_tokens_to_ids('<|speak|>')" in ast.unparse(constructor)
    assert "instance.chunk_speak_token_ids = [instance.speak_token_id]" in ast.unparse(constructor)
    assert attribute_calls(constructor).count("init_tts") == 1
    # This exact unconditional call is bypassed on the model instance during
    # text-only construction. The official files themselves remain unmodified.
    assert any(isinstance(node, ast.Expr) and isinstance(node.value, ast.Call)
               and isinstance(node.value.func, ast.Attribute) and node.value.func.attr == "init_tts"
               for node in constructor.body)
    assert "set" in attribute_calls(method(duplex, "set_session_stop"))
    assert "register_system_prompt" in attribute_calls(method(duplex, "prepare"))
    init_streaming = method(model, "init_streaming_processor")
    assert any(isinstance(node, ast.keyword) and node.arg == "enable_sliding_window"
               and isinstance(node.value, ast.Constant) and node.value.value is True
               for node in ast.walk(init_streaming))
    # Explicit text-only branch returns text with audio_waveform=None.
    branches = [node for node in ast.walk(generate) if isinstance(node, ast.If)
                and isinstance(node.test, ast.UnaryOp) and isinstance(node.test.op, ast.Not)
                and isinstance(node.test.operand, ast.Attribute) and node.test.operand.attr == "generate_audio"]
    assert branches and any(isinstance(node, ast.Return) for node in ast.walk(branches[0]))
    processor = ast.parse((folder / "processing_minicpmo.py").read_text())
    image_processor = class_node(processor, "MiniCPMVImageProcessor")
    grid_source = ast.unparse(method(image_processor, "get_sliced_grid"))
    assert "multiple = min(math.ceil(ratio), max_slice_nums)" in grid_source
    assert "if multiple <= 1 or nerver_split:" in grid_source
    assert "original_width * original_height / (self.scale_resolution * self.scale_resolution)" in grid_source
    mel = class_node(processor, "StreamingMelProcessorExact")
    assert "get_chunk_size" in [n.name for n in mel.body if isinstance(n, ast.FunctionDef)]
    # This independently recomputes the pinned first-core alignment.
    raw = defaults["first_chunk_ms"] * defaults["sample_rate"] // 1000
    first_core = raw // 160 * 160
    assert raw == 16560 and first_core == 16480 and raw - first_core == 80
    utils = ast.parse((folder / "utils.py").read_text())
    decoder = class_node(utils, "StreamDecoder")
    assert "self._unit_history = []" in ast.unparse(method(decoder, "reset"))
    assert "self.cache" in ast.unparse(method(decoder, "get_cache_length"))
    assert "token_ids" in {arg.arg for arg in method(decoder, "embed_tokens").args.args}
    assert "past_key_values=self.cache" in ast.unparse(method(decoder, "feed"))
    decode = method(decoder, "decode")
    rank_preferences = [node for node in ast.walk(decode) if isinstance(node, ast.If)
                        and "listen_rank < listen_top_k" in ast.unparse(node.test)]
    assert rank_preferences and any(isinstance(node, ast.Return) for node in ast.walk(rank_preferences[0]))
    assert "logits[0, self.listen_id] *= listen_prob_scale" in ast.unparse(decode)
    assert any(isinstance(node, ast.Subscript) and isinstance(node.value, ast.Attribute)
               and node.value.attr == "generated_tokens" and isinstance(node.slice, ast.Slice)
               and isinstance(node.slice.lower, ast.UnaryOp)
               and isinstance(node.slice.lower.operand, ast.Name)
               and node.slice.lower.operand.id == "text_repetition_window_size"
               for node in ast.walk(decode))
    config = json.loads((folder / "config.json").read_text())
    assert config["version"] == "4.5" and config["max_position_embeddings"] == 40960
    return {"status": "passed", "model_id": LOCK["model_id"], "revision": LOCK["model_revision"],
            "fingerprinted_files": len(LOCK["reviewed_source_sha256"]),
            "native_pcm_rate": 16000, "unit_samples": 16000, "first_core_samples": first_core,
            "stock_initial_carry_samples": 80, "text_only_branch_verified": True,
            "unconditional_speech_renderer_init_verified": True,
            "native_rank_listening_preference_verified": True, "mel_sliding_window_verified": True,
            "native_pending_logits_speak_decision_verified": True,
            "slicing_and_batching_arguments_verified": True, "native_no_grad_verified": True,
            "size_dependent_slicing_verified": True,
            "native_source_executed": False, "cuda_verified": False}


def verify_references(folder: Path) -> int:
    for name, reference in LOCK.get("reviewed_reference_sources", {}).items():
        if hashlib.sha256((folder / name).read_bytes()).hexdigest() != reference["sha256"]:
            raise ValueError("Official reference fingerprint mismatch: " + name)
    # The official Demo has a different finalize contract. Do not transplant it.
    demo_model = (folder / "demo-model.py").read_text()
    assert "force_listen_override" in demo_model and "def finalize_unit(" in demo_model
    assert "300 seconds" in (folder / "demo-video.md").read_text()
    return len(LOCK.get("reviewed_reference_sources", {}))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-dir", type=Path, help="Existing pinned source directory; no network")
    parser.add_argument("--reference-dir", type=Path, help="Existing pinned official reference files; no network")
    parser.add_argument("--report", type=Path)
    args = parser.parse_args()
    if args.source_dir:
        report = verify(args.source_dir)
        if args.reference_dir:
            report["fingerprinted_official_reference_files"] = verify_references(args.reference_dir)
    else:
        try:
            import certifi
            context = ssl.create_default_context(cafile=certifi.where())
        except ImportError:
            context = ssl.create_default_context()
        with tempfile.TemporaryDirectory(prefix="omnai-source-verification-") as temporary:
            folder = Path(temporary)
            for filename in LOCK["reviewed_source_sha256"]:
                # Host, repository and revision are fixed; no user URL fetcher.
                url = f"https://huggingface.co/openbmb/MiniCPM-o-4_5/resolve/{LOCK['model_revision']}/{filename}"
                with urllib.request.urlopen(url, context=context, timeout=30) as response:
                    data = response.read(2_000_000)
                (folder / filename).write_bytes(data)
            report = verify(folder)
            references = folder / "references"
            references.mkdir()
            for filename, reference in LOCK.get("reviewed_reference_sources", {}).items():
                with urllib.request.urlopen(reference["url"], context=context, timeout=30) as response:
                    data = response.read(2_000_000)
                (references / filename).write_bytes(data)
            report["fingerprinted_official_reference_files"] = verify_references(references)
    rendered = json.dumps(report, indent=2)
    if args.report:
        args.report.write_text(rendered + "\n")
    print(rendered)


if __name__ == "__main__":
    main()