#!/usr/bin/env python3
"""
Query stats for a slice of N source packages from the existing astra.db.

Selects the top N source packages alphabetically from ingested buildinfos,
then counts all artifacts, steps, resources, principals, and edges that
belong to those packages' provenance subgraph.

Usage:
    python3 scripts/query_slice.py --db /data/astra.db --n 3
    python3 scripts/query_slice.py --db /data/astra.db --n 30
    python3 scripts/query_slice.py --db /data/astra.db --n 300
    python3 scripts/query_slice.py --db /data/astra.db --n 3000
"""

import argparse
import os
import re
import sqlite3
import sys

_BUILDINFO_RE = re.compile(r'^artifact:buildinfo:deb/([^@]+)@')


def select_sources(cur, n: int) -> list[str]:
    cur.execute(
        "SELECT DISTINCT astra_id FROM artifacts WHERE kind = 'buildinfo' ORDER BY astra_id"
    )
    seen, sources = set(), []
    for (astra_id,) in cur:
        m = _BUILDINFO_RE.match(astra_id)
        if m:
            name = m.group(1)
            if name not in seen:
                seen.add(name)
                sources.append(name)
            if len(sources) == n:
                break
    return sources


def query_slice(db_path: str, n: int) -> dict:
    con = sqlite3.connect(db_path)
    cur = con.cursor()

    sources = select_sources(cur, n)
    print(f"[slice-{n}] Selected {len(sources)} source packages", file=sys.stderr)
    if sources:
        print(f"  first: {sources[0]}  last: {sources[-1]}", file=sys.stderr)

    # Build IN clause placeholder
    placeholders = ",".join("?" * len(sources))

    # Build step astra_ids for these sources: step:build:deb/<source>@<version>
    # and gitcommit repo last-segment matching source name
    build_step_prefix_conditions = " OR ".join(
        f"astra_id LIKE 'step:build:deb/' || ? || '@%'" for _ in sources
    )

    # --- Artifacts ---
    # buildinfo, tarball for these sources (by step:build prefix match)
    # deb packages: produced by build steps for these sources
    # git-commit / git-file: by repo last-path-segment = source name

    # Buildinfo artifacts
    bi_ids = []
    for src in sources:
        cur.execute(
            "SELECT astra_id FROM artifacts WHERE kind='buildinfo' AND astra_id LIKE ?",
            (f"artifact:buildinfo:deb/{src}@%",),
        )
        bi_ids.extend(r[0] for r in cur.fetchall())

    # Tarball artifacts (via build step → produces → tarball... currently not linked,
    # so count tarballs that appear in buildinfo Checksums via build step edges instead)
    # For now: count tarballs produced by these build steps
    tarball_count = 0
    for src in sources:
        cur.execute(
            """
            SELECT COUNT(DISTINCT e.target) FROM edges e
            JOIN artifacts a ON a.astra_id = e.target AND a.kind = 'tarball'
            WHERE e.source LIKE ? AND e.relation = 'produces'
            """,
            (f"step:build:deb/{src}@%",),
        )
        tarball_count += cur.fetchone()[0]

    # Deb package artifacts produced by build steps for these sources
    deb_count = 0
    for src in sources:
        cur.execute(
            """
            SELECT COUNT(DISTINCT e.target) FROM edges e
            JOIN artifacts a ON a.astra_id = e.target AND a.kind = 'deb'
            WHERE e.source LIKE ? AND e.relation = 'produces'
            """,
            (f"step:build:deb/{src}@%",),
        )
        deb_count += cur.fetchone()[0]

    # Git artifacts: match by repo last path segment = source name
    # gitcommit astra_id: artifact:gitcommit:<host>/<owner>/<repo>@<hash>
    # extract repo name = everything after last '/' before '@'
    git_commit_count = 0
    git_file_count = 0
    for src in sources:
        # repo last segment matches source name exactly
        cur.execute(
            """
            SELECT COUNT(*) FROM artifacts
            WHERE kind = 'git-commit'
              AND astra_id LIKE ? AND astra_id LIKE '%/' || ? || '@%'
            """,
            (f"artifact:gitcommit:%/{src}@%", src),
        )
        git_commit_count += cur.fetchone()[0]
        cur.execute(
            """
            SELECT COUNT(*) FROM artifacts
            WHERE kind = 'git-file'
              AND astra_id LIKE ? AND astra_id LIKE '%/' || ? || '/%'
            """,
            (f"artifact:gitfile:%/{src}/%", src),
        )
        git_file_count += cur.fetchone()[0]

    bi_count = len(bi_ids)
    artifact_count = bi_count + tarball_count + deb_count + git_commit_count + git_file_count

    # --- Steps ---
    step_count = 0
    for src in sources:
        cur.execute(
            "SELECT COUNT(*) FROM steps WHERE astra_id LIKE ?",
            (f"step:build:deb/{src}@%",),
        )
        step_count += cur.fetchone()[0]
        # git steps: step for each gitcommit
        cur.execute(
            "SELECT COUNT(*) FROM steps WHERE astra_id LIKE ? AND astra_id LIKE '%/' || ? || '@%'",
            (f"step:git:%/{src}@%", src),
        )
        step_count += cur.fetchone()[0]

    # --- Resources and Principals ---
    # These are shared (e.g. principal:Debian, archive:debian:...) — count distinct ones
    # referenced in edges that involve our build steps
    resource_ids = set()
    principal_ids = set()
    for src in sources:
        step_pat = f"step:build:deb/{src}@%"
        cur.execute(
            """
            SELECT DISTINCT e.source, e.target FROM edges e
            WHERE e.source LIKE ? OR e.target LIKE ?
            """,
            (step_pat, step_pat),
        )
        for s, t in cur.fetchall():
            for node in (s, t):
                if node.startswith("archive:") or node.startswith("vcs:"):
                    resource_ids.add(node)
                elif node.startswith("principal:"):
                    principal_ids.add(node)

    resource_count = len(resource_ids)
    principal_count = len(principal_ids)

    total_nodes = artifact_count + step_count + resource_count + principal_count

    # --- Edges ---
    edge_count = 0
    for src in sources:
        step_pat = f"step:build:deb/{src}@%"
        cur.execute(
            "SELECT COUNT(*) FROM edges WHERE source LIKE ? OR target LIKE ?",
            (step_pat, step_pat),
        )
        edge_count += cur.fetchone()[0]
        # git edges
        git_pat = f"%/{src}@%"
        cur.execute(
            """
            SELECT COUNT(*) FROM edges
            WHERE (source LIKE 'artifact:gitcommit:%' AND source LIKE ?)
               OR (source LIKE 'step:git:%' AND source LIKE ?)
            """,
            (git_pat, git_pat),
        )
        edge_count += cur.fetchone()[0]

    # --- Complete Evidence % ---
    # Has buildinfo AND gitcommit for the same source
    with_git = sum(1 for src in sources if git_commit_count > 0)
    # More precise: per-source check
    has_git = set()
    for src in sources:
        cur.execute(
            """
            SELECT 1 FROM artifacts
            WHERE kind = 'git-commit'
              AND astra_id LIKE ? AND astra_id LIKE '%/' || ? || '@%'
            LIMIT 1
            """,
            (f"artifact:gitcommit:%/{src}@%", src),
        )
        if cur.fetchone():
            has_git.add(src)

    sources_with_bi = set()
    for src in sources:
        cur.execute(
            "SELECT 1 FROM artifacts WHERE kind='buildinfo' AND astra_id LIKE ? LIMIT 1",
            (f"artifact:buildinfo:deb/{src}@%",),
        )
        if cur.fetchone():
            sources_with_bi.add(src)

    complete = len(sources_with_bi & has_git)
    complete_pct = round(complete * 100.0 / len(sources_with_bi), 1) if sources_with_bi else 0.0

    # Parsed Records = buildinfos + git commits + git files
    parsed_records = bi_count + git_commit_count + git_file_count

    con.close()

    return {
        "n": n,
        "sources": sources,
        "parsed_records": parsed_records,
        "artifacts": artifact_count,
        "by_kind": {
            "buildinfo": bi_count,
            "tarball": tarball_count,
            "deb": deb_count,
            "git-commit": git_commit_count,
            "git-file": git_file_count,
        },
        "steps": step_count,
        "resources": resource_count,
        "principals": principal_count,
        "total_nodes": total_nodes,
        "edges": edge_count,
        "complete_pct": complete_pct,
        "complete_linked": complete,
        "complete_total": len(sources_with_bi),
    }


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:.1f}K"
    return str(n)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--db", required=True)
    ap.add_argument("--n", type=int, required=True, help="number of source packages")
    ap.add_argument("--raw", action="store_true")
    args = ap.parse_args()

    r = query_slice(args.db, args.n)
    fmt = str if args.raw else fmt_num

    print(f"\n=== Slice N={r['n']} ===")
    print(f"Parsed Records : {fmt(r['parsed_records'])}")
    print(f"  (buildinfo={fmt(r['by_kind']['buildinfo'])}, "
          f"git-commit={fmt(r['by_kind']['git-commit'])}, "
          f"git-file={fmt(r['by_kind']['git-file'])})")
    print(f"Artifacts      : {fmt(r['artifacts'])}")
    print(f"  breakdown: " + ", ".join(f"{k}={fmt(v)}" for k, v in 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']} sources with buildinfo+gitcommit)")


if __name__ == "__main__":
    main()