#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "gremlinpython",
#   "python-dotenv",
#   "requests",
# ]
# ///
"""gremlin_lean_report.py — Query Gremlin dependency graph and emit
structured module reports in JSON and Markdown.

Optional: push to AppFloyo Cloud via its REST API using a GoTrue JWT.

Flow:
  Gremlin (46K vertices, 30K edges)
       │
       │  Query all module vertices + 'imports' edges
       ▼
  Build module dependency map
       │
       ├─► JSON report (machine-readable)
       ├─► Markdown report (human-readable, importable into AppFloyo)
       └─► AppFloyo workspace (via --push flag)
"""

from __future__ import annotations

import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import requests

ROOT = Path(__file__).resolve().parents[2]


# ── Gremlin connection ──────────────────────────────────────────────────

def get_gremlin_client():
    """Create Gremlin client from .env.gremlin."""
    env_path = ROOT / ".env.gremlin"
    if not env_path.exists():
        print("ERROR: .env.gremlin not found", file=sys.stderr)
        sys.exit(1)

    lines = env_path.read_text().splitlines()
    cfg = {}
    for line in lines:
        if "=" in line and not line.startswith("#"):
            k, v = line.split("=", 1)
            cfg[k.strip()] = v.strip().strip('"')

    from gremlin_python.driver import client as gremlin_client, serializer
    return gremlin_client.Client(
        cfg.get("GREMLIN_ENDPOINT", "wss://mathblob.gremlin.cosmos.azure.com:443/"),
        "g",
        username=cfg.get("GREMLIN_USERNAME", "/dbs/research/colls/concepts"),
        password=cfg.get("GREMLIN_PASSWORD", ""),
        message_serializer=serializer.GraphSONSerializersV2d0(),
    )


def query_gremlin(query: str, bindings: dict | None = None) -> list[dict]:
    """Run a Gremlin query and return results."""
    c = get_gremlin_client()
    try:
        result = c.submitAsync(query, bindings or {}).result().all().result()
        return list(result)
    finally:
        c.close()


def get_all_modules() -> list[dict]:
    """Get all module vertices with theorem/sorry counts."""
    raw = query_gremlin(
        'g.V().hasLabel("module").project("id","ns","tc","sc","fc")'
        '.by(id).by(coalesce(values("namespace"), constant("")))'
        '.by(coalesce(values("theorem_count"), constant(0)))'
        '.by(coalesce(values("sorry_count"), constant(0)))'
        '.by(coalesce(values("file_path"), constant("")))'
        '.fold()'
    )
    items = raw[0] if raw and isinstance(raw[0], list) else raw
    modules = []
    for r in items if isinstance(items, list) else []:
        if isinstance(r, dict):
            modules.append({
                "id": r.get("id", ""),
                "name": r.get("ns", r["id"].split(".")[-1] if "." in r.get("id", "") else r.get("id", "")),
                "theorem_count": int(r.get("tc", 0)),
                "sorry_count": int(r.get("sc", 0)),
                "file_path": r.get("fc", ""),
            })
    return modules


def get_dependency_map() -> dict[str, list[str]]:
    """Get module dependency graph via 'imports' edges."""
    raw = query_gremlin(
        'g.V().hasLabel("module").project("id","deps")'
        '.by(id).by(out("imports").hasLabel("module").id().fold())'
        '.fold()'
    )
    items = raw[0] if raw and isinstance(raw[0], list) else raw
    result: dict[str, list[str]] = {}
    for r in items if isinstance(items, list) else []:
        if isinstance(r, dict):
            result[r["id"]] = list(r.get("deps", []))
    return result


# ── Report generation ──────────────────────────────────────────────────

def build_report(modules: list[dict], deps: dict[str, list[str]]) -> dict:
    """Build a structured report from Gremlin data."""
    rev_deps: dict[str, int] = {}
    for mod_id, dep_ids in deps.items():
        for dep_id in dep_ids:
            rev_deps[dep_id] = rev_deps.get(dep_id, 0) + 1

    rows = []
    for m in modules:
        mid = m["id"]
        rows.append({
            "id": mid,
            "name": m["name"],
            "theorems": m["theorem_count"],
            "sorries": m["sorry_count"],
            "imports": len(deps.get(mid, [])),
            "imported_by": rev_deps.get(mid, 0),
        })

    rows.sort(key=lambda r: -r["theorems"])

    total_thm = sum(r["theorems"] for r in rows)
    total_sor = sum(r["sorries"] for r in rows)

    return {
        "schema": "gremlin_lean_report_v1",
        "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "source": "Gremlin (Azure Cosmos DB, mathblob)",
        "summary": {
            "total_modules": len(rows),
            "total_theorems": total_thm,
            "total_sorries": total_sor,
            "total_dependency_edges": sum(len(v) for v in deps.values()),
        },
        "by_theorems": rows[:10],
        "most_imported": sorted(rows, key=lambda r: -r["imported_by"])[:10],
        "all_modules": rows,
    }


def emit_json(report: dict, out_path: Path) -> None:
    """Write JSON report."""
    out_path.write_text(json.dumps(report, indent=2) + "\n")
    print(f"  JSON: {out_path}")


def emit_markdown(report: dict, out_path: Path) -> None:
    """Write Markdown report (importable into AppFlowy/Obsidian)."""
    lines = [
        f"# Lean Module Dependency Report",
        f"",
        f"Generated: {report['generated_at']}",
        f"Source: {report['source']}",
        f"",
        f"## Summary",
        f"",
        f"| Metric | Value |",
        f"|--------|-------|",
        f"| Modules | {report['summary']['total_modules']} |",
        f"| Theorems | {report['summary']['total_theorems']} |",
        f"| Sorries | {report['summary']['total_sorries']} |",
        f"| Dependency edges | {report['summary']['total_dependency_edges']} |",
        f"",
        f"## Top 10 Modules by Theorem Count",
        f"",
        f"| # | Module | Theorems | Sorries | Imports | Imported By |",
        f"|---|--------|----------|---------|---------|-------------|",
    ]
    for i, r in enumerate(report["by_theorems"], 1):
        lines.append(f"| {i} | {r['name']} | {r['theorems']} | {r['sorries']} | {r['imports']} | {r['imported_by']} |")

    lines += [
        f"",
        f"## Top 10 Most-Imported Modules",
        f"",
        f"| # | Module | Imported By | Theorems | Sorries |",
        f"|---|--------|-------------|----------|---------|",
    ]
    for i, r in enumerate(report["most_imported"], 1):
        lines.append(f"| {i} | {r['name']} | {r['imported_by']} | {r['theorems']} | {r['sorries']} |")

    lines += [
        f"",
        f"## All Modules ({len(report['all_modules'])})",
        f"",
        f"| Module | Theorems | Sorries | Imports | Imported By |",
        f"|--------|----------|---------|---------|-------------|",
    ]
    for r in report["all_modules"]:
        lines.append(f"| {r['name']} | {r['theorems']} | {r['sorries']} | {r['imports']} | {r['imported_by']} |")

    out_path.write_text("\n".join(lines) + "\n")
    print(f"  MD:  {out_path}")


# ── AppFloyo push ─────────────────────────────────────────────────────

def push_to_appflowy(report: dict) -> None:
    """Push top modules to AppFloyo Cloud workspace."""
    url = os.environ.get("APPFLOWY_URL", "http://100.92.88.64:8000")
    token = os.environ.get("APPFLOWY_TOKEN", "")

    # If no token, try to get one from GoTrue
    if not token:
        gotrue_url = os.environ.get("GOTRUE_URL", "http://100.92.88.64:9999")
        email = os.environ.get("GOTRUE_ADMIN_EMAIL", "admin@researchstack.info")
        password = os.environ.get("GOTRUE_ADMIN_PASSWORD", "admin123")
        try:
            r = requests.post(f"{gotrue_url}/token?grant_type=password",
                              json={"email": email, "password": password}, timeout=10)
            r.raise_for_status()
            token = r.json().get("access_token", "")
            print(f"  Got GoTrue JWT ({len(token)} chars)")
        except Exception as e:
            print(f"  Failed to get GoTrue token: {e}", file=sys.stderr)
            return

    # Get workspace list
    try:
        headers = {"Authorization": f"Bearer {token}"}
        r = requests.get(f"{url}/api/workspace", headers=headers, timeout=10)
        r.raise_for_status()
        workspaces = r.json().get("data", [])
        print(f"  Workspaces: {len(workspaces)}")
    except Exception as e:
        print(f"  Failed to get workspace: {e}", file=sys.stderr)
        return



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

def main():
    out_dir = ROOT / "reports"
    out_dir.mkdir(parents=True, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")

    print("[gremlin] Querying Gremlin...")
    modules = get_all_modules()
    print(f"  Found {len(modules)} module vertices")

    deps = get_dependency_map()
    edge_count = sum(len(v) for v in deps.values())
    print(f"  Loaded {edge_count} import edges")

    print("[gremlin] Building report...")
    report = build_report(modules, deps)

    s = report["summary"]
    print(f"\n  {s['total_modules']} modules, {s['total_theorems']} theorems, {s['total_sorries']} sorries")
    print(f"  {s['total_dependency_edges']} dependency edges")

    print(f"\n  Top 5 by theorem count:")
    for r in report["by_theorems"][:5]:
        print(f"    {r['theorems']:4d} thm  {r['sorries']:3d} sor  "
              f"imp={r['imports']:2d}  dep={r['imported_by']:3d}  | {r['name'][:50]}")

    print(f"\n  Top 5 most-imported:")
    for r in report["most_imported"][:5]:
        print(f"    {r['imported_by']:3d} dependents  | {r['name'][:55]}")

    json_path = out_dir / f"lean_report_{ts}.json"
    md_path = out_dir / f"lean_report_{ts}.md"
    emit_json(report, json_path)
    emit_markdown(report, md_path)

    print(f"\n[gremlin] Reports written to {out_dir}/")
    if os.environ.get("PUSH_APPFLOWY"):
        print("\n[gremlin] Pushing to AppFloyo...")
        push_to_appflowy(report)

    print(f"  To re-push: PUSH_APPFLOWY=1 python3 gremlin_lean_report.py")


if __name__ == "__main__":
    main()