#!/usr/bin/env python3
"""
Ingest exactly N unique source packages (buildinfo + git) into a fresh astra.db.

Iterates through buildinfos alphabetically (one version per source name) and
keeps going until N buildinfo ingestions succeed. Git ingestion is best-effort.

The output DB is written to --out-dir/astra.db (astra uses cwd as data dir).

Usage:
    python3 scripts/ingest_slice.py --n 3 --out-dir /tmp/slice3
    python3 scripts/measure_slice.py --db /tmp/slice3/astra.db --runtime <seconds>
"""

import argparse
import os
import re
import signal
import subprocess
import sys
import time

signal.signal(signal.SIGPIPE, signal.SIG_DFL)

_BINNMU_RE = re.compile(r"\+b\d+$")
_BPO_RE = re.compile(r"~bpo\d+\+\d+$")

_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


# ── buildinfo index ───────────────────────────────────────────────────────────

def build_index(buildinfos_dir: str) -> list[tuple[str, str, str]]:
    """Return sorted list of (source, version, path), one version per source."""
    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
            parts = fname[:-len(".buildinfo")].split("_", 2)
            if len(parts) != 3:
                continue
            source, version, arch = parts
            if arch not in _ACCEPTED_SUFFIXES:
                continue
            prio = 0 if arch 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

    # One version per source name (first alphabetically)
    seen: dict[str, tuple[str, str]] = {}
    for (source, version), path in sorted(index.items()):
        if source not in seen:
            seen[source] = (version, path)

    return [(src, ver, path) for src, (ver, path) in sorted(seen.items())]


# ── tag matching ──────────────────────────────────────────────────────────────

def strip_epoch(v: str) -> str:
    idx = v.find(":")
    return v[idx + 1:] if idx != -1 else v

def strip_binnmu(v: str) -> str:
    return _BINNMU_RE.sub("", v)

def strip_bpo(v: str) -> str:
    return _BPO_RE.sub("", v)

def upstream_from_debian(v: str) -> str | None:
    v = strip_epoch(v)
    return v.rsplit("-", 1)[0] if "-" in v else v

def version_candidates(*versions: str | None) -> list[str]:
    seen, out = set(), []
    for v in versions:
        if not v:
            continue
        v = strip_epoch(v.strip())
        for var in [v, strip_binnmu(v), strip_bpo(v), strip_bpo(strip_binnmu(v))]:
            if var and var not in seen:
                seen.add(var)
                out.append(var)
    return out

def trailing_separator_variants(value: str) -> list[str]:
    v = value.strip().removeprefix("refs/tags/")
    while v.endswith("^{}") or v.endswith("{}"):
        v = v[:-3] if v.endswith("^{}") else v[:-2]
        v = v.strip()
    out = [v] if v else []
    if "/" in v:
        last = v.rsplit("/", 1)[-1]
        if last and last not in out:
            out.append(last)
    for sep in ("+", "~", "_", "-"):
        if sep in v:
            base = v.split(sep, 1)[0]
            if base and base not in out:
                out.append(base)
    return out

def separator_prefix_match(a: str, b: str) -> bool:
    if a == b:
        return True
    if a.startswith(b):
        return len(a) == len(b) or a[len(b)] in "+-._~/"
    if b.startswith(a):
        return len(b) == len(a) or b[len(a)] in "+-._~/"
    return False

def find_matching_tag(tags: list[str], candidates: list[str]) -> str | None:
    for candidate in candidates:
        cvs = trailing_separator_variants(candidate)
        hits = []
        for tag in tags:
            tvs = trailing_separator_variants(tag)
            if any(
                cv == tv or cv in tv or tv in cv or separator_prefix_match(cv, tv)
                for cv in cvs for tv in tvs
            ):
                hits.append(tag)
        if hits:
            hits.sort(key=lambda x: (len(x), x))
            return hits[0]
    return None


# ── git helpers ───────────────────────────────────────────────────────────────

def get_remote_url(clone_path: str) -> str | None:
    r = subprocess.run(
        ["git", "-C", clone_path, "remote", "get-url", "origin"],
        capture_output=True, text=True, timeout=10,
    )
    if r.returncode != 0 or not r.stdout.strip():
        return None
    url = r.stdout.strip()
    from urllib.parse import urlparse, urlunparse
    p = urlparse(url)
    if p.username or p.password:
        url = urlunparse(p._replace(netloc=p.hostname + (f":{p.port}" if p.port else "")))
    return url

def get_remote_tags(clone_path: str) -> list[str]:
    try:
        cp = subprocess.run(
            ["git", "-C", clone_path, "ls-remote", "--tags"],
            capture_output=True, text=True, timeout=120,
        )
    except subprocess.TimeoutExpired:
        print(f"  [timeout] {clone_path}", file=sys.stderr)
        return []
    if cp.returncode != 0:
        return []
    tags = []
    for line in cp.stdout.splitlines():
        parts = line.strip().split()
        if len(parts) == 2 and not parts[1].endswith("^{}"):
            tag = parts[1].removeprefix("refs/tags/")
            if tag:
                tags.append(tag)
    return tags


# ── main ──────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, required=True, help="number of successful buildinfo ingestions")
    ap.add_argument("--out-dir", required=True, help="directory for astra.db (astra uses cwd)")
    ap.add_argument("--buildinfos-dir", default="/data/buildinfos")
    ap.add_argument("--repos-dir", default="/repos")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    # Wipe and reinitialise fresh DB
    if not args.dry_run:
        db_path = os.path.join(args.out_dir, "astra.db")
        if os.path.exists(db_path):
            os.remove(db_path)
            print(f"[slice] removed existing {db_path}", file=sys.stderr)
        r = subprocess.run(["astra", "init"], cwd=args.out_dir, capture_output=True)
        if r.returncode != 0:
            print(f"astra init failed: {r.stderr.decode()}", file=sys.stderr)
            sys.exit(1)
        print(f"[slice] DB initialised at {db_path}", file=sys.stderr)

    print(f"[slice] scanning {args.buildinfos_dir}...", file=sys.stderr)
    all_packages = build_index(args.buildinfos_dir)
    print(f"[slice] {len(all_packages)} unique source packages available", file=sys.stderr)

    ok_bi = fail_bi = ok_git = fail_git = no_clone = no_tag = 0
    ingested: list[str] = []

    start = time.monotonic()

    for source, version, buildinfo_path in all_packages:
        if ok_bi >= args.n:
            break

        # 1. Buildinfo
        cmd = ["astra", "ingest", "debian", "buildinfo", buildinfo_path]
        if args.dry_run:
            print(f"[dry-bi]  {' '.join(cmd)}")
            ok_bi += 1
        else:
            rc = subprocess.run(cmd, capture_output=True, cwd=args.out_dir).returncode
            if rc == 0:
                ok_bi += 1
                ingested.append(source)
            else:
                print(f"  [fail-bi] {source} {version}", file=sys.stderr)
                fail_bi += 1
                continue  # don't attempt git if buildinfo failed

        # 2. Git repo (best-effort)
        clone_path = os.path.join(args.repos_dir, source)
        if not os.path.isdir(clone_path):
            no_clone += 1
            continue

        url = get_remote_url(clone_path)
        if url is None:
            no_clone += 1
            continue

        tags = get_remote_tags(clone_path)
        candidates = version_candidates(version, upstream_from_debian(version))
        tag = find_matching_tag(tags, candidates)
        if tag is None:
            no_tag += 1
            continue

        cmd = ["astra", "ingest", "git", "--tag", tag, url]
        if args.dry_run:
            print(f"[dry-git] {' '.join(cmd)}")
        else:
            rc = subprocess.run(cmd, capture_output=True, cwd=args.out_dir).returncode
            if rc == 0:
                ok_git += 1
            else:
                print(f"  [fail-git] {source} {version} tag={tag}", file=sys.stderr)
                fail_git += 1

    elapsed = time.monotonic() - start
    mins, secs = divmod(elapsed, 60)

    print(f"\nDone: {ok_bi}/{args.n} buildinfos ingested, {fail_bi} failed", file=sys.stderr)
    print(f"git:  {ok_git} ok, {fail_git} failed, {no_clone} no clone, {no_tag} no tag", file=sys.stderr)
    print(f"Run time: {int(mins)}m {secs:.1f}s  ({elapsed:.1f}s)", file=sys.stderr)


if __name__ == "__main__":
    main()
