#!/usr/bin/env python3
"""Enforce Blacksite's module-size debt ratchet."""

from __future__ import annotations

import argparse
import tempfile
import tomllib
from pathlib import Path


def nonblank_lines(path: Path) -> int:
    return sum(1 for line in path.read_text(encoding="utf-8").splitlines() if line.strip())


def module_limit(relative: str, policy: dict) -> int:
    if relative in policy.get("shells", {}):
        return int(policy["limits"]["ui_shell"])
    if relative.startswith("crates/editor/src/ui/"):
        return int(policy["limits"]["ui_module"])
    return int(policy["limits"]["production_rust_module"])


def validate_exception(entry: dict) -> list[str]:
    required = ("path", "issue", "rationale", "maximum", "extraction_target", "expiry_milestone")
    missing = [name for name in required if not entry.get(name)]
    errors = [f"invalid exception for {entry.get('path', '<unknown>')}: missing {name}" for name in missing]
    if entry.get("issue") and not str(entry["issue"]).startswith("#"):
        errors.append(f"invalid exception for {entry['path']}: issue must be a tracker reference")
    return errors


def audit(root: Path, policy: dict) -> tuple[list[str], list[str]]:
    errors: list[str] = []
    notes: list[str] = []
    baselines = {str(path): int(value) for path, value in policy.get("baselines", {}).items()}
    exceptions = {entry.get("path"): entry for entry in policy.get("exceptions", [])}
    for entry in policy.get("exceptions", []):
        errors.extend(validate_exception(entry))

    for path in sorted((root / "crates").glob("**/*.rs")):
        relative = path.relative_to(root).as_posix()
        count = nonblank_lines(path)
        limit = module_limit(relative, policy)
        baseline = baselines.get(relative)
        exception = exceptions.get(relative)

        if baseline is not None and count > baseline:
            errors.append(f"{relative}: {count} nonblank lines exceeds frozen baseline {baseline}")
            continue
        if baseline is not None and count < baseline:
            notes.append(f"{relative}: shrank from baseline {baseline} to {count}")
        if count <= limit:
            continue
        if baseline is not None:
            continue
        if exception is not None and count <= int(exception["maximum"]):
            notes.append(f"{relative}: temporary {exception['issue']} exception ({count}/{exception['maximum']})")
            continue
        errors.append(f"{relative}: {count} nonblank lines exceeds module budget {limit}")
    return errors, notes


def load_policy(path: Path) -> dict:
    with path.open("rb") as handle:
        policy = tomllib.load(handle)
    if policy.get("version") != 1:
        raise SystemExit(f"unsupported architecture policy version in {path}")
    return policy


def self_test() -> int:
    with tempfile.TemporaryDirectory(prefix="blacksite-architecture-audit-") as directory:
        root = Path(directory)
        (root / "crates/editor/src/ui").mkdir(parents=True)
        policy = {
            "limits": {"ui_shell": 2, "ui_module": 3, "production_rust_module": 4},
            "shells": {"crates/editor/src/ui/shell.rs": 2},
            "baselines": {"crates/editor/src/ui/legacy.rs": 5},
            "exceptions": [],
        }
        (root / "crates/editor/src/ui/new.rs").write_text("a\nb\nc\nd\n", encoding="utf-8")
        errors, _ = audit(root, policy)
        assert any("new.rs" in error for error in errors)
        (root / "crates/editor/src/ui/new.rs").write_text("a\nb\n", encoding="utf-8")
        (root / "crates/editor/src/ui/legacy.rs").write_text("a\nb\nc\nd\ne\nf\n", encoding="utf-8")
        errors, _ = audit(root, policy)
        assert any("frozen baseline" in error for error in errors)
        (root / "crates/editor/src/ui/legacy.rs").write_text("a\nb\nc\nd\n", encoding="utf-8")
        errors, notes = audit(root, policy)
        assert not errors and any("shrank" in note for note in notes)
        policy["exceptions"] = [{
            "path": "crates/editor/src/ui/new.rs", "issue": "#1", "rationale": "test",
            "maximum": 4, "extraction_target": "test module", "expiry_milestone": "M2",
        }]
        (root / "crates/editor/src/ui/new.rs").write_text("a\nb\nc\nd\n", encoding="utf-8")
        errors, _ = audit(root, policy)
        assert not errors
        policy["exceptions"][0]["issue"] = "missing-prefix"
        errors, _ = audit(root, policy)
        assert any("tracker reference" in error for error in errors)
    print("PASS architecture-audit self-test")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=("check", "self-test"), nargs="?", default="check")
    parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[2])
    arguments = parser.parse_args()
    if arguments.command == "self-test":
        return self_test()
    root = arguments.root.resolve()
    policy = load_policy(root / ".codex/architecture.toml")
    errors, notes = audit(root, policy)
    for note in notes:
        print(f"NOTE {note}")
    if errors:
        for error in errors:
            print(f"FAIL {error}")
        return 1
    print("PASS architecture-audit")
    return 0


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