#!/usr/bin/env python3
"""Read-only inventory for Chrome's on-device model directory.

The script requires an explicit Chrome user-data root. It does not discover
profiles, read model contents, change policies, or delete files.
"""

from __future__ import annotations

import argparse
import csv
import os
import sys
import tempfile
from dataclasses import asdict, dataclass
from pathlib import Path


MODEL_DIRECTORY = "OptGuideOnDeviceModel"


@dataclass(frozen=True)
class AuditResult:
    scan_root: str
    model_directory: str
    status: str
    file_count: int
    total_bytes: int
    weights_file_count: int
    weights_bytes: int
    skipped_links: int
    errors: int


def _is_link_or_junction(path: Path) -> bool:
    if path.is_symlink():
        return True
    is_junction = getattr(path, "is_junction", None)
    return bool(is_junction and is_junction())


def audit(user_data_root: Path) -> AuditResult:
    root = user_data_root.expanduser()
    if not root.is_dir():
        raise ValueError("the supplied user-data root is not a directory")

    model_root = root / MODEL_DIRECTORY
    if not model_root.exists():
        return AuditResult(
            scan_root="explicit-user-data-root",
            model_directory=MODEL_DIRECTORY,
            status="absent",
            file_count=0,
            total_bytes=0,
            weights_file_count=0,
            weights_bytes=0,
            skipped_links=0,
            errors=0,
        )
    if not model_root.is_dir() or _is_link_or_junction(model_root):
        raise ValueError("the model path is not a regular directory")

    file_count = 0
    total_bytes = 0
    weights_file_count = 0
    weights_bytes = 0
    skipped_links = 0
    errors = 0
    pending = [model_root]

    while pending:
        current = pending.pop()
        try:
            entries = list(current.iterdir())
        except OSError:
            errors += 1
            continue

        for entry in entries:
            try:
                if _is_link_or_junction(entry):
                    skipped_links += 1
                elif entry.is_dir():
                    pending.append(entry)
                elif entry.is_file():
                    size = entry.stat().st_size
                    file_count += 1
                    total_bytes += size
                    if entry.name.casefold() == "weights.bin":
                        weights_file_count += 1
                        weights_bytes += size
            except OSError:
                errors += 1

    if errors:
        status = "incomplete-scan"
    elif weights_file_count:
        status = "present-with-weights"
    elif file_count:
        status = "present-without-weights"
    else:
        status = "present-empty"

    return AuditResult(
        scan_root="explicit-user-data-root",
        model_directory=MODEL_DIRECTORY,
        status=status,
        file_count=file_count,
        total_bytes=total_bytes,
        weights_file_count=weights_file_count,
        weights_bytes=weights_bytes,
        skipped_links=skipped_links,
        errors=errors,
    )


def write_result(result: AuditResult) -> None:
    row = asdict(result)
    writer = csv.DictWriter(sys.stdout, fieldnames=list(row))
    writer.writeheader()
    writer.writerow(row)


def self_test() -> int:
    cases: list[tuple[str, str]] = [
        ("absent", "absent"),
        ("empty", "present-empty"),
        ("metadata-only", "present-without-weights"),
        ("weights", "present-with-weights"),
    ]
    rows: list[dict[str, str]] = []

    with tempfile.TemporaryDirectory(prefix="chrome-ai-audit-") as tmp:
        base = Path(tmp)
        for case, expected in cases:
            root = base / case
            root.mkdir()
            if case != "absent":
                model = root / MODEL_DIRECTORY / "test-version"
                model.mkdir(parents=True)
                if case == "metadata-only":
                    (model / "manifest.json").write_text("{}\n", encoding="utf-8")
                elif case == "weights":
                    (model / "weights.bin").write_bytes(b"x" * 1536)

            actual = audit(root).status
            rows.append(
                {
                    "case": case,
                    "expected": expected,
                    "actual": actual,
                    "result": "PASS" if actual == expected else "FAIL",
                }
            )

    writer = csv.DictWriter(sys.stdout, fieldnames=list(rows[0]))
    writer.writeheader()
    writer.writerows(rows)
    passed = sum(row["result"] == "PASS" for row in rows)
    print(f"summary,{passed}/{len(rows)} PASS")
    return 0 if passed == len(rows) else 1


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Count files and bytes in an explicitly supplied Chrome "
            "OptGuideOnDeviceModel directory without reading file contents."
        )
    )
    parser.add_argument(
        "user_data_root",
        nargs="?",
        type=Path,
        help="Chrome user-data root that may contain OptGuideOnDeviceModel",
    )
    parser.add_argument(
        "--self-test",
        action="store_true",
        help="run four synthetic, temporary test cases",
    )
    args = parser.parse_args()
    if args.self_test == (args.user_data_root is not None):
        parser.error("provide either --self-test or one user-data root")
    return args


def main() -> int:
    args = parse_args()
    if args.self_test:
        return self_test()
    try:
        result = audit(args.user_data_root)
    except (OSError, ValueError) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        return 2
    write_result(result)
    return 0


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