#!/usr/bin/env python3
"""
gen-claims.py — Generate provable marketing claims directly from the product.

This is the single source of truth for every quantitative marketing claim
DonjonSec makes (CVE count, scanner count, compliance-framework count,
AI-backend count, and the per-framework control-count table).

Design principle: numbers are GENERATED FROM THE PRODUCT, never hand-typed.
The website and README consume the emitted register; a CI drift-check
(drift-check.py) fails the build if the live site ever diverges from what
this script produces.

Air-gap safe: stdlib only (plus PyYAML, already a product dependency).
No network calls. Reads the same artifacts the running product reads.

Usage:
    python scripts/gen-claims.py                 # human summary to stdout
    python scripts/gen-claims.py --json          # raw JSON to stdout
    python scripts/gen-claims.py --write          # update docs/claims-register.json + .md
    python scripts/gen-claims.py --check          # exit 1 if register is stale vs product

Exit codes:
    0  success / register in sync (with --check)
    1  register stale vs product (with --check) or a source artifact was missing
"""

from __future__ import annotations

import argparse
import glob
import json
import os
import sqlite3
import sys
from datetime import datetime, timezone
from pathlib import Path

try:
    import yaml
except ImportError:  # pragma: no cover - product always ships PyYAML
    print("FATAL: PyYAML is required (it is a product dependency).", file=sys.stderr)
    sys.exit(1)

# --------------------------------------------------------------------------
# Paths — resolved relative to the repo root so this runs from anywhere.
# --------------------------------------------------------------------------
SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SCRIPT_DIR.parent
VULN_DB = REPO_ROOT / "data" / "vuln_db" / "vuln_intel.db"
FRAMEWORKS_DIR = REPO_ROOT / "config" / "frameworks"
LIB_DIR = REPO_ROOT / "lib"
DOCS_DIR = REPO_ROOT / "docs"
REGISTER_JSON = DOCS_DIR / "claims-register.json"
REGISTER_MD = DOCS_DIR / "claims-register.md"

SCHEMA_VERSION = 1


# --------------------------------------------------------------------------
# Framework control counting
# --------------------------------------------------------------------------
# Each shipped framework YAML uses its own hierarchy, but every one of them
# expresses controls as leaf dicts (carrying a name/description/etc.) reached
# through "container" keys, optionally nested under "grouping" keys. This
# traversal is structural, not per-framework hardcoded, so adding a framework
# YAML automatically updates the count with no code change.

LEAF_CONTAINER_KEYS = {
    "controls", "requirements", "safeguards", "practices", "articles",
    "consumer_rights", "controller_obligations", "processor_obligations",
    "sub_controls", "subrequirements", "provisions", "clauses", "criteria",
    "objectives", "subcategories",
}
GROUP_KEYS = {
    "families", "themes", "functions", "categories", "domains", "levels",
    "chapters", "sections", "types", "sub_categories", "impact_levels",
    "principles",
}
TEXT_FIELDS = {
    "name", "title", "description", "statement", "text", "requirement",
    "objective", "guidance", "intent",
}
SKIP_KEYS = {
    "sources", "summary", "metadata", "policy_mapping", "related_standards",
    "impact_levels",
}


def _entries(value):
    if isinstance(value, dict):
        return list(value.values())
    if isinstance(value, list):
        return list(value)
    return []


def _count_leaves(node) -> int:
    """Count control leaves rooted at a single candidate node."""
    if not isinstance(node, dict):
        return 0
    sub = 0
    for key, value in node.items():
        if key in LEAF_CONTAINER_KEYS and isinstance(value, (dict, list)):
            for entry in _entries(value):
                sub += _count_leaves(entry)
    if sub > 0:
        return sub  # grouping node — its leaves are counted above
    # Leaf only if it actually carries descriptive control text.
    return 1 if any(f in node for f in TEXT_FIELDS) else 0


def count_controls(doc) -> int:
    """Count control leaves across an entire framework document."""
    total = 0

    def descend(obj):
        nonlocal total
        if isinstance(obj, dict):
            for key, value in obj.items():
                if key in SKIP_KEYS:
                    continue
                if key in LEAF_CONTAINER_KEYS and isinstance(value, (dict, list)):
                    for entry in _entries(value):
                        total += _count_leaves(entry)
                elif key in GROUP_KEYS and isinstance(value, (dict, list)):
                    for entry in _entries(value):
                        descend(entry)
                elif isinstance(value, (dict, list)):
                    descend(value)
        elif isinstance(obj, list):
            for entry in obj:
                descend(entry)

    descend(doc)
    return total


# --------------------------------------------------------------------------
# Product introspection
# --------------------------------------------------------------------------
def get_cve_count() -> int:
    if not VULN_DB.exists():
        raise FileNotFoundError(f"vuln intel DB not found: {VULN_DB}")
    conn = sqlite3.connect(f"file:{VULN_DB}?mode=ro", uri=True)
    try:
        cur = conn.cursor()
        cur.execute("SELECT COUNT(*) FROM nvd_cves")
        return int(cur.fetchone()[0])
    finally:
        conn.close()


def get_scanner_count():
    """Read the orchestrator's scanner registry — the product's own source of truth.

    Parsed statically with ast (no import / no execution) so the orchestrator's
    heavy scanner imports never load. Adding a scanner row to _SCANNER_REGISTRY
    automatically updates the count with no change here. Air-gap safe.
    """
    import ast

    src = (LIB_DIR / "orchestrator.py").read_text(encoding="utf-8")
    tree = ast.parse(src, filename="orchestrator.py")
    registry_node = None
    for node in tree.body:
        if isinstance(node, ast.Assign):
            for tgt in node.targets:
                if isinstance(tgt, ast.Name) and tgt.id == "_SCANNER_REGISTRY":
                    registry_node = node.value
    if registry_node is None or not isinstance(registry_node, ast.List):
        raise ValueError("_SCANNER_REGISTRY list not found in orchestrator.py")

    ids = []
    for elt in registry_node.elts:
        # Each row is a tuple; the first element is the scanner id string const.
        if isinstance(elt, ast.Tuple) and elt.elts:
            first = elt.elts[0]
            if isinstance(first, ast.Constant) and isinstance(first.value, str):
                ids.append(first.value)
    return len(ids), ids


def get_ai_backend_count():
    """Count AI backends from the product's own AIEngine.get_supported_backends().

    ai_engine.py is import-safe standalone (relative/absolute import fallback,
    no network I/O at import time), so we load it in isolation and call the
    real method — the authoritative source of truth.
    """
    if str(LIB_DIR) not in sys.path:
        sys.path.insert(0, str(LIB_DIR))
    import importlib.util

    spec = importlib.util.spec_from_file_location(
        "_ai_engine_for_claims", LIB_DIR / "ai_engine.py"
    )
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    backends = mod.AIEngine.get_supported_backends()
    ids = [b["id"] for b in backends]
    return len(ids), ids


def get_framework_table():
    files = sorted(glob.glob(str(FRAMEWORKS_DIR / "*.yaml")) +
                   glob.glob(str(FRAMEWORKS_DIR / "*.yml")))
    table = []
    for path in files:
        with open(path, encoding="utf-8") as fh:
            doc = yaml.safe_load(fh)
        fw = doc.get("framework", {}) if isinstance(doc, dict) else {}
        table.append({
            "file": os.path.basename(path),
            "id": fw.get("id") if isinstance(fw, dict) else None,
            "name": fw.get("name") if isinstance(fw, dict) else None,
            "version": fw.get("version") if isinstance(fw, dict) else None,
            "controls": count_controls(doc),
        })
    return table


# --------------------------------------------------------------------------
# Claim assembly
# --------------------------------------------------------------------------
def build_claims() -> dict:
    cve_count = get_cve_count()
    scanner_count, scanner_ids = get_scanner_count()
    backend_count, backend_ids = get_ai_backend_count()
    fw_table = get_framework_table()
    framework_count = len(fw_table)
    total_controls = sum(row["controls"] for row in fw_table)

    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

    return {
        "schema_version": SCHEMA_VERSION,
        "generated_at": now,
        "generator": "scripts/gen-claims.py",
        "note": (
            "Every number here is generated from the product. Do not hand-edit. "
            "Run scripts/gen-claims.py --write to regenerate."
        ),
        "provable": {
            "cve_signatures": {
                "value": cve_count,
                "source": "data/vuln_db/vuln_intel.db -> SELECT COUNT(*) FROM nvd_cves",
                "marketing_string": f"{cve_count:,}",
                "rounded_marketing_string": f"{cve_count // 1000}K+",
                "provable": True,
            },
            "scanner_count": {
                "value": scanner_count,
                "ids": scanner_ids,
                "source": "lib/orchestrator.py::_SCANNER_REGISTRY",
                "marketing_string": f"{scanner_count} scanners",
                "provable": True,
            },
            "framework_count": {
                "value": framework_count,
                "source": "config/frameworks/*.yaml (file count)",
                "marketing_string": f"{framework_count} compliance frameworks",
                "provable": True,
            },
            "ai_backend_count": {
                "value": backend_count,
                "ids": backend_ids,
                "source": "lib/ai_engine.py::AIEngine.get_supported_backends()",
                "marketing_string": f"{backend_count} AI backends",
                "provable": True,
            },
            "total_controls": {
                "value": total_controls,
                "source": "sum of per-framework control leaves in config/frameworks/*.yaml",
                "marketing_string": f"{total_controls:,} controls",
                "provable": True,
            },
        },
        "framework_table": fw_table,
        # ------------------------------------------------------------------
        # Competitor pricing: NOT generated from the product. These are
        # external, third-party list prices that DonjonSec cannot prove from
        # its own artifacts. They are recorded here so drift-check can detect
        # site/README divergence, but each is explicitly marked unprovable
        # and pending a cited source. See docs/claims-register.md for the
        # reconciliation and the canonical proposed set.
        # ------------------------------------------------------------------
        "competitor_pricing": {
            "provable": False,
            "status": "unsourced — needs a cited source or removal",
            "headline_floor": {
                "value": 110000,
                "marketing_string": "$110K+/yr",
                "basis": (
                    "Documented CONSERVATIVE FLOOR. Sum of the lowest defensible "
                    "annual list price for one representative tool in each replaced "
                    "category (VM + compliance automation + risk quantification). "
                    "Pinned as a floor, not a precise total."
                ),
                "provable": False,
            },
            "canonical_proposed": {
                "note": (
                    "ONE canonical set proposed to replace the site-vs-README "
                    "disagreement. Each entry REQUIRES a cited public source "
                    "(vendor quote, G2/Gartner, public RFP) before it may ship. "
                    "Until sourced, render as 'starting at' ranges or remove."
                ),
                "vendors": [
                    {"vendor": "Tenable", "category": "Vulnerability management",
                     "proposed_floor_usd": 35000, "provable": False,
                     "needs": "cited public list price"},
                    {"vendor": "Qualys", "category": "Vulnerability management",
                     "proposed_floor_usd": 25000, "provable": False,
                     "needs": "cited public list price"},
                    {"vendor": "Rapid7", "category": "Vulnerability management",
                     "proposed_floor_usd": 20000, "provable": False,
                     "needs": "cited public list price"},
                    {"vendor": "RiskLens", "category": "FAIR risk quantification",
                     "proposed_floor_usd": 50000, "provable": False,
                     "needs": "cited public list price"},
                    {"vendor": "Drata", "category": "Compliance automation",
                     "proposed_floor_usd": 30000, "provable": False,
                     "needs": "cited public list price"},
                ],
            },
        },
    }


# --------------------------------------------------------------------------
# Rendering
# --------------------------------------------------------------------------
def render_markdown(claims: dict) -> str:
    p = claims["provable"]
    lines = []
    lines.append("# DonjonSec Claims Register")
    lines.append("")
    lines.append(f"> Generated by `{claims['generator']}` at {claims['generated_at']}.")
    lines.append("> **Do not hand-edit.** Run `python scripts/gen-claims.py --write` to regenerate.")
    lines.append("> Drift between this register and the live site is a CI failure (`drift-check.py`).")
    lines.append("")
    lines.append("## Provable claims (generated from the product)")
    lines.append("")
    lines.append("| Claim | Value | Marketing string | Source / derivation | Provable |")
    lines.append("|-------|-------|------------------|---------------------|----------|")
    lines.append(f"| CVE signatures | {p['cve_signatures']['value']:,} | "
                 f"`{p['cve_signatures']['marketing_string']}` "
                 f"(rounded `{p['cve_signatures']['rounded_marketing_string']}`) | "
                 f"{p['cve_signatures']['source']} | yes |")
    lines.append(f"| Scanners | {p['scanner_count']['value']} | "
                 f"`{p['scanner_count']['marketing_string']}` | "
                 f"{p['scanner_count']['source']} | yes |")
    lines.append(f"| Compliance frameworks | {p['framework_count']['value']} | "
                 f"`{p['framework_count']['marketing_string']}` | "
                 f"{p['framework_count']['source']} | yes |")
    lines.append(f"| AI backends | {p['ai_backend_count']['value']} | "
                 f"`{p['ai_backend_count']['marketing_string']}` | "
                 f"{p['ai_backend_count']['source']} | yes |")
    lines.append(f"| Total controls | {p['total_controls']['value']:,} | "
                 f"`{p['total_controls']['marketing_string']}` | "
                 f"{p['total_controls']['source']} | yes |")
    lines.append("")
    lines.append("## Per-framework control counts (auditor-checkable)")
    lines.append("")
    lines.append("| Framework | id | Version | Controls |")
    lines.append("|-----------|----|---------|----------|")
    for row in claims["framework_table"]:
        lines.append(f"| {row['name'] or row['file']} | `{row['id']}` | "
                     f"{row['version'] or '-'} | {row['controls']} |")
    lines.append(f"| **TOTAL** | | | **{p['total_controls']['value']:,}** |")
    lines.append("")
    cp = claims["competitor_pricing"]
    lines.append("## Competitor pricing (UNPROVABLE — needs sourcing)")
    lines.append("")
    lines.append(f"**Status:** {cp['status']}")
    lines.append("")
    lines.append(f"**Headline floor:** `{cp['headline_floor']['marketing_string']}` — "
                 f"{cp['headline_floor']['basis']}")
    lines.append("")
    lines.append("Proposed ONE canonical set (replaces the site-vs-README disagreement). "
                 "Every entry must carry a cited public source before it ships:")
    lines.append("")
    lines.append("| Vendor | Category | Proposed floor (USD/yr) | Provable | Needs |")
    lines.append("|--------|----------|-------------------------|----------|-------|")
    for v in cp["canonical_proposed"]["vendors"]:
        lines.append(f"| {v['vendor']} | {v['category']} | "
                     f"${v['proposed_floor_usd']:,} | no | {v['needs']} |")
    lines.append("")
    lines.append("> Note: the sum of the proposed floors is greater than the "
                 "`$110K+` headline. The headline is intentionally pinned as a "
                 "conservative floor, NOT the sum, and must be presented that way.")
    lines.append("")
    return "\n".join(lines)


def render_human(claims: dict) -> str:
    p = claims["provable"]
    out = []
    out.append("DonjonSec provable claims (generated from product):")
    out.append(f"  CVE signatures      : {p['cve_signatures']['value']:,}  "
               f"(marketing: {p['cve_signatures']['rounded_marketing_string']})")
    out.append(f"  Scanners            : {p['scanner_count']['value']}")
    out.append(f"  Compliance frameworks: {p['framework_count']['value']}")
    out.append(f"  AI backends         : {p['ai_backend_count']['value']}")
    out.append(f"  Total controls      : {p['total_controls']['value']:,}")
    out.append("")
    out.append("Per-framework control counts:")
    for row in claims["framework_table"]:
        out.append(f"  {row['id'] or row['file']:24} {row['controls']:>5}")
    out.append("")
    out.append("Competitor pricing: UNPROVABLE — needs cited sources (see register).")
    return "\n".join(out)


# --------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------
def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--json", action="store_true", help="emit raw JSON")
    parser.add_argument("--write", action="store_true",
                        help="write docs/claims-register.json and .md")
    parser.add_argument("--check", action="store_true",
                        help="exit 1 if the written register is stale vs the product")
    args = parser.parse_args(argv)

    claims = build_claims()

    if args.check:
        if not REGISTER_JSON.exists():
            print("STALE: claims-register.json does not exist. Run --write.", file=sys.stderr)
            return 1
        existing = json.loads(REGISTER_JSON.read_text(encoding="utf-8"))
        # Compare everything except the timestamp.
        a = dict(claims); a.pop("generated_at", None)
        b = dict(existing); b.pop("generated_at", None)
        if a != b:
            print("STALE: claims-register.json differs from current product values.",
                  file=sys.stderr)
            print("Run: python scripts/gen-claims.py --write", file=sys.stderr)
            return 1
        print("OK: claims-register.json is in sync with the product.")
        return 0

    if args.write:
        DOCS_DIR.mkdir(parents=True, exist_ok=True)
        REGISTER_JSON.write_text(
            json.dumps(claims, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
        )
        REGISTER_MD.write_text(render_markdown(claims), encoding="utf-8")
        print(f"Wrote {REGISTER_JSON}")
        print(f"Wrote {REGISTER_MD}")
        return 0

    if args.json:
        print(json.dumps(claims, indent=2, ensure_ascii=False))
        return 0

    print(render_human(claims))
    return 0


if __name__ == "__main__":
    sys.exit(main())
