#!/usr/bin/env python3
"""
Measure graph statistics from an astra.db for Table 2.

Reports: Parsed Records, Artifacts, Steps, Resources, Principals,
         Total Nodes, Edges, Complete Evidence %, Graph Size.

Usage:
    python3 scripts/measure_slice.py --db /data/astra.db
    python3 scripts/measure_slice.py --db /tmp/astra_slice3.db --runtime 42.5
"""

import argparse
import os
import sqlite3
import sys


def query_one(cur, sql, params=()):
    cur.execute(sql, params)
    row = cur.fetchone()
    return row[0] if row else 0


def measure(db_path: str) -> dict:
    con = sqlite3.connect(db_path)
    cur = con.cursor()

    results = {}

    # Artifacts (all rows in artifacts table)
    results["artifacts"] = query_one(cur, "SELECT COUNT(*) FROM artifacts")

    # Breakdown by kind (for Parsed Records)
    cur.execute("SELECT kind, COUNT(*) FROM artifacts GROUP BY kind ORDER BY kind")
    results["by_kind"] = dict(cur.fetchall())

    # Steps — try separate table first, fall back to edge prefix
    try:
        results["steps"] = query_one(cur, "SELECT COUNT(*) FROM steps")
    except sqlite3.OperationalError:
        results["steps"] = query_one(
            cur,
            "SELECT COUNT(DISTINCT source) FROM edges WHERE source LIKE 'step:%'"
        )

    # Resources — try separate table first, fall back to edge prefix
    try:
        results["resources"] = query_one(cur, "SELECT COUNT(*) FROM resources")
    except sqlite3.OperationalError:
        results["resources"] = query_one(
            cur,
            """
            SELECT COUNT(DISTINCT val) FROM (
                SELECT source AS val FROM edges WHERE source LIKE 'archive:%'
                UNION
                SELECT target AS val FROM edges WHERE target LIKE 'archive:%'
            )
            """
        )

    # Principals — try separate table first, fall back to edge prefix
    try:
        results["principals"] = query_one(cur, "SELECT COUNT(*) FROM principals")
    except sqlite3.OperationalError:
        results["principals"] = query_one(
            cur,
            """
            SELECT COUNT(DISTINCT val) FROM (
                SELECT source AS val FROM edges WHERE source LIKE 'principal:%'
                UNION
                SELECT target AS val FROM edges WHERE target LIKE 'principal:%'
            )
            """
        )

    results["total_nodes"] = (
        results["artifacts"] + results["steps"] +
        results["resources"] + results["principals"]
    )

    # Edges
    results["edges"] = query_one(cur, "SELECT COUNT(*) FROM edges")

    # Complete Evidence %:
    # A source package has complete evidence if it has BOTH a buildinfo artifact
    # AND at least one gitcommit artifact reachable through the provenance graph.
    #
    # Current graph (before intoto links are ingested):
    #   step:build → buildinfo  (produces)
    #   gitcommit artifacts exist independently
    #
    # After intoto ingest, the chain will be:
    #   gitcommit → step:intoto → tarball → step:build → buildinfo
    #
    # For now: count buildinfos that share a build step with at least one gitcommit,
    # OR that have an intoto link edge connecting them to a gitcommit.
    total_buildinfos = query_one(
        cur, "SELECT COUNT(*) FROM artifacts WHERE kind = 'buildinfo'"
    )

    # Try intoto-linked path (after gen_intoto_links + astra ingest intoto)
    linked_via_intoto = query_one(
        cur,
        """
        SELECT COUNT(DISTINCT e2.target)
        FROM edges e1
        JOIN edges e2
            ON e2.source = e1.source AND e2.relation = 'produces'
        WHERE e1.target LIKE 'artifact:buildinfo:%'
          AND e2.target LIKE 'artifact:gitcommit:%'
        """
    )

    if total_buildinfos > 0:
        results["complete_pct"] = round(linked_via_intoto * 100.0 / total_buildinfos, 1)
    else:
        results["complete_pct"] = 0.0
    results["complete_linked"] = linked_via_intoto
    results["complete_total"] = total_buildinfos

    con.close()

    # Graph size
    results["graph_size_bytes"] = os.path.getsize(db_path)

    # Parsed Records = buildinfos + git objects (commits + files)
    bi = results["by_kind"].get("buildinfo", 0)
    gc = results["by_kind"].get("git-commit", 0)
    gf = results["by_kind"].get("git-file", 0)
    results["parsed_records"] = bi + gc + gf

    return results


def fmt_size(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n:.1f}{unit}"
        n /= 1024
    return f"{n:.1f}TB"


def fmt_num(n: int) -> str:
    if n >= 1_000_000:
        return f"{n/1_000_000:.1f}M"
    if n >= 1_000:
        return f"{n/1_000:.0f}K"
    return str(n)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--db", required=True, help="path to astra.db")
    ap.add_argument("--runtime", type=float, help="run time in seconds (from run_slice.sh)")
    ap.add_argument("--raw", action="store_true", help="print raw numbers instead of formatted")
    args = ap.parse_args()

    if not os.path.exists(args.db):
        print(f"error: {args.db} not found", file=sys.stderr)
        sys.exit(1)

    r = measure(args.db)

    fmt = str if args.raw else fmt_num
    size_fmt = str if args.raw else fmt_size

    print("\n=== Table 2 metrics ===")
    print(f"Parsed Records : {fmt(r['parsed_records'])}")
    print(f"  (buildinfo={fmt(r['by_kind'].get('buildinfo',0))}, "
          f"git-commit={fmt(r['by_kind'].get('git-commit',0))}, "
          f"git-file={fmt(r['by_kind'].get('git-file',0))})")
    print(f"Artifacts      : {fmt(r['artifacts'])}")
    print(f"  breakdown: {', '.join(f'{k}={fmt(v)}' for k,v in sorted(r['by_kind'].items()))}")
    print(f"Steps          : {fmt(r['steps'])}")
    print(f"Resources      : {fmt(r['resources'])}")
    print(f"Principals     : {fmt(r['principals'])}")
    print(f"Total Nodes    : {fmt(r['total_nodes'])}")
    print(f"Edges          : {fmt(r['edges'])}")
    print(f"Complete Ev.%  : {r['complete_pct']}%  ({r['complete_linked']}/{r['complete_total']} buildinfos linked to gitcommit)")
    print(f"Graph Size     : {size_fmt(r['graph_size_bytes'])}")
    if args.runtime is not None:
        mins, secs = divmod(args.runtime, 60)
        print(f"Run Time       : {int(mins)}m {secs:.1f}s")
    else:
        print(f"Run Time       : (pass --runtime <seconds>)")
    print()


if __name__ == "__main__":
    main()