#!/usr/bin/env python3

import argparse
import hashlib
import json
import re
import sys
from pathlib import Path
from typing import Dict, Iterable, Tuple


REQUIRED_ROLES = ("model", "lex", "srcvocab", "trgvocab")


class RuntimeDataError(RuntimeError):
    pass


def _version_key(version: object) -> Tuple[Tuple[int, object], ...]:
    parts = re.split(r"(\d+)", str(version))
    return tuple((0, int(part)) if part.isdigit() else (1, part) for part in parts)


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


def _load_records(records_path: Path) -> Iterable[dict]:
    try:
        document = json.loads(records_path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as error:
        raise RuntimeDataError(f"invalid records.json: {error}") from error

    records = document.get("data") if isinstance(document, dict) else None
    if not isinstance(records, list):
        raise RuntimeDataError("invalid records.json: data must be a list")
    return records


def verify_runtime_data(records_path: Path, model_dir: Path) -> Dict[str, Path]:
    candidates = [
        record
        for record in _load_records(records_path)
        if isinstance(record, dict)
        and record.get("sourceLanguage") == "en"
        and record.get("targetLanguage") == "zh-Hans"
    ]

    selected = {}
    for role in REQUIRED_ROLES:
        role_records = [record for record in candidates if record.get("fileType") == role]
        if not role_records:
            raise RuntimeDataError(f"missing required record role: {role}")
        selected[role] = max(role_records, key=lambda record: _version_key(record.get("version", "")))

    verified = {}
    for role, record in selected.items():
        attachment = record.get("attachment")
        filename = attachment.get("filename") if isinstance(attachment, dict) else None
        if not isinstance(filename, str) or not filename:
            raise RuntimeDataError(f"invalid attachment filename for role: {role}")
        decompressed_name = filename[:-4] if filename.endswith(".zst") else filename
        if Path(decompressed_name).name != decompressed_name:
            raise RuntimeDataError(f"unsafe attachment filename for role: {role}")

        model_path = model_dir / decompressed_name
        if not model_path.is_file():
            raise RuntimeDataError(f"missing model file for role {role}: {decompressed_name}")

        expected_size = record.get("decompressedSize")
        if expected_size is not None and model_path.stat().st_size != expected_size:
            raise RuntimeDataError(f"size mismatch for role {role}: {decompressed_name}")

        expected_hash = record.get("decompressedHash")
        if not isinstance(expected_hash, str) or not expected_hash:
            raise RuntimeDataError(f"missing decompressedHash for role: {role}")
        if _sha256(model_path) != expected_hash.lower():
            raise RuntimeDataError(f"hash mismatch for role {role}: {decompressed_name}")
        verified[role] = model_path

    expected_names = {path.name for path in verified.values()}
    actual_names = {entry.name for entry in model_dir.iterdir()}
    extras = sorted(actual_names - expected_names)
    if extras:
        raise RuntimeDataError(f"unexpected model files: {', '.join(extras)}")

    return verified


def main() -> int:
    parser = argparse.ArgumentParser(description="Verify the fixed en to zh-Hans runtime data")
    parser.add_argument("--records", required=True, type=Path)
    parser.add_argument("--model-dir", required=True, type=Path)
    args = parser.parse_args()

    try:
        verified = verify_runtime_data(args.records, args.model_dir)
    except RuntimeDataError as error:
        print(f"runtime data verification failed: {error}", file=sys.stderr)
        return 1

    for role in REQUIRED_ROLES:
        print(f"verified {role}: {verified[role]}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())