#!/usr/bin/env python3
"""
Ingest amd64 buildinfos from debtrace into the running astra container.

Looks up each buildinfo file by reconstructing its path from the build_date
in debtrace (YYYY/MM/DD dir) and globbing for the filename, since the exact
suffix (-buildd etc.) varies.

Requires the astra container to be running with buildinfos already mounted:
    docker run -dit --name astra \
        -v /mnt/barbie/abhi/debtrace/data/buildinfos:/data/buildinfos:ro \
        -v /mnt/barbie/abhi/astra:/data \
        --entrypoint /bin/bash astra

Usage:
    # Dry run - see what would be ingested
    python3 scripts/ingest_from_debtrace.py \
        --db /mnt/barbie/abhi/debtrace/debtrace.db \
        --buildinfos-dir /mnt/barbie/abhi/debtrace/data/buildinfos \
        --dry-run

    # Ingest all amd64 buildinfos
    python3 scripts/ingest_from_debtrace.py \
        --db /mnt/barbie/abhi/debtrace/debtrace.db \
        --buildinfos-dir /mnt/barbie/abhi/debtrace/data/buildinfos

    # Ingest one package
    python3 scripts/ingest_from_debtrace.py \
        --db /mnt/barbie/abhi/debtrace/debtrace.db \
        --buildinfos-dir /mnt/barbie/abhi/debtrace/data/buildinfos \
        --source curl
"""

import argparse
import os
import re
import sqlite3
import subprocess
import sys

_PKG_RE = re.compile(r'^artifact:pkg:deb/debian/([^@]+)@([^?]+)\?arch=')

BUILDINFOS = "/data/buildinfos"

# Accepted arch suffixes (after last _ in filename, before .buildinfo).
# amd64 variants preferred over all variants when both exist for same (source, version).
_AMD64_SUFFIXES = {
    "amd64", "amd64-buildd", "amd64+buildd", "amd64+fix",
    "amd64+workaround", "buildd-amd64", "amd64source",
}
_ALL_SUFFIXES = {"all", "all-buildd"}
_ACCEPTED_SUFFIXES = _AMD64_SUFFIXES | _ALL_SUFFIXES


def build_index(buildinfos_dir: str) -> dict[tuple[str, str], str]:
    """Walk the buildinfos tree and map (source, version) -> path.

    Accepts amd64 and all arch variants. When both exist for the same
    (source, version), amd64 takes priority over all.
    """
    index: dict[tuple[str, str], str] = {}
    priority: dict[tuple[str, str], int] = {}
    for root, dirs, files in os.walk(buildinfos_dir):
        dirs.sort()
        for fname in files:
            if not fname.endswith(".buildinfo"):
                continue
            base = fname[:-len(".buildinfo")]
            parts = base.split("_", 2)
            if len(parts) != 3:
                continue
            source, version, arch_suffix = parts
            if arch_suffix not in _ACCEPTED_SUFFIXES:
                continue
            prio = 0 if arch_suffix in _AMD64_SUFFIXES else 1
            key = (source, version)
            if key not in index or prio < priority.get(key, 999):
                index[key] = os.path.join(root, fname)
                priority[key] = prio
    return index


def find_buildinfo(index: dict, source_name: str, version: str) -> str | None:
    ver = version.split(":", 1)[1] if ":" in version else version
    return index.get((source_name, ver))


def query_buildinfos(db_path: str, astra_db: str, source: str | None) -> list[dict]:
    """
    Return (source_name, version) for every buildinfo whose packages are
    present in astra.db (derived from ingested 'deb' artifacts).
    """
    # Step 1: parse (package, version) from astra.db in Python — fast, indexed
    astra_con = sqlite3.connect(astra_db)
    astra_cur = astra_con.cursor()
    astra_cur.execute("SELECT astra_id FROM artifacts WHERE kind = 'deb'")
    packages: list[tuple[str, str]] = []
    for (astra_id,) in astra_cur:
        m = _PKG_RE.match(astra_id)
        if m:
            packages.append((m.group(1), m.group(2)))
    astra_con.close()
    print(f"[astra] {len(packages)} deb artifacts", file=sys.stderr)

    # Step 2: load into a temp table in debtrace so the join uses indexes
    con = sqlite3.connect(db_path)
    con.execute("CREATE TEMP TABLE _pkgs (package TEXT, version TEXT)")
    con.executemany("INSERT INTO _pkgs VALUES (?, ?)", packages)
    con.execute("CREATE INDEX _pkgs_idx ON _pkgs(package, version)")
    con.row_factory = sqlite3.Row
    cur = con.cursor()

    sql = """
        SELECT DISTINCT s.source_name, s.version
        FROM _pkgs p
        JOIN binary_table bt ON bt.package = p.package AND bt.version = p.version
        JOIN checksum_table ct ON ct.binary_id = bt.binary_id
        JOIN buildinfo_table b
            ON b.buildinfo_id = ct.buildinfo_id
            AND b.architecture IN (
                'amd64', 'amd64 source', 'all amd64 source', 'all amd64',
                'all', 'all source', 'source all'
            )
        JOIN source_table s ON s.source_id = b.source_id
    """
    params = []
    if source:
        sql += " WHERE s.source_name = ?"
        params.append(source)
    sql += " ORDER BY s.source_name, s.version"
    cur.execute(sql, params)
    rows = [dict(r) for r in cur.fetchall()]
    con.close()
    return rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--db", default="/debtrace.db", help="path to debtrace.db")
    ap.add_argument("--astra-db", default="/data/astra.db", help="path to astra.db")
    ap.add_argument("--source", help="filter to one source package")
    ap.add_argument("--dry-run", action="store_true", help="print commands without running")
    args = ap.parse_args()

    print("[index] scanning buildinfos directory...", file=sys.stderr)
    index = build_index(BUILDINFOS)
    print(f"[index] {len(index)} buildinfos indexed", file=sys.stderr)

    rows = query_buildinfos(args.db, args.astra_db, args.source)
    print(f"[debtrace] {len(rows)} entries", file=sys.stderr)

    ok = fail = missing = 0
    for row in rows:
        path = find_buildinfo(index, row["source_name"], row["version"])
        if path is None:
            print(f"[missing] {row['source_name']} {row['version']}", file=sys.stderr)
            missing += 1
            continue

        cmd = ["astra", "ingest", "debian", "buildinfo", path]

        if args.dry_run:
            print(" ".join(cmd))
            continue

        rc = subprocess.run(cmd).returncode
        if rc == 0:
            ok += 1
        else:
            print(f"[fail] {path}", file=sys.stderr)
            fail += 1

    if not args.dry_run:
        print(f"\nDone: {ok} ok, {fail} failed, {missing} not found on disk", file=sys.stderr)


if __name__ == "__main__":
    main()