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

Build a unified map of every data point / concept in the project:
  - Live PostgreSQL databases on neon-64gb (schemas, tables, row counts)
  - Live Gremlin graph on Azure Cosmos DB (vertex/edge label counts)
  - Local repository source files (Lean, Rust, Python, JS/TS, Markdown, configs, etc.)
  - Curated extraction/ideas.json and extraction/math.json

Outputs:
  extraction/data_inventory.json      - database + repo file summary
  extraction/all_concepts_merged.json - unified concept list (one row per concept)

Usage:
  uv run scripts/inventory_everything.py [--skip-db]
"""

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

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "extraction"
OUT.mkdir(exist_ok=True)

SKIP_DIRS = {
    ".git", "node_modules", "__pycache__", "venv", ".venv",
    ".lake", "target", "lean_binned", "scratch", "shared-data",
    "archive", "result", "result-devcontainer", "extraction",
    ".agents", ".claude", ".consolidation-manifests", ".contextstream",
    ".cursor", ".devcontainer", ".devin", ".github", ".headroom",
    ".hermes", ".kilo", ".kilocode", ".opencode", ".roo", ".vscode",
    ".windsurf",
}

# Files we know how to turn into named concepts
SCAN_SUFFIXES = {
    ".lean", ".rs", ".py", ".js", ".ts", ".mjs", ".cjs",
    ".md", ".json", ".yaml", ".yml", ".toml", ".agda",
    ".v", ".sv", ".c", ".h", ".sh", ".tid", ".txt",
    ".cff", ".bib", ".tex",
}

MAX_FILE_BYTES = 5 * 1024 * 1024

LAYER_MAP = {
    "0-Core-Formalism": "core",
    "1-Distributed-Systems": "distributed",
    "2-Search-Space": "search",
    "3-Mathematical-Models": "models",
    "4-Infrastructure": "infrastructure",
    "5-Applications": "applications",
    "6-Documentation": "documentation",
    "scripts": "scripts",
}

MATH_TYPES = {
    "theorem", "definition", "structure", "type", "predicate",
    "operator", "invariant", "constant", "equation", "algorithm",
    "dataset", "model",
}


def short_hash(s: str, n: int = 10) -> str:
    return hashlib.sha1(s.encode()).hexdigest()[:n]


def layer_tag(rel: Path) -> str:
    if len(rel.parts) <= 1:
        return "root"
    top = rel.parts[0]
    return LAYER_MAP.get(top, top)


def safe_id(prefix: str, *parts: str) -> str:
    base = "_".join(str(p) for p in parts)
    base = re.sub(r"[^a-zA-Z0-9_]+", "_", base).strip("_").lower()
    return f"{prefix}_{base}"[:120]


def read_text_limited(path: Path) -> str:
    try:
        size = path.stat().st_size
        if size > MAX_FILE_BYTES:
            return ""
        return path.read_text(errors="replace")
    except Exception:
        return ""


# ── Scanners ─────────────────────────────────────────────────────────────────

LEAN_TOP_RE = re.compile(
    r"^(?:@\[.*?\]\s*)*(?:noncomputable\s+|partial\s+|private\s+|protected\s+)*"
    r"(theorem|lemma|def|structure|inductive|class|abbrev|instance|opaque|axiom)"
    r"\s+(\w+)",
    re.MULTILINE,
)
LEAN_DOC_RE = re.compile(r"/-(.*?)-/", re.DOTALL)
LEAN_LINE_DOC_RE = re.compile(r"--\s*(.*)")


def nearest_doc(text: str, pos: int) -> str:
    """Return the nearest preceding doc comment (block or line) summary."""
    best = ""
    for m in LEAN_DOC_RE.finditer(text):
        if m.end() <= pos:
            lines = [l.strip() for l in m.group(1).strip().splitlines() if l.strip()]
            best = next((l for l in lines if not l.startswith("-") and len(l) > 5), best)
    if not best:
        # fall back to contiguous line comments immediately above
        lines = []
        prev = text.rfind("\n", 0, pos)
        while prev > 0:
            line_start = text.rfind("\n", 0, prev) + 1
            line = text[line_start:prev].strip()
            m = LEAN_LINE_DOC_RE.match(line)
            if m:
                lines.insert(0, m.group(1).strip())
            elif not line:
                break
            else:
                break
            prev = line_start - 1
        best = " ".join(lines)
    return best[:200]


def scan_lean(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    ns = str(rel.with_suffix("")).replace("/", ".")
    results = []
    for m in LEAN_TOP_RE.finditer(text):
        kind, name = m.group(1), m.group(2)
        ctx = text[max(0, m.start() - 80) : m.start() + 200].lower()
        status = "SORRY" if kind == "theorem" and "sorry" in ctx else "DEFINED"
        if kind in ("theorem", "lemma") and status != "SORRY":
            status = "DEFINED"
        desc = nearest_doc(text, m.start()) or f"{kind} {name} in {rel}"
        h = short_hash(f"{rel}:{name}")
        results.append({
            "id": safe_id("math", name, h),
            "type": kind if kind != "abbrev" else "type",
            "name": name,
            "expression": name,
            "latex": name,
            "description": desc,
            "source_file": str(rel),
            "lean_name": f"{ns}.{name}",
            "lean_status": status,
            "tags": [layer_tag(rel), "lean", ns.split(".")[0] if ns else "lean"],
            "related": [],
        })
    return results


RUST_TOP_RE = re.compile(
    r"^(?:pub\s+|crate\s+|async\s+)*"
    r"(fn|struct|enum|trait|impl|type|const|static|mod)\s+(?:<[^>]+>\s*)?(\w+)",
    re.MULTILINE,
)


def scan_rust(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    results = []
    for m in RUST_TOP_RE.finditer(text):
        kind, name = m.group(1), m.group(2)
        h = short_hash(f"{rel}:{name}")
        results.append({
            "id": safe_id("rust", name, h),
            "type": "function" if kind == "fn" else ("structure" if kind in ("struct", "enum") else "definition"),
            "name": name,
            "expression": name,
            "latex": name,
            "description": f"{kind} {name} in {rel}",
            "source_file": str(rel),
            "lean_name": "",
            "lean_status": "NOT_IN_LEAN",
            "tags": [layer_tag(rel), "rust", kind],
            "related": [],
        })
    return results


PY_TOP_RE = re.compile(r"^(?:async\s+)?(def|class)\s+(\w+)", re.MULTILINE)


def scan_python(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    first = text.lstrip()
    file_comment = ""
    if first.startswith("#"):
        file_comment = first.split("\n", 1)[0][:160]
    results = []
    for m in PY_TOP_RE.finditer(text):
        kind, name = m.group(1), m.group(2)
        h = short_hash(f"{rel}:{name}")
        results.append({
            "id": safe_id("py", name, h),
            "type": "function" if kind == "def" else "structure",
            "name": name,
            "expression": name,
            "latex": name,
            "description": file_comment or f"{kind} {name} in {rel}",
            "source_file": str(rel),
            "lean_name": "",
            "lean_status": "NOT_IN_LEAN",
            "tags": [layer_tag(rel), "python", kind],
            "related": [],
        })
    return results


JS_TOP_RE = re.compile(
    r"^(?:export\s+(?:default\s+)?)?(?:async\s+)?(?:function\s+|class\s+|interface\s+|type\s+|enum\s+)(\w+)",
    re.MULTILINE,
)
JS_CONST_RE = re.compile(r"^(?:export\s+)?(?:const|let|var)\s+(\w+)\s*[:=]", re.MULTILINE)


def scan_jsts(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    results = []
    for m in JS_TOP_RE.finditer(text):
        name = m.group(1)
        h = short_hash(f"{rel}:{name}")
        results.append({
            "id": safe_id("js", name, h),
            "type": "function",
            "name": name,
            "expression": name,
            "latex": name,
            "description": f"{m.group(0).strip()} in {rel}",
            "source_file": str(rel),
            "lean_name": "",
            "lean_status": "NOT_IN_LEAN",
            "tags": [layer_tag(rel), "js_ts"],
            "related": [],
        })
    return results[:200]


MD_HEADER_RE = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)


def scan_markdown(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    stem = path.stem.replace("-", "_")
    results = [{
        "id": safe_id("doc", stem, short_hash(str(rel))),
        "type": "concept",
        "name": f"Doc: {path.name}",
        "expression": "",
        "latex": "",
        "description": text[:300].replace("\n", " "),
        "source_file": str(rel),
        "lean_name": "",
        "lean_status": "IDEA",
        "tags": [layer_tag(rel), "markdown", "doc"],
        "related": [],
    }]
    for m in MD_HEADER_RE.finditer(text):
        title = m.group(2).strip()
        if len(title) > 4:
            h = short_hash(f"{rel}:{title}")
            results.append({
                "id": safe_id("doc", re.sub(r"[^a-zA-Z0-9_]", "_", title.lower())[:40], h),
                "type": "concept",
                "name": title,
                "expression": "",
                "latex": "",
                "description": text[m.end():m.end()+250].replace("\n", " "),
                "source_file": str(rel),
                "lean_name": "",
                "lean_status": "IDEA",
                "tags": [layer_tag(rel), "markdown", "section"],
                "related": [],
            })
    return results[:80]


def scan_json(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text or len(text) > 500_000:
        return []
    try:
        data = json.loads(text)
    except Exception:
        return []
    keys = []
    if isinstance(data, dict):
        keys = list(data.keys())[:30]
    h = short_hash(str(rel))
    return [{
        "id": safe_id("json", path.stem, h),
        "type": "concept",
        "name": f"JSON: {path.name}",
        "expression": "",
        "latex": "",
        "description": f"Keys: {', '.join(keys[:10])}{'...' if len(keys) > 10 else ''}",
        "source_file": str(rel),
        "lean_name": "",
        "lean_status": "IDEA",
        "tags": [layer_tag(rel), "json"],
        "related": [],
    }]


YAML_TOP_RE = re.compile(r"^(\w+):", re.MULTILINE)


def scan_yaml(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    keys = YAML_TOP_RE.findall(text)[:20]
    h = short_hash(str(rel))
    return [{
        "id": safe_id("yaml", path.stem, h),
        "type": "concept",
        "name": f"YAML: {path.name}",
        "expression": "",
        "latex": "",
        "description": f"Top keys: {', '.join(keys[:10])}{'...' if len(keys) > 10 else ''}",
        "source_file": str(rel),
        "lean_name": "",
        "lean_status": "IDEA",
        "tags": [layer_tag(rel), "yaml"],
        "related": [],
    }]


def scan_toml(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    sections = re.findall(r"^\[(.+?)\]", text, re.MULTILINE)[:20]
    h = short_hash(str(rel))
    return [{
        "id": safe_id("toml", path.stem, h),
        "type": "concept",
        "name": f"TOML: {path.name}",
        "expression": "",
        "latex": "",
        "description": f"Sections: {', '.join(sections[:10])}{'...' if len(sections) > 10 else ''}",
        "source_file": str(rel),
        "lean_name": "",
        "lean_status": "IDEA",
        "tags": [layer_tag(rel), "toml"],
        "related": [],
    }]


VERILOG_MOD_RE = re.compile(r"^\s*module\s+(\w+)", re.MULTILINE)


def scan_verilog(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    results = []
    for m in VERILOG_MOD_RE.finditer(text):
        name = m.group(1)
        h = short_hash(f"{rel}:{name}")
        results.append({
            "id": safe_id("verilog", name, h),
            "type": "structure",
            "name": name,
            "expression": name,
            "latex": name,
            "description": f"Verilog module {name} in {rel}",
            "source_file": str(rel),
            "lean_name": "",
            "lean_status": "NOT_IN_LEAN",
            "tags": [layer_tag(rel), "verilog"],
            "related": [],
        })
    return results


AGDA_TOP_RE = re.compile(r"^(\w+)\s*:", re.MULTILINE)


def scan_agda(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    names = []
    for m in AGDA_TOP_RE.finditer(text):
        name = m.group(1)
        if name[0].islower():
            continue
        names.append(name)
        if len(names) >= 50:
            break
    h = short_hash(str(rel))
    return [{
        "id": safe_id("agda", path.stem, h),
        "type": "definition",
        "name": path.stem,
        "expression": ", ".join(names[:10]),
        "latex": "",
        "description": f"Agda definitions: {', '.join(names[:10])}{'...' if len(names) > 10 else ''}",
        "source_file": str(rel),
        "lean_name": "",
        "lean_status": "NOT_IN_LEAN",
        "tags": [layer_tag(rel), "agda"],
        "related": [],
    }]


def scan_shell(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    funcs = re.findall(r"^\s*(\w+)\s*\(\)", text, re.MULTILINE)[:20]
    h = short_hash(str(rel))
    desc = text.split("\n")[0][:160] if text else ""
    return [{
        "id": safe_id("sh", path.stem, h),
        "type": "concept",
        "name": f"Shell: {path.name}",
        "expression": "",
        "latex": "",
        "description": f"{desc} Functions: {', '.join(funcs[:10])}",
        "source_file": str(rel),
        "lean_name": "",
        "lean_status": "IDEA",
        "tags": [layer_tag(rel), "shell"],
        "related": [],
    }]


def scan_text(path: Path, rel: Path) -> list:
    text = read_text_limited(path)
    if not text:
        return []
    h = short_hash(str(rel))
    title = ""
    if path.suffix == ".tid":
        m = re.search(r"^title:\s*(.+)$", text, re.MULTILINE)
        if m:
            title = m.group(1).strip()
    return [{
        "id": safe_id("file", path.stem, h),
        "type": "concept",
        "name": title or f"File: {path.name}",
        "expression": "",
        "latex": "",
        "description": text[:250].replace("\n", " "),
        "source_file": str(rel),
        "lean_name": "",
        "lean_status": "IDEA",
        "tags": [layer_tag(rel), path.suffix.lstrip(".")],
        "related": [],
    }]


SCANNERS = {
    ".lean": scan_lean,
    ".rs": scan_rust,
    ".py": scan_python,
    ".js": scan_jsts,
    ".ts": scan_jsts,
    ".mjs": scan_jsts,
    ".cjs": scan_jsts,
    ".md": scan_markdown,
    ".json": scan_json,
    ".yaml": scan_yaml,
    ".yml": scan_yaml,
    ".toml": scan_toml,
    ".v": scan_verilog,
    ".sv": scan_verilog,
    ".agda": scan_agda,
    ".sh": scan_shell,
    ".tid": scan_text,
    ".txt": scan_text,
    ".cff": scan_text,
    ".bib": scan_text,
    ".tex": scan_text,
    ".c": scan_text,
    ".h": scan_text,
}


def should_skip_dir(p: Path) -> bool:
    if p.name in SKIP_DIRS:
        return True
    if p.name.startswith(".") and p.name != ".github":
        return True
    return False


def walk_sources(base: Path):
    stack = [base]
    while stack:
        p = stack.pop()
        if p.is_dir():
            if should_skip_dir(p):
                continue
            try:
                for c in sorted(p.iterdir(), key=lambda x: x.name):
                    stack.append(c)
            except PermissionError:
                continue
        elif p.is_file() and p.suffix.lower() in SCAN_SUFFIXES:
            yield p


# ── Database inventory ───────────────────────────────────────────────────────

def psql_json(db: str, sql: str) -> list:
    cmd = [
        "ssh", "neon-64gb",
        f"podman exec -i arxiv-pg psql -U postgres -d {db} -t -A -F '|' -c \"{sql}\""
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        print(f"  [db {db}] error: {proc.stderr.strip()}", file=sys.stderr)
        return []
    out = []
    for line in proc.stdout.splitlines():
        if not line.strip():
            continue
        out.append(line.split("|"))
    return out


def inventory_postgres() -> dict:
    result = {"host": "neon-64gb", "databases": {}}
    dbs = psql_json(
        "postgres",
        "SELECT datname, pg_size_pretty(pg_database_size(datname)) "
        "FROM pg_database WHERE datname NOT IN ('template0','template1','postgres') "
        "ORDER BY pg_database_size(datname) DESC"
    )
    for row in dbs:
        db_name, size = row[0], row[1]
        tables = psql_json(
            db_name,
            "SELECT schemaname, relname, n_live_tup, "
            "pg_size_pretty(pg_total_relation_size(schemaname||'.'||relname)) "
            "FROM pg_stat_user_tables "
            "WHERE schemaname NOT IN ('pg_catalog','information_schema') "
            "ORDER BY n_live_tup DESC"
        )
        result["databases"][db_name] = {
            "size": size,
            "tables": [
                {
                    "schema": t[0],
                    "table": t[1],
                    "rows": int(t[2]) if t[2].isdigit() else 0,
                    "size": t[3],
                }
                for t in tables
            ],
        }
    return result


def load_gremlin_creds() -> dict:
    env = {}
    creds = ROOT / ".env.gremlin"
    if creds.exists():
        for line in creds.read_text().splitlines():
            if "=" in line and not line.startswith("#"):
                k, _, v = line.partition("=")
                env[k.strip()] = v.strip().strip('"')
    infra = json.loads((ROOT / "infra.json").read_text())
    return {
        "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"),
    }


def inventory_gremlin() -> dict:
    try:
        from gremlin_python.driver import client as gremlin_client
        from gremlin_python.driver.serializer import GraphSONSerializersV2d0
    except ImportError as e:
        return {"error": f"gremlinpython not installed: {e}"}
    creds = load_gremlin_creds()
    gc = gremlin_client.Client(
        creds["endpoint"], "g",
        username=f"/dbs/{creds['database']}/colls/{creds['graph']}",
        password=creds["key"],
        message_serializer=GraphSONSerializersV2d0(),
    )
    try:
        v_total = len(gc.submit("g.V().limit(100000).id()").all().result())
        e_total = len(gc.submit("g.E().limit(100000).id()").all().result())
        v_labels = gc.submit("g.V().groupCount().by(label)").all().result()
        e_labels = gc.submit("g.E().groupCount().by(label)").all().result()
        return {
            "endpoint": creds["endpoint"],
            "graph": creds["graph"],
            "vertices": v_total,
            "edges": e_total,
            "vertex_labels": v_labels[0] if v_labels else {},
            "edge_labels": e_labels[0] if e_labels else {},
        }
    except Exception as e:
        return {"error": str(e)}
    finally:
        gc.close()


# ── Curated files ────────────────────────────────────────────────────────────

def load_curated() -> tuple[list, list]:
    ideas, math = [], []
    ideas_file = ROOT / "extraction" / "ideas.json"
    math_file = ROOT / "extraction" / "math.json"
    if ideas_file.exists():
        ideas = json.loads(ideas_file.read_text()).get("ideas", [])
    if math_file.exists():
        math = json.loads(math_file.read_text()).get("math", [])
    return ideas, math


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

def main():
    ap = argparse.ArgumentParser(description="Inventory every data point in the project")
    ap.add_argument("--skip-db", action="store_true", help="Skip live database inventory")
    ap.add_argument("--skip-repo", action="store_true", help="Skip local repo scan")
    args = ap.parse_args()

    generated_at = datetime.now(timezone.utc).isoformat()
    data_inventory = {
        "schema": "research_stack_data_inventory_v1",
        "generated_at": generated_at,
        "postgres": None,
        "gremlin": None,
    }

    if not args.skip_db:
        print("[1/3] Inventorying PostgreSQL...", flush=True)
        data_inventory["postgres"] = inventory_postgres()
        print("  PostgreSQL done.", flush=True)

        print("[2/3] Inventorying Gremlin...", flush=True)
        data_inventory["gremlin"] = inventory_gremlin()
        print("  Gremlin done.", flush=True)
    else:
        print("[1-2/3] Skipping database inventory.", flush=True)

    all_concepts = []
    file_summary = {"total_files": 0, "by_extension": {}, "by_layer": {}}

    if not args.skip_repo:
        print("[3/3] Scanning local repository...", flush=True)
        curated_ideas, curated_math = load_curated()
        all_concepts.extend(curated_math)
        all_concepts.extend({
            **idea,
            "type": "concept",
            "tags": list(set(idea.get("tags", []) + ["curated", "idea"])),
        } for idea in curated_ideas)

        for path in walk_sources(ROOT):
            rel = path.relative_to(ROOT)
            file_summary["total_files"] += 1
            ext = path.suffix.lower() or "(none)"
            file_summary["by_extension"][ext] = file_summary["by_extension"].get(ext, 0) + 1
            layer = layer_tag(rel)
            file_summary["by_layer"][layer] = file_summary["by_layer"].get(layer, 0) + 1

            scanner = SCANNERS.get(ext)
            if not scanner:
                continue
            try:
                concepts = scanner(path, rel)
            except Exception as exc:
                print(f"  [skip] {rel}: {exc}", file=sys.stderr)
                continue
            all_concepts.extend(concepts)

        print(f"  Scanned {file_summary['total_files']} files, {len(all_concepts)} concepts.", flush=True)
    else:
        print("[3/3] Skipping local repo scan.", flush=True)

    # Deduplicate by id, keeping first occurrence
    seen = set()
    deduped = []
    for c in all_concepts:
        cid = c.get("id")
        if not cid or cid in seen:
            continue
        seen.add(cid)
        deduped.append(c)

    by_type = {}
    by_layer = {}
    for c in deduped:
        by_type[c.get("type", "unknown")] = by_type.get(c.get("type", "unknown"), 0) + 1
        layer = layer_tag(Path(c.get("source_file", "")))
        by_layer[layer] = by_layer.get(layer, 0) + 1

    data_inventory["repo_files"] = file_summary
    data_inventory["concept_count"] = len(deduped)
    data_inventory["concept_summary"] = {"by_type": by_type, "by_layer": by_layer}

    (OUT / "data_inventory.json").write_text(json.dumps(data_inventory, indent=2))
    (OUT / "all_concepts_merged.json").write_text(json.dumps({"math": deduped}, indent=2))

    print("\nWrote:")
    print(f"  extraction/data_inventory.json")
    print(f"  extraction/all_concepts_merged.json ({len(deduped)} concepts)")
    print("Done.")


if __name__ == "__main__":
    main()