#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = []
# ///
"""
lake_build_ingest.py — Run `lake build` and ingest the result into ENE on neon-64gb.

This is a pure-I/O shim: it runs the build, parses the output, and writes rows to
Postgres. It contains no admissibility logic and no Float arithmetic.

Target database: arxiv-pg container on neon-64gb (Tailscale 100.92.88.64).
Tables used:
  ene.sessions     — one row per build invocation
  ene.packages     — one row per build target/receipt
  ene.receipts     — verification receipt bound to the package
  ene.ingest_events — provenance for the ingestion itself

Usage:
  python3 4-Infrastructure/shim/lake_build_ingest.py [target] [--dir LEAN_DIR] [--actor NAME]

Examples:
  python3 4-Infrastructure/shim/lake_build_ingest.py
  python3 4-Infrastructure/shim/lake_build_ingest.py Compiler
  python3 4-Infrastructure/shim/lake_build_ingest.py Semantics.NKHodgeFAMM --actor opencode
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

ROOT = Path("/home/allaun/Research Stack")
LEAN_DIR = ROOT / "0-Core-Formalism/lean/Semantics"
NEON_HOST = "100.92.88.64"
CONTAINER = "arxiv-pg"
DB = "ene"


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _run(cmd: list[str], cwd: Path, timeout: float = 3600.0) -> subprocess.CompletedProcess:
    return subprocess.run(
        cmd,
        cwd=cwd,
        capture_output=True,
        text=True,
        timeout=timeout,
    )


def _ssh_psql(sql: str, timeout: int = 60) -> str:
    """Run a psql command on neon-64gb's arxiv-pg container."""
    escaped = sql.replace('"', '\\"')
    result = subprocess.run(
        [
            "ssh",
            f"allaun@{NEON_HOST}",
            f'podman exec {CONTAINER} psql -U postgres -d {DB} -c "{escaped}"',
        ],
        capture_output=True,
        text=True,
        timeout=timeout,
    )
    if result.returncode != 0:
        raise RuntimeError(f"psql failed: {result.stderr[:500]}")
    return result.stdout


def _ssh_psql_rows(sql: str, timeout: int = 60) -> list[list[str]]:
    """Run a psql query on neon-64gb and return rows."""
    escaped = sql.replace('"', '\\"')
    result = subprocess.run(
        [
            "ssh",
            f"allaun@{NEON_HOST}",
            f'podman exec {CONTAINER} psql -U postgres -d {DB} -t -A -F "|" -c "{escaped}"',
        ],
        capture_output=True,
        text=True,
        timeout=timeout,
    )
    if result.returncode != 0:
        raise RuntimeError(f"psql failed: {result.stderr[:500]}")
    rows = [line.split("|") for line in result.stdout.strip().split("\n") if line]
    return rows


def _git_info() -> dict[str, str]:
    result = _run(["git", "rev-parse", "HEAD"], ROOT)
    commit = result.stdout.strip() if result.returncode == 0 else ""
    result = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"], ROOT)
    branch = result.stdout.strip() if result.returncode == 0 else ""
    result = _run(["git", "status", "--short"], ROOT)
    dirty = bool(result.stdout.strip())
    return {
        "commit": commit,
        "branch": branch,
        "dirty": str(dirty),
    }


def _parse_build_output(stdout: str, stderr: str) -> dict[str, Any]:
    """Extract job count, error count, and warnings from lake build output."""
    combined = stdout + "\n" + stderr
    total: int | None = None
    errors: int | None = None

    # Match "[N/M] Built Module" and final "Built X (Ys)" or "Build completed"
    built_re = re.compile(r"\[(\d+)/(\d+)\] Built")
    for m in built_re.finditer(combined):
        total = int(m.group(2))

    # Lake final summary: "Build completed successfully (N jobs, M errors)"
    summary_re = re.compile(r"Build completed.*?\((\d+) jobs?.*?\)", re.IGNORECASE)
    m = summary_re.search(combined)
    if m:
        total = int(m.group(1))

    # Error count from "X errors" or error markers
    error_re = re.compile(r"(\d+) error", re.IGNORECASE)
    m = error_re.search(combined)
    if m:
        errors = int(m.group(1))

    # Warnings
    warnings = len(re.findall(r"warning:", combined, re.IGNORECASE))

    # Determine status
    if "Build completed successfully" in combined or "Built" in combined:
        status = "success" if not errors else "warnings"
    else:
        status = "failed"

    return {
        "total": total,
        "errors": errors if errors is not None else (0 if status != "failed" else None),
        "warnings": warnings,
        "status": status,
        "raw_output": combined[-20000:],  # cap size
    }


def _insert_session(target: str, git: dict[str, str], build: dict[str, Any], elapsed_ms: int, actor: str) -> str:
    """Insert a row into ene.sessions and return the session id."""
    session_id = hashlib.sha256(f"{git['commit']}:{target}:{_now()}".encode()).hexdigest()[:32]
    title = f"lake build {target}" if target else "lake build"
    summary = json.dumps(
        {
            "target": target,
            "jobs_total": build["total"],
            "jobs_failed": build["errors"],
            "warnings": build["warnings"],
            "status": build["status"],
        },
        separators=(",", ":"),
    )
    sql = f"""
    INSERT INTO ene.sessions
        (id, project_id, title, started_at, ended_at, actor, toolchain,
         summary, memory_hash, promotion_state, build_target, jobs_total,
         jobs_failed, elapsed_ms, git_commit, git_branch)
    VALUES
        ('{session_id}', 'research-stack', '{title}', '{_now()}', '{_now()}',
         '{actor}', 'lake build',
         '{summary}'::jsonb, '', 'held', '{target}',
         {build['total'] if build['total'] is not None else 'NULL'},
         {build['errors'] if build['errors'] is not None else 'NULL'},
         {elapsed_ms}, '{git['commit']}', '{git['branch']}')
    ON CONFLICT (id) DO NOTHING;
    """
    _ssh_psql(sql)
    return session_id


def _insert_package(session_id: str, target: str, git: dict[str, str], build: dict[str, Any]) -> str:
    """Insert a package row representing the build artifact and return its id.

    Build logs get content_hash = NULL so the partial unique index
    ene_pkg_hash_idx (WHERE content_hash IS NOT NULL) does not fire.
    Each build session therefore gets its own package row."""
    pkg_id = hashlib.sha256(f"lake-build:{session_id}:{git['commit']}:{target}".encode()).hexdigest()[:32]
    content = build["raw_output"]
    title = f"lake build {target}" if target else "lake build"
    sql = f"""
    INSERT INTO ene.packages
        (pkg, package_type, title, content, content_hash, element,
         tags, source, created_at, updated_at, session_id, provenance,
         verification_status, promotion_state, domain, archetype, notes)
    VALUES
        ('{pkg_id}', 'build_log', '{title}', '{_escape_sql(content)}',
         NULL, 'METAL',
         '["lake", "build", "lean"]'::jsonb,
         '4-Infrastructure/shim/lake_build_ingest.py',
         '{_now()}', '{_now()}', '{session_id}',
         '{{"git_commit": "{git['commit']}", "git_branch": "{git['branch']}", "dirty": {str(git['dirty']).lower()}}}'::jsonb,
         '{build['status']}', 'held', 'core-formalism', 'build_artifact',
         'Build output captured from lake build {target}');
    """
    _ssh_psql(sql)
    return pkg_id


def _insert_receipt(pkg_id: str, session_id: str, build: dict[str, Any]) -> None:
    """Insert a verification receipt for the build package."""
    receipt_hash = hashlib.sha256(build["raw_output"].encode()).hexdigest()
    input_hash = hashlib.sha256(build["raw_output"].encode()).hexdigest()
    residual = json.dumps(
        {"warnings": build["warnings"], "errors": build["errors"]},
        separators=(",", ":"),
    )
    sql = f"""
    INSERT INTO ene.receipts
        (package_id, receipt_type, receipt_hash, input_hash, output_hash,
         toolchain, verifier, status, residual)
    VALUES
        ('{pkg_id}', 'lake_build', '{receipt_hash}', '{input_hash}',
         '{receipt_hash}', 'lake build', 'opencode', '{build['status']}',
         '{residual}'::jsonb)
    ON CONFLICT (package_id, receipt_type, receipt_hash) DO NOTHING;
    """
    _ssh_psql(sql)


def _insert_ingest_event(pkg_id: str, session_id: str) -> None:
    """Record provenance for the ingestion itself."""
    sql = f"""
    INSERT INTO ene.ingest_events
        (package_id, source_uri, source_type, source_hash, ingest_profile,
         parser_version, extracted_entities, extracted_equations,
         extracted_links, extracted_bibtex, warnings)
    VALUES
        ('{pkg_id}', 'file://4-Infrastructure/shim/lake_build_ingest.py',
         'build_log_ingest', '', 'lake_build_ingest_v1', '1.0.0',
         '{{"session_id": "{session_id}"}}'::jsonb, '[]'::jsonb,
         '[]'::jsonb, '[]'::jsonb, '[]'::jsonb)
    ON CONFLICT DO NOTHING;
    """
    _ssh_psql(sql)


def _escape_sql(s: str) -> str:
    return s.replace("'", "''").replace("\\", "\\\\")


def main() -> int:
    parser = argparse.ArgumentParser(description="Ingest lake build output into ENE on neon-64gb")
    parser.add_argument("target", nargs="?", default="", help="lake build target (default: full build)")
    parser.add_argument("--dir", type=Path, default=LEAN_DIR, help="directory to run lake build in")
    parser.add_argument("--actor", default="opencode", help="actor name for session")
    parser.add_argument("--dry-run", action="store_true", help="print parsed build info without writing")
    args = parser.parse_args()

    actor = args.actor
    target = args.target
    cmd = ["lake", "build"]
    if target:
        cmd.append(target)

    print(f"Running: {' '.join(cmd)} in {args.dir}", file=sys.stderr)
    start = time.time()
    proc = _run(cmd, args.dir)
    elapsed_ms = int((time.time() - start) * 1000)

    git = _git_info()
    build = _parse_build_output(proc.stdout, proc.stderr)

    print(
        f"Build status: {build['status']} | jobs: {build['total']} | "
        f"errors: {build['errors']} | warnings: {build['warnings']} | "
        f"elapsed: {elapsed_ms}ms",
        file=sys.stderr,
    )

    if args.dry_run:
        print(json.dumps({"git": git, "build": build, "elapsed_ms": elapsed_ms}, indent=2))
        return 0

    session_id = _insert_session(target, git, build, elapsed_ms, actor)
    pkg_id = _insert_package(session_id, target, git, build)
    _insert_receipt(pkg_id, session_id, build)
    _insert_ingest_event(pkg_id, session_id)

    print(f"Ingested: session={session_id} package={pkg_id}", file=sys.stderr)
    return 0 if build["status"] != "failed" else 1


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