#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import unicodedata
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable, List, Tuple

CHUNK = 1024 * 1024  # 1 MiB
DEFAULT_EXCLUDES = {"__MACOSX",}

def nfc(s: str) -> str:
    return unicodedata.normalize("NFC", s)

def hash_file(path: Path) -> Tuple[str, str]:
    """Return streaming (md5, sha256) for a file."""
    md5 = hashlib.md5()
    sha = hashlib.sha256()
    with path.open("rb", buffering=CHUNK) as f:
        for chunk in iter(lambda: f.read(CHUNK), b""):
            md5.update(chunk)
            sha.update(chunk)
    return md5.hexdigest(), sha.hexdigest()

def iter_files(root: Path, follow_symlinks: bool = False) -> Iterable[Path]:
    """Yield files under root, skipping symlinks unless allowed."""
    # os.walk gives control over symlinks; pathlib.rglob does not.
    for dirpath, dirnames, filenames in os.walk(root, followlinks=follow_symlinks):
        # Skip hidden dot-directories at the top level? Leave behavior to caller.
        for name in filenames:
            p = Path(dirpath) / name
            try:
                # Skip if symlinked file and not following symlinks
                if not follow_symlinks and p.is_symlink():
                    continue
                if p.is_file():
                    yield p
            except OSError:
                # permissions/races: skip
                continue

def build_manifest(folder: Path, excludes: set[str]) -> dict:
    """Create a manifest dict for a firmware folder (recursive)."""
    files: List[dict] = []
    relpaths: List[str] = []

    for p in iter_files(folder, follow_symlinks=False):
        # Exclude by top-level folder name (e.g., "__MACOSX") or dot dirs at root
        parts = p.relative_to(folder).parts
        if not parts:
            continue
        top = parts[0]
        if top.startswith(".") or top in excludes:
            continue

        rel = nfc(Path(*parts).as_posix())
        try:
            stats = p.stat()
            md5, sha256 = hash_file(p)
        except (OSError, PermissionError) as e:
            print(f"[warn] Skipping {p}: {e}", file=sys.stderr)
            continue

        files.append({
            "path": rel,
            "size": stats.st_size,
            "mtime": int(stats.st_mtime),
            "md5": md5,
            "sha256": sha256,
        })
        relpaths.append(rel)

    # Stable, case-insensitive order
    files.sort(key=lambda f: f["path"].lower())

    # Stream the tree hash to avoid building a huge string
    tree = hashlib.sha256()
    for f in files:
        line = f'{f["sha256"]}  {f["path"]}\n'.encode("utf-8")
        tree.update(line)
    tree_sha256 = tree.hexdigest()

    return {
        "firmware": folder.name,
        "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "algorithms": ["sha256", "md5"],
        "file_count": len(files),
        "tree_sha256": tree_sha256,
        "files": files,
    }

def main(argv: List[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Build per-folder firmware manifests.")
    ap.add_argument("--root", type=Path, default=Path(__file__).resolve().parent,
                    help="Root directory containing firmware subfolders (default: script dir)")
    ap.add_argument("--single", action="store_true",
                    help="Treat --root itself as the firmware folder (no subfolder scan).")
    ap.add_argument("--exclude", action="append", default=[],
                    help="Top-level folder names to exclude (can repeat).")
    ap.add_argument("--dry-run", action="store_true", help="Scan but do not write JSON files.")
    args = ap.parse_args(argv)

    root = args.root
    excludes = DEFAULT_EXCLUDES | set(args.exclude)

    if not root.is_dir():
        print(f"[!] Root is not a directory: {root}", file=sys.stderr)
        return 2

    print(f"Working directory: {root}")

    if args.single:
        subfolders = [root]
    else:
        subfolders = sorted([p for p in root.iterdir() if p.is_dir()],
                            key=lambda p: p.name.lower())

    if not subfolders:
        print("[!] No firmware folders found here.")
        return 1

    total_files = 0
    for folder in subfolders:
        name = folder.name
        if name.startswith(".") or name.upper() in {"__MACOSX"}:
            continue

        # Skip empty
        if not any(iter_files(folder)):
            print(f" - {name}: [skipped] (empty)")
            continue

        print(f" - {name}: hashing...", end="", flush=True)
        manifest = build_manifest(folder, excludes=excludes)
        total_files += manifest["file_count"]

        outfile = (root if not args.single else folder.parent) / f"{name}.json"
        if args.dry_run:
            print(f" [dry-run] would write {outfile.name} ({manifest['file_count']} files)")
        else:
            outfile.write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
            print(f" wrote {outfile.name} ({manifest['file_count']} files)")

    print(f"\nDone. Total files hashed: {total_files}")
    return 0

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