#!/usr/bin/env python3

import argparse
import hashlib
import os
from pathlib import Path
from typing import List, Set, Tuple


ARCH_DIRS = ("arm64", "la64", "rv64", "ppc64le")
GROUPS = (
    ("COMMON", Path("src/dynarec"), set(ARCH_DIRS)),
    ("ARM64", Path("src/dynarec/arm64"), set()),
    ("LOONGARCH64", Path("src/dynarec/la64"), set()),
    ("RISCV64", Path("src/dynarec/rv64"), set()),
    ("PPC64LE", Path("src/dynarec/ppc64le"), set()),
)


def iter_files(base: Path, excluded_dirs: Set[str]) -> List[Path]:
    files: List[Path] = []
    for path in sorted(base.rglob("*")):
        if not path.is_file():
            continue
        rel_parts = path.relative_to(base).parts
        if excluded_dirs and rel_parts and rel_parts[0] in excluded_dirs:
            continue
        files.append(path)
    return files


def hash_files(root: Path, files: List[Path]) -> int:
    digest = hashlib.sha256()
    for path in files:
        rel_path = path.relative_to(root).as_posix().encode("utf-8")
        digest.update(rel_path)
        digest.update(b"\0")
        digest.update(path.read_bytes())
        digest.update(b"\0")
    value = int.from_bytes(digest.digest()[:3], "big")
    return value or 1


def render_header(root: Path) -> str:
    values: List[Tuple[str, int]] = []
    for name, rel_base, excluded_dirs in GROUPS:
        base = root / rel_base
        files = iter_files(base, excluded_dirs)
        if not files:
            raise RuntimeError(f"no files found for {name.lower()} hash in {base}")
        values.append((name, hash_files(root, files)))
    hashes = dict(values)

    lines = [
        "/* Auto-generated by tools/gen_dynacache_hashes.py. Do not edit. */",
        "#ifndef __DYNACACHE_HASHES_H__",
        "#define __DYNACACHE_HASHES_H__",
        "",
        f"#define DYNAREC_VERSION 0x{hashes['COMMON']:06X}u",
        "",
        "#if defined(ARM64)",
        f"#define ARCH_VERSION 0x{hashes['ARM64']:06X}u",
        "#elif defined(RV64)",
        f"#define ARCH_VERSION 0x{hashes['RISCV64']:06X}u",
        "#elif defined(LA64)",
        f"#define ARCH_VERSION 0x{hashes['LOONGARCH64']:06X}u",
        "#elif defined(PPC64LE)",
        f"#define ARCH_VERSION 0x{hashes['PPC64LE']:06X}u",
        "#else",
        "#error Unsupported architecture",
        "#endif",
        "",
        "#endif",
    ]
    return "\n".join(lines) + "\n"


def write_if_changed(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.exists() and path.read_text(encoding="utf-8") == content:
        return
    tmp_path = path.with_suffix(f"{path.suffix}.{os.getpid()}.tmp")
    tmp_path.write_text(content, encoding="utf-8")
    tmp_path.replace(path)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    root = Path(args.root).resolve()
    output = Path(args.output).resolve()
    write_if_changed(output, render_header(root))
    return 0


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