"""Verify the original Swift export without changing or loading its weights."""
import argparse
import hashlib
import json
import math
import struct
import sys
from datetime import datetime, timezone
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument("--root", required=True)
parser.add_argument("--manifest-sha256", required=True)
args = parser.parse_args()
root = Path(args.root)


def digest(path):
    result = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(16 * 1024 * 1024), b""):
            result.update(block)
    return result.hexdigest()


manifest_path = root / "EXPORT_MANIFEST.json"
assert digest(manifest_path) == args.manifest_sha256, "Export provenance mismatch"
manifest = json.loads(manifest_path.read_text())
assert manifest["dtype"] == "bfloat16"
shards = manifest["output_shards"]
assert len(shards) == 18
index = json.loads((root / "model.safetensors.index.json").read_text())
expected_shards = {item["file"] for item in shards}
assert set(index["weight_map"].values()) == expected_shards
assert {f.name for f in root.glob("*.safetensors")} == expected_shards
files, tensors = [], {}
for item in shards:
    name = item["file"]
    assert Path(name).name == name
    path = root / name
    size = path.stat().st_size
    assert size == item["bytes"], f"Size mismatch: {name}"
    sha = digest(path)
    assert sha == item["sha256"], f"SHA256 mismatch: {name}"
    with path.open("rb") as stream:
        header_size = struct.unpack("<Q", stream.read(8))[0]
        assert 0 < header_size < 16 * 1024 * 1024
        header = json.loads(stream.read(header_size))
    ranges = []
    for tensor_name, metadata in header.items():
        if tensor_name == "__metadata__":
            continue
        assert tensor_name not in tensors, f"Duplicate tensor: {tensor_name}"
        assert index["weight_map"][tensor_name] == name
        assert metadata["dtype"] == "BF16", tensor_name
        start, end = metadata["data_offsets"]
        count = math.prod(metadata["shape"])
        assert end - start == count * 2
        ranges.append((start, end))
        tensors[tensor_name] = dict(metadata, shard=name, parameters=count, bytes=end-start)
    ranges.sort()
    assert ranges[0][0] == 0
    assert all(left[1] == right[0] for left, right in zip(ranges, ranges[1:]))
    assert ranges[-1][1] + 8 + header_size == size
    files.append({"name": name, "bytes": size, "sha256": sha, "kind": "bf16_shard"})
    print(f"VERIFIED {name} {size} bytes", file=sys.stderr, flush=True)
assert set(tensors) == set(index["weight_map"])
assert sum(v["bytes"] for v in tensors.values()) == manifest["total_tensor_bytes"]
assert index["metadata"]["total_size"] == manifest["total_tensor_bytes"]
for name, sha in manifest["asset_sha256"].items():
    path = root / name
    assert path.is_file(), f"Missing asset: {name}"
    actual = digest(path)
    assert actual == sha, f"Asset hash mismatch: {name}"
    files.append({"name": name, "bytes": path.stat().st_size, "sha256": actual, "kind": "source_asset"})
recorded = {item["name"] for item in files}
for path in sorted(root.iterdir()):
    if path.is_file() and path.name not in recorded:
        files.append({"name": path.name, "bytes": path.stat().st_size, "sha256": digest(path), "kind": "additional_source_file"})
result = {
    "status": "PASS", "recorded_at": datetime.now(timezone.utc).isoformat(),
    "source_root": str(root), "manifest_sha256": args.manifest_sha256,
    "shard_count": len(shards), "shard_bytes": sum(item["bytes"] for item in shards),
    "source_file_bytes": sum(item["bytes"] for item in files),
    "tensor_count": len(tensors), "parameter_count": sum(v["parameters"] for v in tensors.values()),
    "files": files, "tensors": tensors,
    "config": json.loads((root / "config.json").read_text()),
    "custom_modeling_files": [path.name for path in root.glob("*.py")],
    "special_tokens_map_present": (root / "special_tokens_map.json").is_file(),
    "original_unchanged": True,
}
print(json.dumps(result, indent=2))