#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.11"
# dependencies = ["gremlinpython>=3.7", "psycopg2-binary", "requests"]
# ///
# INFRA:LIVE gremlin
# INFRA:LIVE postgres
# INFRA:LIVE contextstream
"""
distribute_extraction.py

Reads curated extraction/ideas.json and extraction/math.json and distributes to:
  - Gremlin  (ideas → graph vertices + edges)
  - PostgreSQL (math → math_objects table on arxiv DB)

By default uses the hand-curated files. Use --from-merged to ingest the
auto-extracted all_concepts_merged.json instead.

AppFlowy and ContextStream distribution handled separately
(AppFlowy needs REST API credentials; ContextStream via MCP session).

Usage:
  uv run scripts/distribute_extraction.py [--dry-run] [--from-merged] [--target gremlin|postgres|all]
"""

import json
import os
import re
import sys
import argparse
import time
from pathlib import Path

ROOT = Path(__file__).parent.parent

# ── Load infra.json first — abort if service is DEAD ─────────────────────────

def load_infra() -> dict:
    p = ROOT / "infra.json"
    if not p.exists():
        sys.exit("ERROR: infra.json not found. Run from project root.")
    with open(p) as f:
        return json.load(f)

def require_live(infra: dict, service: str):
    s = infra["services"].get(service, {})
    status = s.get("status", "UNKNOWN")
    if status != "LIVE":
        sys.exit(f"ERROR: service '{service}' is {status} per infra.json. Aborting.")

# ── Gremlin distribution (ideas → vertices + edges) ──────────────────────────

def gremlin_submit(gc, query, max_retries=5, base_delay=0.5):
    """Submit a Gremlin query with exponential backoff for rate limiting."""
    for attempt in range(max_retries):
        try:
            return gc.submit(query).all().result()
        except Exception as exc:
            err = str(exc).lower()
            is_rate_limit = any(tag in err for tag in [
                "requestratetoolarge", "too many requests", "429",
                "request rate is large", "throttled"
            ])
            if not is_rate_limit or attempt == max_retries - 1:
                raise
            delay = base_delay * (2 ** attempt)
            print(f"  [retry {attempt + 1}/{max_retries}] rate limited, sleeping {delay:.2f}s")
            sys.stdout.flush()
            time.sleep(delay)


def distribute_gremlin(ideas: list, infra: dict, dry_run: bool):
    require_live(infra, "gremlin")
    creds_file = ROOT / infra["services"]["gremlin"].get("credentials", ".env.gremlin")

    env = {}
    if creds_file.exists():
        for line in creds_file.read_text().splitlines():
            if "=" in line and not line.startswith("#"):
                k, _, v = line.partition("=")
                env[k.strip()] = v.strip().strip('"')

    endpoint = env.get("GREMLIN_ENDPOINT") or infra["services"]["gremlin"]["endpoint"]
    key      = env.get("GREMLIN_PASSWORD") or env.get("GREMLIN_KEY", "")
    database = infra["services"]["gremlin"].get("database", "research")
    graph    = infra["services"]["gremlin"].get("graph", "concepts")

    print(f"\n[Gremlin] endpoint={endpoint}  graph={graph}  ideas={len(ideas)}")

    def vertex_label(idea):
        lbl = idea.get('label', '')
        if lbl and re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", lbl) and len(lbl) <= 40:
            return lbl
        return "concept"

    if dry_run:
        for idea in ideas:
            print(f"  DRY-RUN vertex: id={idea['id']}  label={vertex_label(idea)}")
        return

    try:
        from gremlin_python.driver import client as gremlin_client
        from gremlin_python.driver.serializer import GraphSONSerializersV2d0
    except ImportError as e:
        print(f"  [skip] gremlinpython import failed: {e}")
        print("  Run: uv pip install gremlinpython")
        return

    gc = gremlin_client.Client(
        endpoint,
        "g",
        username=f"/dbs/{database}/colls/{graph}",
        password=key,
        message_serializer=GraphSONSerializersV2d0(),
    )

    for idea in ideas:
        name = idea.get('name', '')
        title = idea.get('title', '')
        summary = idea.get('summary', '') or idea.get('description', '')
        props = {
            "title":   title or name,
            "name":    name or title,
            "summary": summary[:500],
            "status":  idea.get("status", ""),
            "type":    idea.get("type", ""),
            "tags":    ",".join(idea.get("tags", [])),
            "source":  idea.get("source", idea.get("source_file", "")),
        }
        prop_str = " ".join(f".property('{k}', '{v}')" for k, v in props.items()
                            if v)
        label = vertex_label(idea)
        query = f"g.V().has('id', '{idea['id']}').fold().coalesce(unfold(), addV('{label}').property('id', '{idea['id']}').property('pk', '{idea['id']}')){prop_str}"
        try:
            gremlin_submit(gc, query)
            print(f"  vertex upserted: {idea['id']} ({label[:40]})")
            sys.stdout.flush()
        except Exception as exc:
            err = str(exc)[:120]
            print(f"  ERROR {idea['id']} {err}")
            sys.stdout.flush()

    # Add edges
    for idea in ideas:
        for edge in idea.get("edges", []):
            query = (f"g.V().has('id', '{idea['id']}').has('pk', '{idea['id']}')"
                     f".coalesce("
                     f"  outE('{edge['rel']}').where(inV().has('id', '{edge['to']}').has('pk', '{edge['to']}')).fold(),"
                     f"  addE('{edge['rel']}').to(g.V().has('id', '{edge['to']}').has('pk', '{edge['to']}'))"
                     f")")
            try:
                gremlin_submit(gc, query)
                print(f"  edge: {idea['id']} --[{edge['rel']}]--> {edge['to']}")
            except Exception as exc:
                print(f"  edge skip ({idea['id']}→{edge['to']}): {exc}")

    gc.close()

# ── PostgreSQL distribution (math → math_objects table) ──────────────────────

def distribute_postgres(math_items: list, infra: dict, dry_run: bool, sql_only: bool = False):
    require_live(infra, "postgres")
    svc = infra["services"]["postgres"]
    host = svc.get("host", "neon-64gb")

    print(f"\n[PostgreSQL] host={host}  math_items={len(math_items)}")

    if dry_run:
        for m in math_items:
            print(f"  DRY-RUN row: id={m['id']}  type={m['type']}")
        return

    import subprocess, shlex

    def esc(s):
        return str(s or "").replace("'", "''")

    # Build complete SQL
    sql_lines = [
        "CREATE TABLE IF NOT EXISTS math_objects (",
        "    id TEXT PRIMARY KEY, type TEXT NOT NULL, name TEXT NOT NULL,",
        "    expression TEXT, latex TEXT, description TEXT,",
        "    source_file TEXT, lean_name TEXT, lean_status TEXT,",
        "    tags TEXT[], related TEXT[], created_at TIMESTAMP DEFAULT NOW()",
        ");",
        ""
    ]

    for m in math_items:
        tags    = "{" + ",".join(esc(t) for t in m.get("tags", []))    + "}"
        related = "{" + ",".join(esc(r) for r in m.get("related", [])) + "}"
        sql_lines.append(f"""
            INSERT INTO math_objects
              (id,type,name,expression,latex,description,source_file,lean_name,lean_status,tags,related)
            VALUES (
              '{esc(m["id"])}','{esc(m["type"])}','{esc(m["name"])}',
              '{esc(m.get("expression"))}','{esc(m.get("latex"))}',
              '{esc(m.get("description"))}','{esc(m.get("source_file"))}',
              '{esc(m.get("lean_name"))}','{esc(m.get("lean_status"))}',
              '{tags}','{related}'
            )
            ON CONFLICT (id) DO UPDATE SET
              expression=EXCLUDED.expression, latex=EXCLUDED.latex,
              description=EXCLUDED.description, lean_status=EXCLUDED.lean_status,
              tags=EXCLUDED.tags;
        """.strip())

    sql_file = ROOT / "extraction" / "math_postgres.sql"
    sql_file.write_text("\n".join(sql_lines))
    print(f"  SQL written to {sql_file}")

    if sql_only:
        print(f"  --sql-only mode: file written, not executed. Pipe via SSH:")
        print(f"     ssh {host} 'podman exec -i arxiv-pg psql -U postgres -d arxiv' < {sql_file}")
        return

    # Execute via SSH in one pass
    cmd = ["ssh", host, "podman exec -i arxiv-pg psql -U postgres -d arxiv"]
    proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    stdout, stderr = proc.communicate(input="\n".join(sql_lines))

    if proc.returncode != 0:
        print(f"  ERROR: {stderr.strip()}")
        return

    print("  math_objects table ensured")
    for m in math_items:
        print(f"  upserted: {m['id']} ({m.get('lean_status','')})")

# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description="Distribute extraction JSON to live infra")
    ap.add_argument("--dry-run",     action="store_true", help="Print actions without executing")
    ap.add_argument("--from-merged", action="store_true", help="Use auto-extracted all_concepts_merged.json instead of curated ideas.json/math.json")
    ap.add_argument("--target",      default="all",       help="gremlin|postgres|all")
    ap.add_argument("--sql-only",    action="store_true", help="Write SQL file but don't execute (postgres only)")
    ap.add_argument("--limit",       type=int, default=0,  help="Process only the first N items per target (0 = all)")
    ap.add_argument("--offset",      type=int, default=0,  help="Skip the first N items per target")
    args = ap.parse_args()

    infra    = load_infra()
    math_types = {"theorem","definition","structure","type","predicate",
                  "operator","invariant","constant","equation","algorithm",
                  "dataset","model"}

    if args.from_merged:
        merged   = json.loads((ROOT / "extraction/all_concepts_merged.json").read_text())["math"]
        math_obj = [x for x in merged if x.get("type") in math_types]
        ideas    = [x for x in merged if x.get("type") == "concept"]
    else:
        ideas    = json.loads((ROOT / "extraction/ideas.json").read_text())["ideas"]
        math_obj = [x for x in json.loads((ROOT / "extraction/math.json").read_text())["math"]
                    if x.get("type") in math_types]

    def slice(items):
        start = args.offset
        end = None if args.limit == 0 else start + args.limit
        return items[start:end]

    ideas = slice(ideas)
    math_obj = slice(math_obj)

    print(f"Loaded {len(ideas)} ideas, {len(math_obj)} math objects (offset={args.offset}, limit={args.limit})")
    print(f"Dry run: {args.dry_run}  From merged: {args.from_merged}  Target: {args.target}  SQL-only: {args.sql_only}")

    if args.target in ("gremlin", "all"):
        distribute_gremlin(ideas, infra, args.dry_run)

    if args.target in ("postgres", "all"):
        distribute_postgres(math_obj, infra, args.dry_run, args.sql_only)

    print("\nDone.")

if __name__ == "__main__":
    main()