#!/usr/bin/env python3
import sys
import os
import yaml
import hashlib

def sha256sum(filename, buf_size=65536):
    sha256 = hashlib.sha256()
    with open(filename, 'rb') as f:
        while True:
            data = f.read(buf_size)
            if not data:
                break
            sha256.update(data)
    return sha256.hexdigest()

def verify_model(model_yaml_path):
    if not os.path.isfile(model_yaml_path):
        print(f"❌ Model YAML not found: {model_yaml_path}")
        sys.exit(1)

    with open(model_yaml_path, 'r', encoding='utf-8') as f:
        model_data = yaml.safe_load(f)

    base_dir = os.path.dirname(model_yaml_path)
    all_ok = True

    for fmt in model_data.get("formats", []):
        for variant in fmt.get("variants", []):
            for file_path in variant.get("files", []):
                checksum_expected = variant.get("checksums", {}).get(file_path)
                abs_path = os.path.join(base_dir, file_path)

                if not os.path.isfile(abs_path):
                    print(f"❌ Missing file: {abs_path}")
                    all_ok = False
                    continue

                if not checksum_expected:
                    print(f"⚠️ No checksum for {file_path}, skipping verification.")
                    continue

                checksum_actual = sha256sum(abs_path)
                if checksum_actual.lower() == checksum_expected.lower():
                    print(f"✅ {file_path} OK")
                else:
                    print(f"❌ {file_path} checksum mismatch! Expected {checksum_expected}, got {checksum_actual}")
                    all_ok = False

    if all_ok:
        print("✅ All files verified successfully.")
    else:
        print("❌ Verification failed.")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <path-to-model.yaml>")
        sys.exit(1)

    verify_model(sys.argv[1])