#!/usr/bin/env python3
"""Build a deterministic, content-hashed inventory of ignored local assets."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path


DEFAULT_INCLUDES = ("data", "artifacts", "mlruns")


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


def build_inventory(root: Path, includes: tuple[str, ...]) -> dict:
    root = root.resolve()
    entries: list[dict] = []
    summaries: dict[str, dict] = {}
    for include in includes:
        relative = Path(include)
        if (
            relative.is_absolute()
            or ".." in relative.parts
            or len(relative.parts) != 1
        ):
            raise ValueError(f"unsafe include: {include}")
        directory = root / relative
        file_count = 0
        total_bytes = 0
        if directory.exists() and not directory.is_dir():
            raise ValueError(f"include is not a directory: {include}")
        if directory.is_dir():
            for path in sorted(directory.rglob("*")):
                if path.is_symlink():
                    raise ValueError(
                        f"symlink is not allowed in inventory: "
                        f"{path.relative_to(root)}"
                    )
                if not path.is_file():
                    continue
                size = path.stat().st_size
                entries.append(
                    {
                        "path": path.relative_to(root).as_posix(),
                        "size": size,
                        "sha256": _sha256(path),
                    }
                )
                file_count += 1
                total_bytes += size
        summaries[include] = {
            "available": directory.is_dir(),
            "files": file_count,
            "bytes": total_bytes,
        }
    body = {
        "schema_version": "quant-os-local-assets-inventory/1.0",
        "includes": list(includes),
        "summaries": summaries,
        "entries": entries,
    }
    body["inventory_sha256"] = hashlib.sha256(
        json.dumps(
            body,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            allow_nan=False,
        ).encode("utf-8")
    ).hexdigest()
    return body


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("root", type=Path)
    parser.add_argument(
        "--include",
        action="append",
        dest="includes",
        help="one top-level ignored directory; repeat as needed",
    )
    parser.add_argument("--output", type=Path)
    parser.add_argument("--compare", type=Path)
    args = parser.parse_args()
    includes = tuple(args.includes or DEFAULT_INCLUDES)
    try:
        inventory = build_inventory(args.root, includes)
    except (OSError, ValueError) as exc:
        print(f"inventory failed: {exc}", file=sys.stderr)
        return 2
    encoded = (
        json.dumps(
            inventory,
            ensure_ascii=False,
            sort_keys=True,
            indent=2,
            allow_nan=False,
        )
        + "\n"
    )
    if args.compare is not None:
        try:
            expected = json.loads(args.compare.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as exc:
            print(f"cannot read comparison inventory: {exc}", file=sys.stderr)
            return 2
        if inventory != expected:
            print("local asset inventory differs", file=sys.stderr)
            return 1
    if args.output is None:
        print(encoded, end="")
    else:
        args.output.write_text(encoded, encoding="utf-8")
        print(
            json.dumps(
                {
                    "ok": True,
                    "files": len(inventory["entries"]),
                    "inventory_sha256": inventory["inventory_sha256"],
                    "output": args.output.name,
                },
                sort_keys=True,
            )
        )
    return 0


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