#!/usr/bin/env python3
"""Auto-generate OWASP Agentic Top 10 coverage artefacts.

Walks `agent_audit_kit.rules.builtin.RULES` and emits:

- `docs/owasp-agentic-coverage.md` — the human-facing coverage table.
- `public/owasp-agentic-coverage.json` — machine-readable artefact the
  OWASP Agentic Reference-Tool submission packet links to. Schema:

      {
        "schema_version": "1",
        "last_updated": "<ISO8601 UTC>",
        "aak_version": "<pyproject version>",
        "rule_count": <int>,
        "coverage": [
          {
            "asi_id": "ASI01",
            "title": "Goal Hijack",
            "rule_density": <int>,
            "rules": [
              {"id": "...", "severity": "...",
               "cve_references": [...], "aicm_references": [...]}
            ]
          },
          ...
        ]
      }

- README.md block between `<!-- owasp-coverage:start -->` and
  `<!-- owasp-coverage:end -->` — compact summary.

`tests/test_owasp_agentic_coverage.py` is the regression fence; missing
slots still cause `main()` to exit non-zero.
"""

from __future__ import annotations

import argparse
import json
import pathlib
import sys
from collections import defaultdict
from datetime import datetime, timezone

try:  # Python 3.11+ stdlib
    import tomllib  # type: ignore[import-not-found]
except ModuleNotFoundError:  # pragma: no cover
    import tomli as tomllib  # type: ignore[no-redef]


REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
DOC_PATH = REPO_ROOT / "docs" / "owasp-agentic-coverage.md"
README_PATH = REPO_ROOT / "README.md"
PYPROJECT = REPO_ROOT / "pyproject.toml"
DEFAULT_JSON_PATH = REPO_ROOT / "public" / "owasp-agentic-coverage.json"

_MARKER_START = "<!-- owasp-coverage:start -->"
_MARKER_END = "<!-- owasp-coverage:end -->"

# Canonical OWASP Agentic Top 10 2026 titles. Source:
# https://genai.owasp.org/resource/owasp-top-10-for-agentic-applications-for-2026/
ASI_TITLES: dict[str, str] = {
    "ASI01": "Goal Hijack",
    "ASI02": "Tool Misuse",
    "ASI03": "Memory Poisoning",
    "ASI04": "Identity & Privilege Abuse",
    "ASI05": "Cascading Failures",
    "ASI06": "Unauthorized Capability Acquisition",
    "ASI07": "Plan Injection",
    "ASI08": "Agent Communication Poisoning",
    "ASI09": "Resource Abuse",
    "ASI10": "Supply-Chain",
}


def build_coverage() -> dict[str, list[str]]:
    from agent_audit_kit.rules.builtin import RULES

    cov: dict[str, list[str]] = defaultdict(list)
    for rid, rule in sorted(RULES.items()):
        for asi in rule.owasp_agentic_references:
            cov[asi].append(rid)
    return dict(cov)


def render_markdown(coverage: dict[str, list[str]]) -> str:
    lines: list[str] = [
        "# OWASP Top 10 for Agentic Applications — 2026 coverage",
        "",
        "Auto-generated by `scripts/gen_owasp_coverage.py`. Every AAK rule "
        "that claims an `ASI<nn>` reference lands in the corresponding row. "
        "`tests/test_owasp_agentic_coverage.py` is the regression fence — "
        "it fails CI if any row has zero rules.",
        "",
        "Source: <https://genai.owasp.org/resource/owasp-top-10-for-agentic-applications-for-2026/>",
        "",
        "| ASI | Title | # rules | Rule IDs |",
        "| --- | --- | --- | --- |",
    ]
    for asi in sorted(ASI_TITLES):
        rules = coverage.get(asi, [])
        title = ASI_TITLES[asi]
        rule_text = ", ".join(f"`{r}`" for r in rules) or "_no rules — CI will fail_"
        lines.append(f"| **{asi}** | {title} | {len(rules)} | {rule_text} |")
    lines.append("")
    return "\n".join(lines)


def render_readme_snippet(coverage: dict[str, list[str]]) -> str:
    """Compact summary suitable for the README marker block."""
    lines = ["| ASI | Title | # rules |", "| --- | --- | --- |"]
    for asi in sorted(ASI_TITLES):
        rules = coverage.get(asi, [])
        lines.append(f"| **{asi}** | {ASI_TITLES[asi]} | {len(rules)} |")
    return "\n".join(lines)


def _read_version() -> str:
    try:
        with PYPROJECT.open("rb") as fh:
            data = tomllib.load(fh)
        return str(data["project"]["version"])
    except (OSError, KeyError):
        return "unknown"


def build_public_json(coverage: dict[str, list[str]]) -> dict:
    from agent_audit_kit.rules.builtin import RULES

    entries = []
    for asi in sorted(ASI_TITLES):
        rule_ids = coverage.get(asi, [])
        rules_out = []
        for rid in rule_ids:
            rule = RULES.get(rid)
            if rule is None:
                continue
            rules_out.append({
                "id": rid,
                "severity": rule.severity.value,
                "cve_references": sorted(rule.cve_references),
                "aicm_references": sorted(rule.aicm_references),
            })
        entries.append({
            "asi_id": asi,
            "title": ASI_TITLES[asi],
            "rule_density": len(rule_ids),
            "rules": rules_out,
        })
    return {
        "schema_version": "1",
        "last_updated": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "aak_version": _read_version(),
        "rule_count": len(RULES),
        "coverage": entries,
    }


def _update_readme_marker(snippet: str) -> bool:
    if not README_PATH.is_file():
        return False
    text = README_PATH.read_text(encoding="utf-8")
    if _MARKER_START not in text or _MARKER_END not in text:
        return False
    before, _, rest = text.partition(_MARKER_START)
    _, _, after = rest.partition(_MARKER_END)
    rebuilt = (
        f"{before}{_MARKER_START}\n{snippet}\n{_MARKER_END}{after}"
    )
    if rebuilt != text:
        README_PATH.write_text(rebuilt, encoding="utf-8")
        return True
    return False


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--json",
        type=pathlib.Path,
        default=DEFAULT_JSON_PATH,
        help="Path to write the machine-readable coverage JSON",
    )
    parser.add_argument(
        "--skip-json",
        action="store_true",
        help="Skip writing the JSON artefact (keeps legacy MD/README behaviour)",
    )
    # Default to [] not sys.argv when invoked as `main()` from tests.
    args = parser.parse_args(argv if argv is not None else [])

    coverage = build_coverage()
    missing = [asi for asi in ASI_TITLES if not coverage.get(asi)]
    text = render_markdown(coverage)
    DOC_PATH.parent.mkdir(parents=True, exist_ok=True)
    DOC_PATH.write_text(text, encoding="utf-8")
    sys.stdout.write(f"wrote {DOC_PATH} ({len(text)} bytes)\n")
    if _update_readme_marker(render_readme_snippet(coverage)):
        sys.stdout.write("updated README.md owasp-coverage marker\n")

    if not args.skip_json:
        payload = build_public_json(coverage)
        args.json.parent.mkdir(parents=True, exist_ok=True)
        args.json.write_text(
            json.dumps(payload, indent=2, sort_keys=False) + "\n",
            encoding="utf-8",
        )
        sys.stdout.write(f"wrote {args.json}\n")

    if missing:
        sys.stderr.write(
            "coverage gap: no rules tag "
            + ", ".join(missing)
            + ". Add at least one rule in agent_audit_kit/rules/builtin.py.\n"
        )
        return 1
    return 0


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