#!/usr/bin/env python3
"""
rrc_dual_query.py — Dual-query bridge: Neon (arxiv_papers) + Gremlin (module graph)

Wraps RRC classification output with:
  - Neon:   related papers from arxiv_papers matching kernel keywords
  - Gremlin: Lean modules that implement the matched kernel + their import chain

Usage (standalone):
    python3 4-Infrastructure/shim/rrc_dual_query.py \
        --name "sidon_test" --equation "|A| <= sqrt(2N)" --route "number_theory"

Import:
    from rrc_dual_query import enrich_classification
"""

from __future__ import annotations

import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Any

# ── Paths ─────────────────────────────────────────────────────────────────────

ROOT     = Path(__file__).resolve().parent.parent.parent
ENV_FILE = ROOT / ".env.gremlin"

# ── Neon connection (SSH → podman → psql) ─────────────────────────────────────

NEON_HOST = "neon-64gb"
CONTAINER = "arxiv-pg"
DB        = "arxiv"

def neon_query(sql: str, timeout: int = 60) -> list[list[str]]:
    result = subprocess.run(
        ["ssh", NEON_HOST,
         f"podman exec {CONTAINER} psql -U postgres -d {DB} -t -A -F '|' -c \"{sql}\""],
        capture_output=True, text=True, timeout=timeout,
    )
    if result.returncode != 0:
        return []
    return [line.split("|") for line in result.stdout.strip().split("\n") if line]

# ── Gremlin connection (mathblob) ─────────────────────────────────────────────

_gremlin_client = None

def _load_env():
    if not ENV_FILE.exists():
        raise FileNotFoundError(f"{ENV_FILE} not found — run setup_mathblob.py first")
    for line in ENV_FILE.read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            k, _, v = line.partition("=")
            os.environ.setdefault(k.strip(), v.strip())

def gremlin_client():
    global _gremlin_client
    if _gremlin_client is None:
        from gremlin_python.driver import client as gc, serializer
        _load_env()
        _gremlin_client = gc.Client(
            os.environ["GREMLIN_ENDPOINT"], "g",
            username=os.environ["GREMLIN_USERNAME"],
            password=os.environ["GREMLIN_PASSWORD"],
            message_serializer=serializer.GraphSONSerializersV2d0(),
        )
    return _gremlin_client

def gremlin_query(q: str, bindings: dict = None) -> list[Any]:
    try:
        c = gremlin_client()
        return c.submitAsync(q, bindings or {}).result().all().result()
    except Exception as e:
        print(f"  GREMLIN ERR: {e!s:.120}", file=sys.stderr)
        return []

# ── Kernel → keyword + module-name mapping ────────────────────────────────────

KERNEL_MAP: dict[str, dict] = {
    "diophantine": {
        "neon_keywords": ["diophantine", "integer solution", "ramanujan", "nagell", "goormaghtigh"],
        "module_patterns": ["Goormaghtigh", "Diophantine", "Erdos", "RRC"],
    },
    "combinatorics": {
        "neon_keywords": ["combinatorics", "extremal", "sidon", "erdős", "sumset"],
        "module_patterns": ["Sidon", "Erdos", "Combinatorics", "RRC"],
    },
    "sidon": {
        "neon_keywords": ["sidon", "B2 set", "sum-free", "additive combinatorics"],
        "module_patterns": ["Sidon", "SidonCollision", "SidonSets"],
    },
    "geometry": {
        "neon_keywords": ["riemannian", "geodesic", "manifold", "curvature", "topology"],
        "module_patterns": ["Geometry", "Manifold", "Topology", "Euclidean"],
    },
    "graph_reconstruction": {
        "neon_keywords": ["graph reconstruction", "deck", "ulam", "kelly"],
        "module_patterns": ["Graph", "Braid", "YangMills"],
    },
    "dataset": {
        "neon_keywords": ["dataset", "corpus", "benchmark"],
        "module_patterns": ["Corpus", "Emit", "RRC"],
    },
}

def _kernel_for(classification: dict) -> str:
    """Extract kernel stage name from RRC classification result."""
    stage = classification.get("stage", "") or classification.get("kernel", "")
    for key in KERNEL_MAP:
        if key in stage.lower():
            return key
    return "combinatorics"  # default

# ── Neon: fetch related papers ─────────────────────────────────────────────────

def neon_related_papers(kernel: str, limit: int = 5) -> list[dict]:
    cfg = KERNEL_MAP.get(kernel, {})
    keywords = cfg.get("neon_keywords", [])
    if not keywords:
        return []

    conditions = " OR ".join(
        f"(lower(title) LIKE '%{kw}%' OR lower(abstract) LIKE '%{kw}%')"
        for kw in keywords
    )
    sql = (
        f"SELECT paper_id, title, substring(abstract, 1, 200) "
        f"FROM arxiv_papers WHERE {conditions} LIMIT {limit}"
    )
    rows = neon_query(sql)
    return [{"paper_id": r[0], "title": r[1], "abstract": r[2]} for r in rows if len(r) >= 3]

# ── Gremlin: fetch implementing modules ───────────────────────────────────────

def gremlin_implementing_modules(kernel: str) -> list[dict]:
    cfg = KERNEL_MAP.get(kernel, {})
    patterns = cfg.get("module_patterns", [])
    if not patterns:
        return []

    results = []
    seen = set()
    for pat in patterns:
        # Try TextP.containing (Cosmos DB may not support it) → fall back to
        # fetching all internal module IDs and filtering client-side.
        rows = gremlin_query(
            "g.V().hasLabel('module').has('kind','internal')"
            ".has('id', TextP.containing(pat)).limit(10).values('id')",
            {"pat": pat},
        )
        if not rows:
            # Client-side fallback: pull all internal module ids, filter locally
            all_ids = gremlin_query(
                "g.V().hasLabel('module').has('kind','internal').values('id')"
            )
            rows = [mid for mid in all_ids if pat.lower() in mid.lower()][:10]

        for mod_id in rows:
            if mod_id not in seen:
                seen.add(mod_id)
                results.append({"module": mod_id, "pattern": pat})

    return results

def gremlin_import_chain(module_id: str, depth: int = 3) -> list[str]:
    """Return the transitive imports of a module up to `depth` hops."""
    rows = gremlin_query(
        "g.V().has('module','id',mid)"
        ".repeat(out('imports')).times(depth).emit()"
        ".has('kind','internal').dedup().limit(20).values('id')",
        {"mid": module_id, "depth": depth},
    )
    return rows

def gremlin_dependents(module_id: str) -> list[str]:
    """Return modules that import this one (in-edges)."""
    rows = gremlin_query(
        "g.V().has('module','id',mid).in('imports').limit(15).values('id')",
        {"mid": module_id},
    )
    return rows

# ── Combined enrichment ────────────────────────────────────────────────────────

def enrich_classification(classification: dict) -> dict:
    """
    Takes RRC classification output dict, queries Neon + Gremlin,
    returns enriched receipt.

    classification should have at least: name, equation, stage/kernel, score.
    """
    kernel = _kernel_for(classification)

    papers   = neon_related_papers(kernel)
    modules  = gremlin_implementing_modules(kernel)

    # For each implementing module, get its import chain
    module_details = []
    for m in modules[:3]:  # limit to top 3 to keep latency low
        mid = m["module"]
        chain      = gremlin_import_chain(mid)
        dependents = gremlin_dependents(mid)
        module_details.append({
            "module":     mid,
            "imports":    chain,
            "depended_by": dependents,
        })

    return {
        **classification,
        "kernel_matched": kernel,
        "neon_papers":    papers,
        "lean_modules":   module_details,
    }

# ── CLI ───────────────────────────────────────────────────────────────────────

def main():
    import argparse
    sys.path.insert(0, str(Path(__file__).resolve().parent))
    from rrc_self_classify import classify_equation

    parser = argparse.ArgumentParser(description="RRC dual-query: classify + enrich")
    parser.add_argument("--name",     required=True)
    parser.add_argument("--equation", required=True)
    parser.add_argument("--route",    default="")
    parser.add_argument("--json",     action="store_true", help="Output raw JSON")
    args = parser.parse_args()

    print(f"Classifying: {args.name!r}...")
    classification = classify_equation(args.name, args.equation, args.route)

    print(f"Enriching with Neon + Gremlin (kernel: {_kernel_for(classification)})...")
    enriched = enrich_classification(classification)

    if args.json:
        print(json.dumps(enriched, indent=2, default=str))
        return

    print(f"\n── Classification ───────────────────────────────────────")
    print(f"  Name   : {enriched.get('name', args.name)}")
    print(f"  Kernel : {enriched['kernel_matched']}")
    print(f"  Stage  : {enriched.get('stage', '?')}")
    print(f"  Score  : {enriched.get('score', '?')}")

    print(f"\n── Related papers (Neon/arxiv) ──────────────────────────")
    papers = enriched.get("neon_papers", [])
    if papers:
        for p in papers:
            print(f"  [{p['paper_id']}] {p['title'][:70]}")
    else:
        print("  (none found — check neon-64gb SSH connection)")

    print(f"\n── Lean modules (Gremlin/mathblob) ─────────────────────")
    mods = enriched.get("lean_modules", [])
    if mods:
        for m in mods:
            print(f"  {m['module']}")
            if m["imports"]:
                print(f"    imports: {', '.join(m['imports'][:5])}")
            if m["depended_by"]:
                print(f"    used by: {', '.join(m['depended_by'][:5])}")
    else:
        print("  (no matching modules found in graph)")

if __name__ == "__main__":
    main()