#!/usr/bin/env python3
"""Verify a Q-Prefer adapter without loading the 4B base model."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path

from safetensors import safe_open

from qprefer_reward.constants import (
    BASE_MODEL_ID,
    BASE_MODEL_REVISION,
    EXPECTED_SPECIAL_TOKEN_IDS,
    PUBLISHED_ADAPTER_SHA256,
    PUBLISHED_ADAPTER_TENSOR_COUNT,
    PUBLISHED_RM_HEAD_SHAPE,
    PUBLISHED_SPECIAL_EMBEDDINGS_SHA256,
    SPECIAL_TOKENS,
)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("adapter", type=Path)
    parser.add_argument("--allow-unknown-checkpoint", action="store_true")
    return parser.parse_args()


def verify_manifest(root: Path) -> None:
    manifest_path = root / "artifact_manifest.json"
    if not manifest_path.is_file():
        raise FileNotFoundError(f"required artifact manifest is missing: {manifest_path}")

    manifest = json.loads(manifest_path.read_text())
    expected_metadata = {
        "format_version": 1,
        "base_model": BASE_MODEL_ID,
        "base_revision": BASE_MODEL_REVISION,
        "special_tokens": list(SPECIAL_TOKENS),
        "special_token_ids": list(EXPECTED_SPECIAL_TOKEN_IDS),
        "supported_dimensions": ["visual_quality", "text_alignment"],
        "unsupported_dimensions": ["motion_quality"],
    }
    mismatches = {
        key: (expected, manifest.get(key))
        for key, expected in expected_metadata.items()
        if manifest.get(key) != expected
    }
    if mismatches:
        raise RuntimeError(f"unexpected artifact manifest metadata: {mismatches}")

    files = manifest.get("files")
    if not isinstance(files, dict) or not files:
        raise RuntimeError("artifact manifest has no file checksums")
    required = {
        "adapter_config.json",
        "adapter_model.safetensors",
        "chat_template.jinja",
        "README.md",
        "special_token_embeddings.safetensors",
        "tokenizer.json",
        "tokenizer_config.json",
    }
    missing = required.difference(files)
    if missing:
        raise RuntimeError(f"artifact manifest is missing required files: {sorted(missing)}")

    for filename, record in files.items():
        if Path(filename).name != filename or filename == manifest_path.name:
            raise RuntimeError(f"invalid artifact filename in manifest: {filename!r}")
        if not isinstance(record, dict):
            raise RuntimeError(f"invalid manifest record for {filename!r}")
        path = root / filename
        if not path.is_file():
            raise FileNotFoundError(f"manifest file is missing: {path}")
        observed_bytes = path.stat().st_size
        observed_sha = sha256(path)
        if record.get("bytes") != observed_bytes or record.get("sha256") != observed_sha:
            raise RuntimeError(
                f"artifact integrity check failed for {filename}: "
                f"expected bytes/sha256={record.get('bytes')}/{record.get('sha256')}, "
                f"observed={observed_bytes}/{observed_sha}"
            )


def main() -> None:
    args = parse_args()
    root = args.adapter.expanduser().resolve()
    weights = root / "adapter_model.safetensors"
    config_path = root / "adapter_config.json"
    if not weights.is_file() or not config_path.is_file():
        raise FileNotFoundError("adapter_model.safetensors or adapter_config.json is missing")

    verify_manifest(root)

    observed_sha = sha256(weights)
    if observed_sha != PUBLISHED_ADAPTER_SHA256 and not args.allow_unknown_checkpoint:
        raise RuntimeError(
            "adapter checksum mismatch: "
            f"expected={PUBLISHED_ADAPTER_SHA256}, observed={observed_sha}"
        )

    config = json.loads(config_path.read_text())
    expected_config = {
        "peft_type": "LORA",
        "r": 64,
        "lora_alpha": 128,
        "lora_dropout": 0.05,
    }
    mismatches = {
        key: (expected, config.get(key))
        for key, expected in expected_config.items()
        if config.get(key) != expected
    }
    if mismatches:
        raise RuntimeError(f"unexpected adapter configuration: {mismatches}")
    if "rm_head" not in config.get("modules_to_save", []):
        raise RuntimeError("adapter config does not preserve rm_head")

    with safe_open(weights, framework="pt", device="cpu") as handle:
        tensor_keys = list(handle.keys())
        reward_head_keys = [key for key in tensor_keys if key.endswith("rm_head.weight")]
        if len(tensor_keys) != PUBLISHED_ADAPTER_TENSOR_COUNT:
            raise RuntimeError(
                "unexpected adapter tensor count: "
                f"expected={PUBLISHED_ADAPTER_TENSOR_COUNT}, observed={len(tensor_keys)}"
            )
        if len(reward_head_keys) != 1:
            raise RuntimeError(f"expected one rm_head.weight, found {reward_head_keys}")
        shape = tuple(handle.get_tensor(reward_head_keys[0]).shape)
    if shape != PUBLISHED_RM_HEAD_SHAPE:
        raise RuntimeError(
            f"unexpected rm_head shape: expected={PUBLISHED_RM_HEAD_SHAPE}, observed={shape}"
        )

    embeddings = root / "special_token_embeddings.safetensors"
    if not embeddings.is_file():
        raise FileNotFoundError(f"required special-token embeddings are missing: {embeddings}")
    embedding_sha = sha256(embeddings)
    if embedding_sha != PUBLISHED_SPECIAL_EMBEDDINGS_SHA256:
        raise RuntimeError(
            "special-token embedding checksum mismatch: "
            f"expected={PUBLISHED_SPECIAL_EMBEDDINGS_SHA256}, observed={embedding_sha}"
        )
    with safe_open(embeddings, framework="pt", device="cpu") as handle:
        embedding_shape = tuple(handle.get_tensor("special_token_embeddings").shape)
    if embedding_shape != PUBLISHED_RM_HEAD_SHAPE:
        raise RuntimeError(
            "special-token embedding shape must equal (3, hidden_size): "
            f"expected={PUBLISHED_RM_HEAD_SHAPE}, observed={embedding_shape}"
        )

    print("Q-Prefer artifact verified")
    print(f"adapter: {root}")
    print(f"sha256: {observed_sha}")
    print(f"tensors: {len(tensor_keys)}")
    print(f"rm_head: {shape}")
    print(f"special embeddings sha256: {embedding_sha}")


if __name__ == "__main__":
    main()