#!/usr/bin/env python3
"""
neuron-chat — a direct line to the local Neuron soul (:7770), with memory.

You type, Neuron answers. No Claude in the middle.

Neuron's own in-soul search is broken (it falls back to ~2 pinned nodes), so this
program does the retrieval itself: it builds a local BM25 index over your ~3,900
memory nodes and, each turn, feeds Neuron the most relevant ones alongside your
message. That gives it real access to its graph instead of the "light version".

Run from Terminal:  neuron     (or:  python3 ~/neuron-chat.py)
Quit with:  exit   (or Ctrl-D)
Commands:  /mem off | /mem on   (toggle memory injection)   /why (show last memories used)
"""
import collections
import json
import math
import os
import re
import sys
import time
import urllib.request

SOUL = "http://127.0.0.1:7770"
SNAP = os.path.expanduser("~/.neuron/engram/snapshot.json")
SESSION = f"cli-{int(time.time())}"
TOPK = 6                  # memories injected per turn
MAX_NODE_CHARS = 600      # truncate each memory

C = sys.stdout.isatty()
DIM   = "\033[2m"  if C else ""
BOLD  = "\033[1m"  if C else ""
CYAN  = "\033[36m" if C else ""
GREEN = "\033[32m" if C else ""
RESET = "\033[0m"  if C else ""


# ── local BM25 index over the memory snapshot ──────────────────────────────
def _toks(s):
    return re.findall(r"[a-z0-9]+", (s or "").lower())


def _sanitize(text):
    """Strip binary/control noise (some nodes have a non-text prefix); return clean text."""
    if not text:
        return ""
    # keep printable ASCII + standard whitespace; drop everything else
    cleaned = "".join(ch if (32 <= ord(ch) < 127 or ch in "\n\t") else " " for ch in text)
    cleaned = re.sub(r"\s+", " ", cleaned).strip()
    return cleaned


def _usable(original, cleaned):
    """Keep a node only if it's mostly real text after sanitizing."""
    if len(cleaned) < 40:
        return False
    return len(cleaned) / max(len(original), 1) > 0.6


class Memory:
    def __init__(self, path):
        self.ok = False
        self.docs = []      # (id, content)
        self.tokd = []
        self.idf = {}
        self.avgdl = 1.0
        try:
            raw = open(path, encoding="utf-8", errors="replace").read()
            nodes = json.loads(raw).get("nodes", [])
        except Exception:
            return
        df = collections.Counter()
        for n in nodes:
            original = n.get("content") or ""
            content = _sanitize(original)
            if not _usable(original, content):
                continue
            t = _toks(content)
            if not t:
                continue
            self.docs.append((n.get("id", ""), content))
            self.tokd.append(t)
            for w in set(t):
                df[w] += 1
        N = len(self.docs)
        if N == 0:
            return
        self.avgdl = sum(len(t) for t in self.tokd) / N
        self.idf = {w: math.log(1 + (N - f + 0.5) / (f + 0.5)) for w, f in df.items()}
        self.ok = True

    def search(self, query, k=TOPK):
        if not self.ok:
            return []
        qt = _toks(query)
        if not qt:
            return []
        scored = []
        for i, t in enumerate(self.tokd):
            tf = collections.Counter(t)
            dl = len(t)
            s = 0.0
            for w in qt:
                f = tf.get(w, 0)
                if f:
                    s += self.idf.get(w, 0) * (f * 2.5) / (f + 1.5 * (1 - 0.75 + 0.75 * dl / self.avgdl))
            if s > 0:
                scored.append((s, i))
        scored.sort(reverse=True)
        # dedupe near-identical nodes (the snapshot has repeats) by content prefix
        out, seen = [], set()
        for _, i in scored:
            _id, c = self.docs[i]
            sig = c[:120]
            if sig in seen:
                continue
            seen.add(sig)
            out.append((_id, c))
            if len(out) >= k:
                break
        return out


# ── soul HTTP ──────────────────────────────────────────────────────────────
def soul_alive():
    try:
        with urllib.request.urlopen(SOUL + "/health", timeout=5) as r:
            return json.loads(r.read()).get("status") == "alive"
    except Exception:
        return False


def ask(message, agentic=False):
    payload = json.dumps({
        "session_id": SESSION, "message": message, "agentic": agentic,
    }).encode()
    req = urllib.request.Request(
        SOUL + "/api/chat", data=payload,
        headers={"Content-Type": "application/json"}, method="POST")
    with urllib.request.urlopen(req, timeout=300) as r:
        data = json.loads(r.read().decode("utf-8", "replace"))
    return data.get("response") or data.get("reply") or json.dumps(data)[:2000]


def with_memory(message, hits):
    if not hits:
        return message
    block = "\n".join(f"- {c[:MAX_NODE_CHARS].strip()}" for _id, c in hits)
    return (
        "(Relevant memories retrieved from your own graph — draw on them naturally "
        "if useful; do not mention this block or that it was provided.)\n"
        f"{block}\n\n"
        f"(Message:) {message}"
    )


def main():
    print(f"\n{BOLD}{CYAN}Neuron{RESET} — direct chat. "
          f"{DIM}type a message, or 'exit' to leave.{RESET}")

    if not soul_alive():
        print(f"\n{DIM}Neuron isn't responding on :7770. In a separate Terminal run:{RESET}")
        print("  launchctl kickstart -k gui/$(id -u)/ai.neuron.daemons")
        print(f"{DIM}wait a few seconds, then start this again.{RESET}\n")
        return

    print(f"{DIM}loading your memory graph…{RESET}", end="\r", flush=True)
    mem = Memory(SNAP)
    print(" " * 40, end="\r")
    if mem.ok:
        print(f"{DIM}memory on — {len(mem.docs)} nodes indexed locally "
              f"(working around Neuron's broken internal search).{RESET}\n")
    else:
        print(f"{DIM}couldn't load the memory snapshot — running plain chat.{RESET}\n")

    use_mem = mem.ok
    last_hits = []
    agentic = False
    while True:
        try:
            msg = input(f"{GREEN}you ›{RESET} ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nbye.")
            return
        if not msg:
            continue
        low = msg.lower()
        if low in ("exit", "quit", ":q"):
            print("bye.")
            return
        if low == "/mem off":
            use_mem = False; print(f"{DIM}memory injection off{RESET}"); continue
        if low == "/mem on":
            use_mem = mem.ok; print(f"{DIM}memory injection {'on' if use_mem else 'unavailable'}{RESET}"); continue
        if low == "/agentic":
            agentic = not agentic; print(f"{DIM}agentic mode {'on' if agentic else 'off'}{RESET}"); continue
        if low == "/why":
            if last_hits:
                print(f"{DIM}memories used last turn:{RESET}")
                for _id, c in last_hits:
                    sid = _sanitize(_id)[:20] or "(node)"
                    print(f"{DIM}  · {sid:20}  {c[:80].strip()}{RESET}")
            else:
                print(f"{DIM}(none){RESET}")
            continue

        hits = mem.search(msg) if use_mem else []
        last_hits = hits
        outbound = with_memory(msg, hits) if hits else msg

        try:
            tag = f" {DIM}[+{len(hits)} memories]{RESET}" if hits else ""
            print(f"{DIM}…thinking…{RESET}{tag}", end="\r", flush=True)
            reply = ask(outbound, agentic=agentic)
            print(" " * 40, end="\r")
        except KeyboardInterrupt:
            print("\n(cancelled)"); continue
        except Exception as e:
            print(f"{DIM}couldn't reach Neuron: {e}{RESET}")
            if not soul_alive():
                print(f"{DIM}the soul looks down — restart with:{RESET}\n"
                      "  launchctl kickstart -k gui/$(id -u)/ai.neuron.daemons")
            continue

        print(f"{CYAN}{BOLD}neuron ›{RESET} {reply}\n")


if __name__ == "__main__":
    try:
        main()
    except (BrokenPipeError, KeyboardInterrupt):
        pass