#!/usr/bin/env python3
"""Strict immutable provenance gates for interchangeable MOSS-SFT components.

The combined pipeline always uses original OpenMOSS MOSS-Image-SFT architecture,
tokenizer, scheduler and sampler. Component paths are local staged weight folders
and must have a matching complete source-revision receipt; never select a model
based only on a filename, a loose remote branch, or an unrelated Qwen version.
"""
from pathlib import Path
import hashlib
import json
import torch

REV="fc5de0d56fc910cbca6f6625a70f5f582b26deb1"

def sha(path):
    h=hashlib.sha256()
    with Path(path).open("rb") as f:
        for piece in iter(lambda:f.read(4*1024*1024),b""):
            h.update(piece)
    return h.hexdigest()

def inspect_qwen(root,verify_sha=True):
    root=Path(root).resolve()
    manifest=root/"COMPLETE_ENCODER.json"
    if not manifest.is_file():raise RuntimeError("Missing full encoder export receipt")
    m=json.loads(manifest.read_text())
    if (m.get("status")!="COMPLETE_RELOADED_REAL_TORCHAO_FP8WO_ENCODER"
        or m.get("source_revision")!=REV or m.get("converted_linear_weights")!=248
        or m.get("qwen_text_encoder_stored_shards")!=4
        or m.get("quantization")!="Float8WeightOnlyConfig"):
        raise RuntimeError("Not the validated original-Qwen3.5 FP8 weight-only encoder")
    enc=root/"text_encoder"
    assert_file=enc/"config.json"
    if not assert_file.is_file() or sha(assert_file)!=m["config_sha256"]:
        raise RuntimeError("Encoder architecture config changed")
    total=0
    for x in m["shards"]:
        p=enc/x["file"]
        if not p.is_file() or p.stat().st_size!=x["bytes"]:
            raise RuntimeError("Incomplete stored Qwen FP8WO shard: "+x["file"])
        if verify_sha and sha(p)!=x["sha256"]:
            raise RuntimeError("Encoder stored shard SHA-256 mismatch: "+x["file"])
        total+=x["bytes"]
    if total!=m["total_text_encoder_bytes"]:
        raise RuntimeError("Qwen FP8WO total size mismatch")
    index=json.loads((enc/"pytorch_model.bin.index.json").read_text())
    if len(index["weight_map"])!=m["source_keys"]:
        raise RuntimeError("Qwen weight shard index incomplete")
    return m

def inspect_vae(root,verify_sha=True):
    root=Path(root).resolve()
    p=root/"COMPLETE_VAE.json"
    if not p.is_file():raise RuntimeError("Missing VAE precision export receipt")
    m=json.loads(p.read_text())
    if (m.get("status")!="COMPLETE_DTYPE_CAST_FROM_PINNED_ORIGINAL"
        or m.get("source_revision")!=REV or m.get("dtype") not in ("fp16","bf16")):
        raise RuntimeError("Not an original Flux2 VAE precision checkpoint")
    vae=root/"vae"/"diffusion_pytorch_model.safetensors"
    config=root/"vae"/"config.json"
    if not vae.is_file() or vae.stat().st_size!=m["output_bytes"]:
        raise RuntimeError("VAE weights incomplete")
    if not config.is_file() or sha(config)!=m["config_sha256"]:
        raise RuntimeError("VAE config changed")
    if verify_sha and sha(vae)!=m["output_sha256"]:
        raise RuntimeError("VAE weights SHA-256 mismatch")
    return (m,torch.float16 if m["dtype"]=="fp16" else torch.bfloat16)