#!/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] [--rrc TRACE_JSON]

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
  python3 4-Infrastructure/shim/lake_build_ingest.py Compiler --rrc Compiler.trace.json

When --rrc is given, the shim looks for <target>.trace.json (or the explicit path),
runs the canonical `pist-classify-trace` executable, then runs `rrc-watchdog` with
its predicted shape. The resulting RRC alignment receipt is stored alongside the
build receipt. Both steps are pure I/O wrappers around Lean executables.
"""
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"
PIST_CLASSIFY = LEAN_DIR / ".lake" / "build" / "bin" / "pist-classify-trace"
RRC_WATCHDOG = LEAN_DIR / ".lake" / "build" / "bin" / "rrc-watchdog"


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 _find_rrc_trace(target: str, explicit_path: Path | None, lean_dir: Path) -> Path | None:
    """Resolve the RRC trace fixture path.

    If an explicit path is given, use it. Otherwise look for <target>.trace.json
    in the Lean directory. Returns None if not found.
    """
    if explicit_path is not None:
        if explicit_path.is_absolute():
            return explicit_path if explicit_path.exists() else None
        candidate = lean_dir / explicit_path
        if candidate.exists():
            return candidate
        candidate = ROOT / explicit_path
        return candidate if candidate.exists() else None
    if not target:
        return None
    candidate = lean_dir / f"{target}.trace.json"
    return candidate if candidate.exists() else None


def _run_pist_classify(trace_path: Path, lean_dir: Path) -> dict[str, Any]:
    """Run the canonical pist-classify-trace executable on a trace fixture."""
    if not PIST_CLASSIFY.exists():
        raise RuntimeError(f"pist-classify-trace not found at {PIST_CLASSIFY}")
    proc = subprocess.run(
        [str(PIST_CLASSIFY), str(trace_path)],
        cwd=lean_dir,
        capture_output=True,
        text=True,
        timeout=60,
    )
    if proc.returncode != 0:
        raise RuntimeError(
            f"pist-classify-trace failed: {proc.stderr[:500]}"
        )
    return json.loads(proc.stdout)


def _run_rrc_watchdog(pist_label: str, exact_label: str | None, rrc_shape: str, lean_dir: Path) -> dict[str, Any]:
    """Run the canonical rrc-watchdog executable for RRC alignment.

    The executable returns JSON and exits 0 when score >= 86.
    We capture exit code as part of the result but do not treat non-zero as failure.
    """
    if not RRC_WATCHDOG.exists():
        raise RuntimeError(f"rrc-watchdog not found at {RRC_WATCHDOG}")
    cmd = [str(RRC_WATCHDOG), "--pist-label", pist_label]
    if exact_label:
        cmd.extend(["--exact-label", exact_label])
    cmd.extend(["--rrc-shape", rrc_shape])
    proc = subprocess.run(
        cmd,
        cwd=lean_dir,
        capture_output=True,
        text=True,
        timeout=60,
    )
    try:
        result = json.loads(proc.stdout)
    except json.JSONDecodeError as e:
        raise RuntimeError(
            f"rrc-watchdog produced non-JSON output: {proc.stdout[:500]}"
        ) from e
    result["exit_code"] = proc.returncode
    return result


def _run_rrc_alignment(target: str, trace_path: Path, lean_dir: Path) -> dict[str, Any] | None:
    """Run PIST classification + RRC alignment gate on a trace fixture.

    Returns a dict with the PIST output, RRC alignment output, and computed
    receipt hash. Returns None if any step cannot be completed.
    """
    try:
        pist = _run_pist_classify(trace_path, lean_dir)
    except Exception as e:
        print(f"RRC PIST classification skipped: {e}", file=sys.stderr)
        return None

    predicted_shape = pist.get("predicted_rrc_shape", "HoldForUnlawfulOrUnderspecifiedShape")
    tactic_family = pist.get("tactic_family", "unknown")

    # rrc-watchdog expects a shape argument. Use the predicted shape as both
    # the PIST label and the declared RRC shape, so the gate tests self-alignment.
    try:
        rrc = _run_rrc_watchdog(
            predicted_shape,
            None,
            _shape_to_watchdog_arg(predicted_shape),
            lean_dir,
        )
    except Exception as e:
        print(f"RRC alignment gate skipped: {e}", file=sys.stderr)
        return None

    canonical = json.dumps(
        {"pist": pist, "rrc": rrc},
        sort_keys=True,
        separators=(",", ":"),
    )
    return {
        "pist": pist,
        "rrc": rrc,
        "predicted_shape": predicted_shape,
        "tactic_family": tactic_family,
        "canonical_json": canonical,
        "receipt_hash": hashlib.sha256(canonical.encode()).hexdigest(),
    }


def _shape_to_watchdog_arg(shape: str) -> str:
    """Map pist-classify-trace shape strings to rrc-watchdog CLI arguments."""
    mapping = {
        "SignalShapedRouteCompiler": "signalShapedRouteCompiler",
        "ProjectableGeometryTopology": "projectableGeometryTopology",
        "CognitiveLoadField": "cognitiveLoadField",
        "CadForceProbeReceipt": "cadForceProbeReceipt",
        "LogogramProjection": "logogramProjection",
        "HoldForUnlawfulOrUnderspecifiedShape": "holdForUnlawfulOrUnderspecifiedShape",
    }
    return mapping.get(shape, "holdForUnlawfulOrUnderspecifiedShape")


def _insert_rrc_receipt(
    pkg_id: str,
    session_id: str,
    alignment: dict[str, Any],
    target: str,
) -> None:
    """Insert an RRC alignment receipt bound to the build package."""
    pist = alignment["pist"]
    rrc = alignment["rrc"]
    residual = json.dumps(
        {
            "target": target,
            "predicted_shape": alignment["predicted_shape"],
            "tactic_family": alignment["tactic_family"],
            "alignment_status": rrc.get("alignmentStatus", "unknown"),
            "alignment_score": rrc.get("score"),
            "alignment_warnings": rrc.get("warnings", []),
            "pist_spectral_radius": pist.get("spectral_radius_q16"),
            "rrc_exit_code": rrc.get("exit_code"),
        },
        separators=(",", ":"),
    )
    sql = f"""
    INSERT INTO ene.receipts
        (package_id, receipt_type, receipt_hash, input_hash, output_hash,
         toolchain, verifier, status, residual)
    VALUES
        ('{pkg_id}', 'rrc_alignment', '{alignment['receipt_hash']}',
         '{hashlib.sha256(alignment['canonical_json'].encode()).hexdigest()}',
         '{alignment['receipt_hash']}',
         'pist-classify-trace; rrc-watchdog', 'opencode',
         '{rrc.get('alignmentStatus', 'unknown')}',
         '{residual}'::jsonb)
    ON CONFLICT (package_id, receipt_type, receipt_hash) DO NOTHING;
    """
    _ssh_psql(sql)



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")
    parser.add_argument("--rrc", type=Path, nargs="?", const=True, default=None,
                        help="enable RRC alignment. Uses <target>.trace.json in LEAN_DIR if no path is given.")
    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,
    )

    # Resolve RRC trace fixture if requested.
    rrc_alignment: dict[str, Any] | None = None
    if args.rrc is not None:
        explicit_path = args.rrc if isinstance(args.rrc, Path) else None
        trace_path = _find_rrc_trace(target, explicit_path, args.dir)
        if trace_path:
            print(f"RRC trace fixture: {trace_path}", file=sys.stderr)
            rrc_alignment = _run_rrc_alignment(target, trace_path, args.dir)
            if rrc_alignment:
                print(
                    f"RRC alignment: {rrc_alignment['rrc'].get('alignmentStatus')} "
                    f"score={rrc_alignment['rrc'].get('score')}",
                    file=sys.stderr,
                )
        else:
            print(
                f"RRC trace fixture not found for target '{target}' (use --rrc <path>)",
                file=sys.stderr,
            )

    if args.dry_run:
        payload = {"git": git, "build": build, "elapsed_ms": elapsed_ms}
        if rrc_alignment:
            payload["rrc_alignment"] = rrc_alignment
        print(json.dumps(payload, 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)
    if rrc_alignment:
        _insert_rrc_receipt(pkg_id, session_id, rrc_alignment, target)
    _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())