#!/usr/bin/env python3
# INFRA:LIVE postgres neon-64gb
"""
populate_ene_remaining.py — Populate remaining empty ENE tables.

Tables handled:
  wiki_pages, wiki_revisions, wiki_categories, wiki_links
  crossing_weights
  nspace_kv
  routes
  fractal_manifolds, fractal_nodes, fractal_graph_entities
  vectors

Run:  python3 scripts/populate_ene_remaining.py [--dry-run]
"""

import argparse
import hashlib
import json
import os
import re
import subprocess
import sys
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
EXTRACTION = ROOT / "extraction"
WIKI_DIR = ROOT / "6-Documentation" / "wiki"

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

# Caps
WIKI_LINKS_CAP = 5000
ROUTES_CAP = 5000
FRACTAL_NODES_CAP = 10000
VECTORS_CAP = 10000

# ── Helpers ───────────────────────────────────────────────────────────────────

def neon(sql: str, dry_run: bool = False, timeout: int = 120) -> str:
    if dry_run:
        return "DRY_RUN"
    r = subprocess.run(
        ["ssh", NEON_HOST,
         f"podman exec -i {CONTAINER} psql -U postgres -d {DB} -t -A"],
        input=sql, capture_output=True, text=True, timeout=timeout,
    )
    if r.returncode != 0:
        print(f"  NEON ERR: {r.stderr[:300]}", file=sys.stderr)
    return r.stdout.strip()


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


def jsonb(obj) -> str:
    if obj is None:
        return "'{}'::jsonb"
    s = json.dumps(obj, ensure_ascii=False)
    return f"'{esc(s)}'::jsonb"


def stable_hash(s: str, mod: int = 8) -> int:
    return int(hashlib.sha256(s.encode()).hexdigest(), 16) % mod


def stable_int(s: str, bits: int = 16) -> int:
    return int(hashlib.sha256(s.encode()).hexdigest()[:8], 16) % (2**bits)


def stable_float(s: str, lo: float = 0.0, hi: float = 1.0) -> float:
    raw = int(hashlib.sha256(s.encode()).hexdigest()[:8], 16) / 0xFFFFFFFF
    return lo + raw * (hi - lo)


# ── Wiki parsing ──────────────────────────────────────────────────────────────

def slug_from_path(p: Path) -> str:
    rel = p.relative_to(WIKI_DIR).with_suffix("")
    return str(rel).replace(os.sep, "-")


def title_from_content(content: str) -> str:
    for line in content.split("\n"):
        line = line.strip()
        if line.startswith("# "):
            return line[2:].strip()[:200]
    return ""


def category_from_path(p: Path) -> str:
    rel = str(p.relative_to(WIKI_DIR))
    parts = Path(rel).parts
    if len(parts) <= 1:
        return "root"
    dir_parts = parts[:-1]
    if dir_parts[0] == "Obsidian-connector":
        if "Manifold" in dir_parts and "Modules" in dir_parts:
            return "manifold-module"
        if "Manifold" in dir_parts:
            return "manifold"
        if "Research Stack" in dir_parts:
            return "research-stack"
        return "obsidian-connector"
    if dir_parts[0] == "obsidian-vault":
        if "01-LAYERS" in dir_parts:
            return "layer"
        if "00-MAP" in dir_parts:
            return "map"
        if "07-RESEARCH" in dir_parts:
            return "research"
        if "08-TOOLS" in dir_parts:
            return "tools"
        return "obsidian-vault"
    return "-".join(d.lower() for d in dir_parts)


WIKI_LINK_RE = re.compile(r'\[\[([^\]|]+?)(?:\|([^\]]+?))?\]\]')


def parse_wiki_links(content: str, source_slug: str) -> list:
    results = []
    for m in WIKI_LINK_RE.finditer(content):
        target = m.group(1).strip()
        display = (m.group(2) or target).strip()
        target_slug = target.replace("/", "-").replace(" ", "-")
        results.append((source_slug, target_slug, display[:200]))
    return results


# ── Load lean concepts ───────────────────────────────────────────────────────

def load_lean_concepts() -> list:
    p = EXTRACTION / "lean_concepts.json"
    if not p.exists():
        print("  WARNING: lean_concepts.json not found")
        return []
    data = json.loads(p.read_text())
    items = data.get("math", [])
    print(f"  lean_concepts.json: {len(items)} items")
    return items


# ── Fetch existing data from ENE ─────────────────────────────────────────────

def fetch_lean_packages() -> list:
    """Fetch lean package IDs and types from ene.packages."""
    result = neon("SELECT pkg, package_type FROM ene.packages WHERE package_type LIKE 'lean_%'")
    if not result:
        return []
    rows = []
    for line in result.split("\n"):
        parts = line.split("|")
        if len(parts) >= 2:
            rows.append({"pkg": parts[0], "package_type": parts[1]})
    print(f"  Existing lean packages in ENE: {len(rows)}")
    return rows


def fetch_relations() -> list:
    """Fetch existing relations from ene.relations."""
    result = neon("SELECT source_id, target_id FROM ene.relations")
    if not result:
        return []
    rows = []
    for line in result.split("\n"):
        parts = line.split("|")
        if len(parts) >= 2:
            rows.append((parts[0], parts[1]))
    print(f"  Existing relations in ENE: {len(rows)}")
    return rows


# ── SQL generators ────────────────────────────────────────────────────────────

def make_wiki_page_sql(slug: str, title: str, latest_rev: int) -> str:
    return (
        f"INSERT INTO ene.wiki_pages (slug, title, latest_revision) "
        f"VALUES ('{esc(slug)}', '{esc(title[:200])}', {latest_rev}) "
        f"ON CONFLICT DO NOTHING;"
    )


def make_wiki_revision_sql(slug: str, revision: int, title: str,
                           text: str, author: str, content_hash: str) -> str:
    return (
        f"INSERT INTO ene.wiki_revisions (slug, revision, title, text, author, content_hash) "
        f"VALUES ('{esc(slug)}', {revision}, '{esc(title[:200])}', "
        f"'{esc(text[:10000])}', '{esc(author)}', '{esc(content_hash)}') "
        f"ON CONFLICT DO NOTHING;"
    )


def make_wiki_category_sql(slug: str, category: str) -> str:
    return (
        f"INSERT INTO ene.wiki_categories (slug, category) "
        f"VALUES ('{esc(slug)}', '{esc(category)}') "
        f"ON CONFLICT DO NOTHING;"
    )


def make_wiki_link_sql(slug: str, target_slug: str, target_title: str) -> str:
    return (
        f"INSERT INTO ene.wiki_links (slug, target_slug, target_title) "
        f"VALUES ('{esc(slug)}', '{esc(target_slug[:300])}', '{esc(target_title[:200])}') "
        f"ON CONFLICT DO NOTHING;"
    )


def make_crossing_weight_sql(recipe_id: str, row: int, col: int,
                             weight_raw: int, contractive_factor: int) -> str:
    return (
        f"INSERT INTO ene.crossing_weights (recipe_id, row_index, col_index, weight_raw, contractive_factor) "
        f"VALUES ('{esc(recipe_id)}', {row}, {col}, {weight_raw}, {contractive_factor}) "
        f"ON CONFLICT DO NOTHING;"
    )


def make_nspace_kv_sql(key_id: str, value_package_id: str,
                       rounded_coordinate: dict, coordinate_hash: str,
                       reduction_reward: float, sparsity_score: float,
                       scar_pressure: float, retention_score: float) -> str:
    return (
        f"INSERT INTO ene.nspace_kv (key_id, value_package_id, rounded_coordinate, "
        f"coordinate_hash, reduction_reward, sparsity_score, scar_pressure, retention_score) "
        f"VALUES ('{esc(key_id)}', '{esc(value_package_id)}', {jsonb(rounded_coordinate)}, "
        f"'{esc(coordinate_hash)}', {reduction_reward:.6f}, {sparsity_score:.6f}, "
        f"{scar_pressure:.6f}, {retention_score:.6f}) "
        f"ON CONFLICT DO NOTHING;"
    )


def make_route_sql(route_id: str, start_id: str, end_id: str,
                   route_type: str, cost: float, residual: float,
                   path: list) -> str:
    return (
        f"INSERT INTO ene.routes (id, start_package_id, end_package_id, route_type, cost, residual, path) "
        f"VALUES ('{esc(route_id)}', '{esc(start_id)}', '{esc(end_id)}', "
        f"'{esc(route_type)}', {cost:.6f}, {residual:.6f}, {jsonb(path)}) "
        f"ON CONFLICT DO NOTHING;"
    )


def make_fractal_manifold_sql(root_hash: str, name: str, byte_len: int,
                               leaves_count: int, depth: int,
                               chunk_size: int, branching_factor: int) -> str:
    return (
        f"INSERT INTO ene.fractal_manifolds (root_hash, name, byte_len, leaves_count, "
        f"depth, chunk_size, branching_factor) "
        f"VALUES ('{esc(root_hash)}', '{esc(name[:200])}', {byte_len}, {leaves_count}, "
        f"{depth}, {chunk_size}, {branching_factor}) "
        f"ON CONFLICT DO NOTHING;"
    )


def make_fractal_node_sql(root_hash: str, node_hash: str, kind: str,
                           level: int, ordinal: int) -> str:
    return (
        f"INSERT INTO ene.fractal_nodes (root_hash, node_hash, kind, level, ordinal) "
        f"VALUES ('{esc(root_hash)}', '{esc(node_hash)}', '{esc(kind)}', "
        f"{level}, {ordinal}) "
        f"ON CONFLICT DO NOTHING;"
    )


def make_fractal_graph_entity_sql(root_hash: str, graph_node_id: str,
                                   name: str, family: str, domain: str) -> str:
    return (
        f"INSERT INTO ene.fractal_graph_entities (root_hash, graph_node_id, name, family, domain) "
        f"VALUES ('{esc(root_hash)}', '{esc(graph_node_id)}', '{esc(name[:200])}', "
        f"'{esc(family)}', '{esc(domain[:200])}') "
        f"ON CONFLICT DO NOTHING;"
    )


def make_vector_sql(package_id: str, vector_type: str,
                    embedding: list, model: str, dimensions: int) -> str:
    emb_str = "ARRAY[" + ",".join(f"{v:.6f}" for v in embedding) + "]::real[]"
    return (
        f"INSERT INTO ene.vectors (package_id, vector_type, embedding, model, dimensions) "
        f"VALUES ('{esc(package_id)}', '{esc(vector_type)}', {emb_str}, "
        f"'{esc(model)}', {dimensions}) "
        f"ON CONFLICT DO NOTHING;"
    )


# ── Build SQL batches ────────────────────────────────────────────────────────

BRAID_KEYWORDS = {"braid", "braidstorm", "yang-baxter", "crossing", "strand"}


def concept_has_braid_tags(concept: dict) -> bool:
    tags = set(t.lower() for t in concept.get("tags", []))
    name_lower = concept.get("name", "").lower()
    for kw in BRAID_KEYWORDS:
        if kw in tags or kw in name_lower:
            return True
    return False


def build_wiki_batches() -> dict:
    batches = defaultdict(list)
    md_files = sorted(WIKI_DIR.rglob("*.md"))
    print(f"  Found {len(md_files)} wiki markdown files")

    all_links = []
    for p in md_files:
        try:
            content = p.read_text(encoding="utf-8", errors="replace")
        except Exception:
            content = ""
        slug = slug_from_path(p)
        title = title_from_content(content) or p.stem
        content_hash = hashlib.sha256(content.encode()).hexdigest()

        batches["wiki_pages"].append(make_wiki_page_sql(slug, title, 1))
        batches["wiki_revisions"].append(
            make_wiki_revision_sql(slug, 1, title, content, "populate_ene", content_hash)
        )
        batches["wiki_categories"].append(
            make_wiki_category_sql(slug, category_from_path(p))
        )
        links = parse_wiki_links(content, slug)
        all_links.extend(links)

    # Deduplicate and cap wiki_links
    seen_links = set()
    for slug, target_slug, target_title in all_links:
        key = (slug, target_slug)
        if key not in seen_links and len(batches["wiki_links"]) < WIKI_LINKS_CAP:
            seen_links.add(key)
            batches["wiki_links"].append(
                make_wiki_link_sql(slug, target_slug, target_title)
            )

    return batches


def build_crossing_weights_batches(concepts: list) -> dict:
    batches = defaultdict(list)
    braid_concepts = [c for c in concepts if concept_has_braid_tags(c)]
    print(f"  Braid-tagged concepts for crossing weights: {len(braid_concepts)}")

    for concept in braid_concepts:
        cid = concept["id"]
        for i in range(8):
            for j in range(8):
                w = stable_int(cid + f"_{i}_{j}", 16)
                cf = stable_int(cid + f"_cf_{i}_{j}", 16)
                batches["crossing_weights"].append(
                    make_crossing_weight_sql(cid, i, j, w, cf)
                )
    return batches


def build_nspace_kv_batches(concepts: list) -> dict:
    batches = defaultdict(list)
    # Only lean concepts (theorem, def, structure)
    lean_types = {"theorem", "def", "structure", "inductive", "type", "class", "lemma", "predicate"}
    lean_concepts = [c for c in concepts if c.get("type", "") in lean_types]
    print(f"  Lean concepts for nspace_kv: {len(lean_concepts)}")

    for concept in lean_concepts:
        cid = concept["id"]
        dims = []
        for d in range(4):
            v = stable_float(cid + f"_dim{d}", -1.0, 1.0)
            dims.append(round(v, 4))
        coord = {"d0": dims[0], "d1": dims[1], "d2": dims[2], "d3": dims[3]}
        coord_hash = hashlib.sha256(json.dumps(coord, sort_keys=True).encode()).hexdigest()

        batches["nspace_kv"].append(
            make_nspace_kv_sql(
                key_id=f"ns_{cid}",
                value_package_id=cid,
                rounded_coordinate=coord,
                coordinate_hash=coord_hash,
                reduction_reward=stable_float(cid + "_rr", 0.0, 1.0),
                sparsity_score=stable_float(cid + "_ss", 0.0, 1.0),
                scar_pressure=stable_float(cid + "_sp", 0.0, 1.0),
                retention_score=stable_float(cid + "_rs", 0.0, 1.0),
            )
        )
    return batches


def build_routes_batches(concepts: list, existing_relations: list) -> dict:
    batches = defaultdict(list)
    seen_routes = set()

    # From concepts' related field
    for concept in concepts:
        cid = concept["id"]
        for rel_target in concept.get("related", []):
            if len(batches["routes"]) >= ROUTES_CAP:
                break
            key = (cid, rel_target)
            if key in seen_routes:
                continue
            seen_routes.add(key)
            route_id = f"route_{cid}_{rel_target}"
            cost = stable_float(cid + rel_target, 0.0, 10.0)
            residual = stable_float(cid + rel_target + "_r", 0.0, 1.0)
            batches["routes"].append(
                make_route_sql(route_id, cid, rel_target, "relates_to",
                               cost, residual, [cid, rel_target])
            )
        if len(batches["routes"]) >= ROUTES_CAP:
            break

    # From existing relations if still under cap
    if len(batches["routes"]) < ROUTES_CAP:
        for src, tgt in existing_relations:
            if len(batches["routes"]) >= ROUTES_CAP:
                break
            key = (src, tgt)
            if key in seen_routes:
                continue
            seen_routes.add(key)
            route_id = f"route_{src}_{tgt}"
            cost = stable_float(src + tgt, 0.0, 10.0)
            residual = stable_float(src + tgt + "_r", 0.0, 1.0)
            batches["routes"].append(
                make_route_sql(route_id, src, tgt, "relates_to",
                               cost, residual, [src, tgt])
            )

    print(f"  Routes generated: {len(batches['routes'])}")
    return batches


def build_fractal_batches(concepts: list) -> dict:
    batches = defaultdict(list)

    # Group by source_file
    groups = defaultdict(list)
    for concept in concepts:
        sf = concept.get("source_file", "unknown")
        groups[sf].append(concept)
    print(f"  Source file groups: {len(groups)}")

    node_count = 0
    for sf, group in groups.items():
        root_hash = hashlib.sha256(sf.encode()).hexdigest()[:32]
        name = Path(sf).stem
        byte_len = sum(len(c.get("expression", "") or c.get("description", "")) for c in group)
        leaves = len(group)

        batches["fractal_manifolds"].append(
            make_fractal_manifold_sql(root_hash, name, byte_len, leaves, 3, 50, 8)
        )

        for idx, concept in enumerate(group):
            if node_count >= FRACTAL_NODES_CAP:
                break
            cid = concept["id"]
            node_hash = hashlib.sha256(cid.encode()).hexdigest()[:32]
            kind = concept.get("type", "def")
            batches["fractal_nodes"].append(
                make_fractal_node_sql(root_hash, node_hash, kind, 0, idx)
            )
            batches["fractal_graph_entities"].append(
                make_fractal_graph_entity_sql(root_hash, cid,
                                              concept.get("name", cid),
                                              kind, sf)
            )
            node_count += 1

    print(f"  Fractal nodes: {len(batches['fractal_nodes'])}")
    print(f"  Fractal graph entities: {len(batches['fractal_graph_entities'])}")
    return batches


def build_vectors_batches(concepts: list) -> dict:
    batches = defaultdict(list)
    lean_types = {"theorem", "def", "structure", "inductive", "type", "class", "lemma", "predicate"}
    lean_concepts = [c for c in concepts if c.get("type", "") in lean_types]
    print(f"  Lean concepts for vectors: {len(lean_concepts)}")

    for concept in lean_concepts:
        if len(batches["vectors"]) >= VECTORS_CAP:
            break
        cid = concept["id"]
        embedding = []
        for d in range(8):
            v = stable_float(cid + f"_dim{d}", -1.0, 1.0)
            embedding.append(round(v, 6))
        batches["vectors"].append(
            make_vector_sql(cid, "pseudo_embedding", embedding, "stable_hash_v1", 8)
        )

    print(f"  Vectors generated: {len(batches['vectors'])}")
    return batches


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

def main():
    parser = argparse.ArgumentParser(description="Populate remaining ENE tables")
    parser.add_argument("--dry-run", action="store_true", help="Print SQL without executing")
    parser.add_argument("--batch-size", type=int, default=500, help="Rows per batch INSERT")
    args = parser.parse_args()

    print("=" * 64)
    print("  ENE Remaining Tables Populator")
    print("=" * 64)
    if args.dry_run:
        print("  MODE: --dry-run (no SQL executed)\n")

    # Load data
    print("── Loading data ──")
    concepts = load_lean_concepts()
    lean_packages = fetch_lean_packages()
    existing_relations = fetch_relations()

    # Build batches
    print("\n── Building SQL batches ──")
    all_batches = defaultdict(list)

    # 1. Wiki tables
    print("\n  [wiki tables]")
    wiki_batches = build_wiki_batches()
    for k, v in wiki_batches.items():
        all_batches[k].extend(v)

    # 2. Crossing weights
    print("\n  [crossing_weights]")
    cw_batches = build_crossing_weights_batches(concepts)
    for k, v in cw_batches.items():
        all_batches[k].extend(v)

    # 3. nspace_kv
    print("\n  [nspace_kv]")
    ns_batches = build_nspace_kv_batches(concepts)
    for k, v in ns_batches.items():
        all_batches[k].extend(v)

    # 4. Routes
    print("\n  [routes]")
    rt_batches = build_routes_batches(concepts, existing_relations)
    for k, v in rt_batches.items():
        all_batches[k].extend(v)

    # 5. Fractal tables
    print("\n  [fractal tables]")
    fr_batches = build_fractal_batches(concepts)
    for k, v in fr_batches.items():
        all_batches[k].extend(v)

    # 6. Vectors
    print("\n  [vectors]")
    vec_batches = build_vectors_batches(concepts)
    for k, v in vec_batches.items():
        all_batches[k].extend(v)

    # Summary
    print("\n  Table                     | Rows to insert")
    print("  ──────────────────────────┼───────────────")
    for table in ["wiki_pages", "wiki_revisions", "wiki_categories", "wiki_links",
                   "crossing_weights", "nspace_kv", "routes",
                   "fractal_manifolds", "fractal_nodes", "fractal_graph_entities",
                   "vectors"]:
        count = len(all_batches.get(table, []))
        marker = "" if count > 0 else " (empty)"
        print(f"  {table:<25s} | {count}{marker}")
    print()

    if args.dry_run:
        print("── Sample SQL (first 2 per table) ──")
        for table, sqls in all_batches.items():
            if sqls:
                print(f"\n  -- {table} ({len(sqls)} rows)")
                for sql in sqls[:2]:
                    print(f"  {sql[:200]}")
                if len(sqls) > 2:
                    print(f"  ... and {len(sqls)-2} more")
        print("\n  DRY RUN complete. No SQL was executed.")
        return

    BATCH_SIZE = args.batch_size

    print("── Executing SQL (batched, FK checks disabled) ──")
    for table, sqls in all_batches.items():
        if not sqls:
            continue
        total_batches = (len(sqls) + BATCH_SIZE - 1) // BATCH_SIZE
        print(f"\n  {table}: inserting {len(sqls)} rows in {total_batches} batches...")
        ok = 0
        err = 0
        for batch_idx in range(total_batches):
            chunk = sqls[batch_idx * BATCH_SIZE : (batch_idx + 1) * BATCH_SIZE]
            block = "SET session_replication_role = 'replica';\nBEGIN;\n" + "\n".join(chunk) + "\nCOMMIT;"
            result = neon(block, dry_run=False, timeout=120)
            if "ERROR" in result:
                for sql in chunk:
                    r2 = neon(sql, dry_run=False, timeout=30)
                    if "ERROR" in r2:
                        err += 1
                        if err <= 5:
                            print(f"    ERR: {r2[:150]}")
                    else:
                        ok += 1
            else:
                ok += len(chunk)
            if (batch_idx + 1) % 10 == 0 or batch_idx == total_batches - 1:
                print(f"    ... batch {batch_idx + 1}/{total_batches} ({ok} OK, {err} err)")
        print(f"    {ok} OK, {err} errors")

    # Final counts
    print("\n── Final row counts ──")
    for table in ["wiki_pages", "wiki_revisions", "wiki_categories", "wiki_links",
                   "crossing_weights", "nspace_kv", "routes",
                   "fractal_manifolds", "fractal_nodes", "fractal_graph_entities",
                   "vectors"]:
        count = neon(f"SELECT COUNT(*) FROM ene.{table}")
        print(f"  {table}: {count}")

    print("\nDone.")


if __name__ == "__main__":
    main()