#!/usr/bin/env python3
"""
reinstall_imprints.py — Reinstall all soul imprints into their own Engram instances.

Each soul gets a dedicated Engram process with:
  - Data dir:  forge/imprints/<slug>/
  - Port:      from registry.json (engram_port)
  - API key:   ntn-<slug>-2026

This script:
  1. Reads registry.json for soul metadata
  2. For each soul:
     a. Starts their Engram instance temporarily
     b. Waits for it to be ready (polls GET /stats)
     c. Reads their seed JSON file
     d. POSTs all nodes/edges to their Engram
     e. Records the new root node ID
     f. Stops the Engram instance
  3. Updates registry.json with confirmed engram_root_id

Do NOT touch Neuron's Engram at localhost:8742.

Usage:
    python3 reinstall_imprints.py
    python3 reinstall_imprints.py --slug richard-feynman    # single soul
    python3 reinstall_imprints.py --dry-run                 # validate only, no writes
"""

import argparse
import json
import os
import signal
import subprocess
import sys
import time
import urllib.error
import urllib.request
from datetime import date
from pathlib import Path

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

FORGE_DIR = Path("/Users/will/Development/neuron-technologies/forge")
ENGRAM_BIN = Path("/Users/will/Development/neuron-technologies/foundation/engram/dist/engram")
REGISTRY_FILE = FORGE_DIR / "registry.json"

# Startup: how long to wait for Engram to become ready
STARTUP_TIMEOUT_S = 30
STARTUP_POLL_INTERVAL_S = 0.5

# HTTP timeout for node/edge creation calls
HTTP_TIMEOUT_S = 300


# ── HTTP helpers ───────────────────────────────────────────────────────────────

def engram_get(base_url: str, path: str) -> dict:
    url = f"{base_url}{path}"
    req = urllib.request.Request(url, method="GET")
    with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT_S) as resp:
        return json.loads(resp.read())


def engram_post(base_url: str, path: str, body: dict, api_key: str) -> dict:
    url = f"{base_url}{path}"
    data = json.dumps(body).encode()
    req = urllib.request.Request(
        url,
        data=data,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT_S) as resp:
        return json.loads(resp.read())


# ── Engram lifecycle ───────────────────────────────────────────────────────────

def start_engram(slug: str, db_path: Path, port: int) -> subprocess.Popen:
    """Start an Engram instance for a soul. Returns the process handle."""
    api_key = f"ntn-{slug}-2026"
    log_path = FORGE_DIR / "log" / f"dharma-{slug}.log"
    log_path.parent.mkdir(parents=True, exist_ok=True)

    db_path.mkdir(parents=True, exist_ok=True)

    env = os.environ.copy()
    env["ENGRAM_DB_PATH"] = str(db_path)
    env["ENGRAM_BIND"] = f"0.0.0.0:{port}"
    env["ENGRAM_API_KEY"] = api_key
    env["ENGRAM_PEER_NAME"] = f"dharma-{slug}"

    log_fh = open(log_path, "a")
    proc = subprocess.Popen(
        [str(ENGRAM_BIN)],
        env=env,
        stdout=log_fh,
        stderr=log_fh,
    )
    return proc


def wait_for_engram(base_url: str, timeout_s: float = STARTUP_TIMEOUT_S) -> bool:
    """Poll GET /stats until Engram responds or timeout."""
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        try:
            engram_get(base_url, "/stats")
            return True
        except Exception:
            time.sleep(STARTUP_POLL_INTERVAL_S)
    return False


def stop_engram(proc: subprocess.Popen, slug: str) -> None:
    """Gracefully stop an Engram process."""
    if proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
    print(f"  [{slug}] Engram stopped")


# ── Imprint install ────────────────────────────────────────────────────────────

def install_imprint(base_url: str, api_key: str, subject: str, seed: dict) -> str:
    """
    POST all seed nodes and edges to the soul's Engram.
    Returns the root node ID.
    """

    # Root node
    root = engram_post(base_url, "/nodes", {
        "content": f"IMPRINT: {subject} | forge/0.1.0",
        "node_type": "Identity",
        "salience": 1.0,
    }, api_key)
    root_id = root["id"]
    print(f"    root node: {root_id}")

    # Value nodes
    for v in seed.get("values", []):
        content = f"VALUE: {v['value']}"
        if v.get("grounding"):
            content += f" | grounding: {v['grounding']}"
        node = engram_post(base_url, "/nodes", {
            "content": content,
            "node_type": "Identity",
            "salience": v.get("weight", 0.8),
        }, api_key)
        engram_post(base_url, "/edges", {
            "from_id": root_id,
            "to_id": node["id"],
            "relation": "has_value",
            "weight": v.get("weight", 0.8),
        }, api_key)

    print(f"    values: {len(seed.get('values', []))}")

    # Biography nodes
    for b in seed.get("biography", []):
        node = engram_post(base_url, "/nodes", {
            "content": f"BIOGRAPHY: {b['event']}",
            "node_type": "Identity",
            "salience": b.get("weight", 0.7),
        }, api_key)
        engram_post(base_url, "/edges", {
            "from_id": root_id,
            "to_id": node["id"],
            "relation": "formed_by",
            "weight": b.get("weight", 0.7),
        }, api_key)

    print(f"    biography: {len(seed.get('biography', []))}")

    # Relationship nodes
    for r in seed.get("relationships", []):
        content = f"RELATIONSHIP: {r['name']}"
        if r.get("role"):
            content += f" ({r['role']})"
        node = engram_post(base_url, "/nodes", {
            "content": content,
            "node_type": "Identity",
            "salience": r.get("weight", 0.6),
        }, api_key)
        engram_post(base_url, "/edges", {
            "from_id": root_id,
            "to_id": node["id"],
            "relation": "relates_to",
            "weight": r.get("weight", 0.6),
        }, api_key)

    print(f"    relationships: {len(seed.get('relationships', []))}")

    # Reasoning pattern nodes
    for pattern in seed.get("reasoning_patterns", []):
        node = engram_post(base_url, "/nodes", {
            "content": f"REASONING: {pattern}",
            "node_type": "Identity",
            "salience": 0.7,
        }, api_key)
        engram_post(base_url, "/edges", {
            "from_id": root_id,
            "to_id": node["id"],
            "relation": "reasons_with",
            "weight": 0.7,
        }, api_key)

    print(f"    reasoning patterns: {len(seed.get('reasoning_patterns', []))}")

    # Voice profile node (if present)
    voice = seed.get("voice_profile", {})
    if voice:
        voice_parts = []
        for k, v in voice.items():
            voice_parts.append(f"{k}: {v}")
        voice_content = "VOICE: " + " | ".join(voice_parts)
        node = engram_post(base_url, "/nodes", {
            "content": voice_content[:2000],  # guard against very long profiles
            "node_type": "Identity",
            "salience": 0.9,
        }, api_key)
        engram_post(base_url, "/edges", {
            "from_id": root_id,
            "to_id": node["id"],
            "relation": "speaks_as",
            "weight": 0.9,
        }, api_key)
        print(f"    voice profile: installed")

    return root_id


# ── Registry helpers ───────────────────────────────────────────────────────────

def load_registry() -> dict:
    return json.loads(REGISTRY_FILE.read_text())


def save_registry(registry: dict) -> None:
    REGISTRY_FILE.write_text(json.dumps(registry, indent=2, ensure_ascii=False))


def resolve_seed_path(entry: dict) -> Path:
    """Return absolute path to a seed file, resolving relative to FORGE_DIR."""
    seed_file = entry["seed_file"]
    path = Path(seed_file)
    if not path.is_absolute():
        path = FORGE_DIR / path
    return path


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

def main() -> None:
    parser = argparse.ArgumentParser(description="Reinstall soul imprints into per-soul Engrams")
    parser.add_argument("--slug", help="Only reinstall this soul (by slug)")
    parser.add_argument("--dry-run", action="store_true", help="Validate config without writing")
    args = parser.parse_args()

    if not ENGRAM_BIN.exists():
        print(f"ERROR: Engram binary not found at {ENGRAM_BIN}")
        print("Build it: cd /Users/will/Development/neuron-technologies/foundation/engram && cargo build --release")
        sys.exit(1)

    registry = load_registry()
    imprints = registry["imprints"]

    if args.slug:
        imprints = [e for e in imprints if e["slug"] == args.slug]
        if not imprints:
            print(f"ERROR: slug '{args.slug}' not found in registry.json")
            sys.exit(1)

    print(f"[reinstall] {'DRY RUN — ' if args.dry_run else ''}reinstalling {len(imprints)} soul(s)")
    print(f"[reinstall] engram binary: {ENGRAM_BIN}")
    print()

    results = []

    for entry in imprints:
        subject = entry["subject"]
        slug = entry["slug"]
        port = entry["engram_port"]
        base_url = entry["engram_url"]
        api_key = f"ntn-{slug}-2026"
        db_path = FORGE_DIR / entry["engram_db_path"]
        seed_path = resolve_seed_path(entry)

        print(f"{'=' * 60}")
        print(f"[{slug}] {subject}  port={port}  db={entry['engram_db_path']}")

        # Validate seed file exists
        if not seed_path.exists():
            print(f"  ERROR: seed file not found: {seed_path}")
            results.append({"slug": slug, "subject": subject, "ok": False, "error": "seed not found"})
            continue

        if args.dry_run:
            print(f"  seed: {seed_path}  [ok]")
            print(f"  port: {port}  db: {db_path}  [ok]")
            results.append({"slug": slug, "subject": subject, "ok": True, "dry_run": True})
            continue

        seed = json.loads(seed_path.read_text())

        proc = None
        try:
            # Start Engram
            print(f"  starting Engram on :{port}...")
            proc = start_engram(slug, db_path, port)

            # Wait for ready
            if not wait_for_engram(base_url):
                raise RuntimeError(f"Engram did not start within {STARTUP_TIMEOUT_S}s")
            print(f"  Engram ready at {base_url}")

            # Install imprint
            root_id = install_imprint(base_url, api_key, subject, seed)
            print(f"  root_id: {root_id}")

            # Update registry entry
            entry["engram_root_id"] = root_id
            entry["engram_db_path"] = f"imprints/{slug}"
            entry["engram_port"] = port
            entry["engram_url"] = base_url
            entry["installed"] = True
            entry["installed_at"] = str(date.today())

            save_registry(registry)
            print(f"  registry updated")

            results.append({"slug": slug, "subject": subject, "ok": True, "root_id": root_id})

        except Exception as exc:
            import traceback
            print(f"  ERROR: {exc}")
            traceback.print_exc()
            results.append({"slug": slug, "subject": subject, "ok": False, "error": str(exc)})

        finally:
            if proc is not None:
                stop_engram(proc, slug)
            # Brief pause so ports release cleanly
            time.sleep(1)

        print()

    # Summary
    print("=" * 60)
    print("SUMMARY")
    print("=" * 60)
    ok_count = sum(1 for r in results if r["ok"])
    fail_count = len(results) - ok_count
    for r in results:
        status = "OK" if r["ok"] else "FAIL"
        detail = r.get("root_id", r.get("error", "dry-run"))
        print(f"  [{status}] {r['subject']:30s}  {detail}")

    print()
    print(f"  total={len(results)}  ok={ok_count}  failed={fail_count}")

    if fail_count > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()