#!/usr/bin/env python3
"""
Check or update the codegen stamp file.

Uses only the Python standard library (hashlib, pathlib, sys) so it can
run without a virtual environment.

Modes:
    --check   Exit 0 if stamp is up-to-date, exit 1 if stale/missing.
    --update  Recompute the hash and write it to the stamp file.

Usage:
    python update_codegen_stamp.py --check  <stamp_file> <input_files...>
    python update_codegen_stamp.py --update <stamp_file> <input_files...>
"""

import hashlib
import sys
from pathlib import Path


def compute_combined_hash(input_files: list[str]) -> str:
    """Compute a combined SHA-256 hash of all input files.

    Algorithm: compute each file's SHA-256 hex digest, concatenate them
    all, then SHA-256 the concatenation.
    """
    parts = []
    for filepath in input_files:
        file_hash = hashlib.sha256(Path(filepath).read_bytes()).hexdigest()
        parts.append(file_hash)

    combined = "".join(parts)
    return hashlib.sha256(combined.encode()).hexdigest()


def read_stamp_hash(stamp_file: str) -> str:
    """Read the COMBINED_HASH from an existing stamp file, or '' if missing."""
    path = Path(stamp_file)
    if not path.exists():
        return ""
    for line in path.read_text().splitlines():
        if line.startswith("COMBINED_HASH="):
            return line.split("=", 1)[1]
    return ""


def main():
    if len(sys.argv) < 4 or sys.argv[1] not in ("--check", "--update"):
        print(
            f"Usage: {sys.argv[0]} --check|--update <stamp_file> <input_files...>",
            file=sys.stderr,
        )
        sys.exit(2)

    mode = sys.argv[1]
    stamp_file = sys.argv[2]
    input_files = sys.argv[3:]

    current_hash = compute_combined_hash(input_files)

    if mode == "--check":
        stamp_hash = read_stamp_hash(stamp_file)
        if current_hash == stamp_hash:
            sys.exit(0)
        else:
            sys.exit(1)

    # --update
    with open(stamp_file, "w") as fp:
        fp.write(
            "# Auto-generated by protocol autogen - do not edit manually.\n"
            "# This file tracks input hashes to avoid unnecessary code regeneration.\n"
            "# It should be checked into version control alongside the generated files.\n"
        )
        fp.write(f"COMBINED_HASH={current_hash}\n")


if __name__ == "__main__":
    main()