tui/engine.py

"""TinyLiquid Analyst terminal engine.

Purpose-built core for the forensic/dark-web research assistant: chat in the
analyst/skeptic personas with multi-turn memory, optional retrieval-grounded
chat (RAG over the library), SOP-driven case analysis, an /agent task loop,
library search, document reading, and persistent case files. Both the TUI
(tui/analyst.py) and the headless CLI (tui/cli.py) run on this engine.

Guardrails: research/OSINT only. Dark-web actions require an explicit Tor
proxy and the crawler's stop rules; nothing illegal is ever in scope.
"""

import json
import re
import threading
from pathlib import Path

import torch

from model.config import TinyLiquidConfig
from model.tiny_liquid import TinyLiquid
from model.utils import latest_ckpt
from data.tokenizer import load_tokenizer
from research.index import TinyIndex
from research.structured import analyst_report
from research.agent import load_sop, run_case, SOP_ALIASES
from research import orchestrator as orch
from research import websearch as ws
from research import workspace as wkspc
from research.helix import HelixMemory
from research.user_journal import UserJournal
import urllib.parse

ROOT = Path(__file__).resolve().parents[1]
CASE_DIR = ROOT / "cases"

PERSONA_TOK = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "none": ""}
PERSONA_ID = {"analyst": 1, "skeptic": 2, "none": 0}
MEMORY_TURNS = 6          # last N turns injected as conversation history
MAX_HIST_TOK = 768        # cap history tokens inside the context window

BANNER = (
    "TinyLiquid Analyst - on-device forensic research terminal\n"
    "Authorized research/OSINT only. Outputs are decision support, never a verdict.\n"
    "Dark-web actions: only via an explicit Tor proxy, with rate limits and stop rules.\n"
)


class Case:
    def __init__(self, name="default"):
        self.name = name
        self.ledger = []
        self.chat = []

    def add(self, role, text):
        self.chat.append({"role": role, "text": text})

    def note(self, text):
        self.ledger.append(text)

    def save(self, path: Path):
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps({"name": self.name, "ledger": self.ledger,
                                    "chat": self.chat}, indent=2), encoding="utf-8")

    def load(self, path: Path):
        d = json.loads(path.read_text(encoding="utf-8"))
        self.name = d.get("name", self.name)
        self.ledger = d.get("ledger", [])
        self.chat = d.get("chat", [])


class AnalystEngine:
    def __init__(self, ckpt="ckpt/v8_lora/best.pt", tok_path="data/tokenizer.json",
                 library="data/library", threads=8):
        torch.set_num_threads(threads)
        self.tok = load_tokenizer(tok_path)
        self.ckpt = latest_ckpt(ckpt)
        assert self.ckpt, f"no checkpoints in {ckpt}"
        sd = torch.load(self.ckpt, map_location="cpu")
        cfg = TinyLiquidConfig(vocab_size=self.tok.get_vocab_size(),
                               **{k: v for k, v in sd["config"].items() if k != "vocab_size"})
        self.model = TinyLiquid(cfg)
        self.model.load_state_dict(sd["model"])
        self.model.eval()
        self.persona = "analyst"
        self.memory = True
        self.rag = False
        self.case = Case()
        self.library = str(Path(library))
        self.tor = False
        self.helix_on = True
        self.memdb = HelixMemory(str(ROOT / "data" / "helix_memory.jsonl"))
        self.user_journal = UserJournal()
        self.index = self._build_index(library)

    def _build_index(self, library):
        idx = TinyIndex()
        for f in sorted(Path(library).glob("*.txt")):
            idx.add(f.stem, str(f), f.read_text(encoding="utf-8", errors="ignore"))
        for d in (ROOT / "corpus" / "raw",):
            if d.exists():
                for f in sorted(d.glob("*.txt")):
                    idx.add(f"{d.name}/{f.stem}", str(f),
                            f.read_text(encoding="utf-8", errors="ignore"))
        idx.build()
        return idx

    # ---- prompt building ----
    def _rag_context(self, text):
        if not self.rag:
            return ""
        hits = self.index.query(text, k=2)
        if not hits:
            return ""
        parts = ["Library context:"]
        for key, score in hits:
            for k, path, doc in self.index.docs:
                if k == key:
                    parts.append(f"[{key}] " + doc[:400].replace("\n", " "))
                    break
        return "\n".join(parts) + "\n"

    def _history_prompt(self, text):
        """Multi-turn memory: last N case turns -> a chat prefix."""
        if not self.memory:
            return ""
        turns = self.case.chat[-(MEMORY_TURNS * 2):] if self.case.chat else []
        parts = []
        used = 0
        for t in reversed(turns):
            seg = t["text"]
            if used + len(seg) > MAX_HIST_TOK:
                break
            parts.append((t["role"], seg))
            used += len(seg)
        parts.reverse()
        out = []
        for role, seg in parts:
            if role == "user":
                out.append("<|user|>")
            else:
                out.append("<|assistant|>")
            out.append(seg)
        return "".join(out)

    def _build_prompt(self, text):
        text2 = self.user_journal.context() + text
        ctx = self._rag_context(text2)
        hist = self._history_prompt(text2)
        p = "<|user|>" + (ctx + text2 if ctx else text2) + "<|assistant|>"
        return hist + p

    # ---- chat / analysis ----
    def generate(self, text, persona=None, max_new=180, temp=0.6, on_token=None):
        persona = persona or self.persona
        p_token = PERSONA_TOK[persona]
        prompt = p_token + "<|user|>" + text + "<|assistant|>"
        ids = self.tok.encode(prompt).ids
        out = self.model.generate(self.tok, ids, persona_id=PERSONA_ID[persona],
                                  max_new=max_new, temperature=temp, top_k=40,
                                  repetition_penalty=1.4, no_repeat_ngram_size=4,
                                  on_token=on_token)
        return self.tok.decode(out[len(ids):]).strip()

    def chat(self, text, on_token=None):
        self.case.add("user", text)
        self.user_journal.note_thread(text)
        prompt = self._build_prompt(text)
        ids = self.tok.encode(prompt).ids
        out = self.model.generate(self.tok, ids, persona_id=PERSONA_ID[self.persona],
                                  max_new=180, temperature=0.6, top_k=40,
                                  repetition_penalty=1.4, no_repeat_ngram_size=4,
                                  on_token=on_token)
        reply = self.tok.decode(out[len(ids):]).strip()
        self.case.add("assistant", reply)
        return reply

    def analyze(self, text, sop=None):
        sop_text = load_sop(sop, text) if sop else load_sop(None, text)
        prior = ""
        if self.helix_on:
            rec = self.memdb.recall(text, sop_text)
            if rec:
                prior = ("\nPRIOR ANALYSIS OF THIS CASE (recalled):\n" + rec["reasoning"]
                         + "\n-- use as a prior; do not assume it is still correct.\n\n")
        user = self.user_journal.context() + prior + sop_text + "\n\nTASK: " + text
        self.case.add("user", "ANALYSIS: " + text)
        report = analyst_report(self.model, self.tok, user, persona_id=1,
                                max_scratch=120, max_reason=60)
        skeptic = self.generate(
            f"Act as the skeptic. Attack this conclusion:\nClaim: {text}\n"
            f"Conclusion: {report.get('verdict', '')} {report.get('reasoning', '')}",
            persona="skeptic", max_new=100)
        report["skeptic"] = skeptic
        report["sop"] = sop or "auto"
        for q in self._open_questions(report.get("reasoning", "")):
            note = "OPEN: " + q
            if note not in self.case.ledger:
                self.case.ledger.append(note)
        if self.helix_on:
            self.memdb.write(
                claim=text,
                evidence=sop_text if isinstance(sop_text, str) else "",
                verdict=report.get("verdict", ""),
                confidence=report.get("confidence", ""),
                reasoning=report.get("reasoning", ""),
            )
        self.case.add("assistant", json.dumps(report, ensure_ascii=False))
        self.user_journal.note_thread(text)
        return report

    def recall(self, text):
        if not self.helix_on:
            return "long-term memory is off (/mem on)"
        rec = self.memdb.recall(text)
        if not rec:
            return "(no prior record for this case in long-term memory)"
        return ("prior verdict: {0} ({1})\nreasoning: {2}\nsource text: {3}".format(
            rec["verdict"], rec["confidence"], rec["reasoning"],
            (rec.get("claim", "") or "")[:200]))

    def mem_stats(self):
        return {"helix": self.helix_on, **self.memdb.stats()}

    # ---- gap ledger: my special touch ----
    _OPEN_KW = ("missing", "what would settle", "what would change", "what is needed",
                "what is required", "unsupported detail", "not in the record",
                "no record of", "no document", "outstanding")

    def _open_questions(self, text):
        """Pull the open threads out of a verdict's reasoning (pure suit logic)."""
        if not text:
            return []
        out = []
        sents = re.split(r"(?<=[.!?])\s+", text.replace("\n", " ").strip())
        for s in sents:
            low = s.lower()
            if any(k in low for k in self._OPEN_KW):
                q = s.strip()
                if q and q not in out:
                    out.append(q)
        return out

    def gaps(self):
        opens = [l for l in self.case.ledger if l.startswith("OPEN:")]
        return opens or ["(no open questions on this case yet - run /case analyses)"]

    # ---- dual mind: two cognitive minds fused into one opinion ----
    def _run_minds(self, doc):
        """Two minds, one model: analyst + skeptic passes with own memory pools."""
        from research.fusion import run_two_pass
        return run_two_pass(self.model, self.tok, doc,
                            memory=self.memdb if self.helix_on else None)

    def _voice(self, instruction, max_new=150):
        return self.generate(instruction, persona="analyst", max_new=max_new, temp=0.6)

    def opinion(self, text, evidence=None):
        from research.verify import deterministic_verdict
        from research.fusion import fuse, opinion_text
        doc = text if not evidence else "{0}\nEvidence: {1}".format(text, evidence)
        rv = deterministic_verdict(doc)
        rule = rv if rv["verdict"] in ("supports", "refutes", "not enough information") else None
        a, s = self._run_minds(doc)
        op = fuse(a, s, rule=rule, sources=getattr(self, "_last_sources", ()))
        op["spoken"] = opinion_text(op)
        self.case.add("user", "OPINION: " + text)
        for q in op["open_questions"]:
            note = "OPEN: " + q
            if note not in self.case.ledger:
                self.case.ledger.append(note)
        self.case.add("assistant", op["spoken"])
        return op

    def research(self, question, tor=None, pulls=2):
        """Research-partner loop: search+pull docs -> verify -> two minds -> reply."""
        from research.verify import deterministic_verdict
        from research.fusion import fuse, opinion_text
        tor = self.tor if tor is None else bool(tor)

        if self.helix_on:
            rec = self.memdb.recall(question, "")
            if rec:
                reply = self._voice(
                    "A prior record exists for this exact case. Restate your "
                    "assessment of {0} using: verdict {1} ({2}); reasoning {3}".format(
                        question, rec["verdict"], rec.get("confidence", ""),
                        (rec.get("reasoning") or "")[:220]))
                return {"source": "memory", "question": question, "reply": reply,
                        "opinion": (rec.get("reasoning") or "")[:300],
                        "gaps": [], "sources": [], "artifact": None}

        pulled = []
        try:
            pulled = ws.pull(question, n=pulls, tor=tor, library_dir=self.library)
            self._reindex()
        except Exception:
            pulled = []
        self._last_sources = [p.get("file", "") for p in pulled if p.get("file")]

        evidence = ""
        for key, score in self.index.query(question, k=2):
            for k, path, doc in self.index.docs:
                if k == key:
                    evidence += doc[:600].replace("\n", " ") + " "
                    break
        doc = question if not evidence.strip() else question + "\nEvidence: " + evidence[:1200]

        rv = deterministic_verdict(doc)
        rule = rv if rv["verdict"] in ("supports", "refutes", "not enough information") else None
        a, s = self._run_minds(doc)
        op = fuse(a, s, rule=rule, sources=self._last_sources)
        op["spoken"] = opinion_text(op)

        self.case.add("user", "RESEARCH: " + question)
        for q in op["open_questions"]:
            note = "OPEN: " + q
            if note not in self.case.ledger:
                self.case.ledger.append(note)
        art = None
        try:
            res = wkspc.synthesize(self.case.ledger, title="research_" + question[:40],
                                   lines=self.case.ledger)
            art = res[0] if res else None
        except Exception:
            art = None

        srcs = ", ".join(self._last_sources[:4]) or "library index"
        reply = self._voice(
            "Research outcome for: {0}\n{1}\nSources: {2}\n"
            "Reply to the user in your own voice: what the record shows, the "
            "discrepancy you found, your assessment, and what would change it.".format(
                question, op["spoken"], srcs))
        self.case.add("assistant", reply)
        return {"question": question, "reply": reply, "opinion": op["spoken"],
                "verdict": op["verdict"], "confidence": op["confidence"],
                "sources": self._last_sources[:6], "gaps": op["open_questions"],
                "artifact": art}

    def agent(self, task, max_steps=5):
        """Procedural agent loop: model issues SEARCH/READ/NOTE/VERDICT actions."""
        sop_text = load_sop(None, task)
        plan, ledger = run_case(self.model, self.tok, task, self.index,
                                sop_text, max_steps=max_steps)
        report = analyst_report(self.model, self.tok, task, persona_id=1,
                                max_scratch=90, max_reason=50)
        report["steps"] = plan
        self.case.add("user", "AGENT TASK: " + task)
        self.case.add("assistant", json.dumps(report, ensure_ascii=False))
        self.user_journal.note_thread(text)
        return report

    def agents(self, task, n=4, max_steps=5):
        """Parallel research swarm: n agents under distinct angles, merged."""
        sop_text = load_sop(None, task)
        lock = threading.Lock()
        results = orch.run_parallel(self.model, self.tok, task, self.index,
                                    sop_text, n=n, lock=lock,
                                    max_steps=max_steps, library=self.library)
        merged = orch.synthesize(task, results, self._build_index(self.library))
        merged["analyst"] = analyst_report(self.model, self.tok, task, persona_id=1,
                                           max_scratch=90, max_reason=50)
        self.case.add("user", "AGENTS TASK (" + str(n) + " parallel): " + task)
        self.case.add("assistant", json.dumps(merged, ensure_ascii=False))
        return merged

    # ---- research tools ----
    def search(self, query, k=5):
        return self.index.query(query, k)

    def read(self, key):
        for k, path, text in self.index.docs:
            if k == key:
                return f"<{k} ({path})>\n" + text[:2500]
        return "(document not found)"

    def sop_text(self, sop):
        if sop:
            return load_sop(sop, "")
        files = sorted((ROOT / "research" / "sop_library").glob("*.md"))
        return "\n\n".join(f.read_text(encoding="utf-8") for f in files)

    # ---- case files ----
    def save_case(self, name):
        self.case.name = name
        self.case.save(CASE_DIR / f"{name}.json")
        return str(CASE_DIR / f"{name}.json")

    def load_case(self, name):
        path = CASE_DIR / f"{name}.json"
        if not path.exists():
            return f"no case file: {path}"
        self.case = Case(name)
        self.case.load(path)
        return f"loaded {path} ({len(self.case.chat)} msgs, {len(self.case.ledger)} notes)"


    # ---- live web / dark-web retrieval (client hands) ----
    def _reindex(self):
        self.index = self._build_index(self.library)

    def tor_status_msg(self):
        return ws.tor_status()[1]

    def web_search(self, query, n=8):
        return ws.search_web(query, limit=n)

    def web_fetch(self, url):
        doc = ws.fetch(url, tor=self.tor)
        slug = urllib.parse.urlparse(url).path.rstrip("/").rsplit("/", 1)[-1] or "doc"
        path = ws.save_doc(self.library, slug, doc["title"], doc["content"])
        self._reindex()
        doc["file"] = str(path)
        return doc

    def web_pull(self, query, n=3):
        res = ws.pull(query, n=n, tor=self.tor, library_dir=self.library)
        self._reindex()
        return res

    # ---- research workspace (analytic artifacts: the sandbox) ----
    def synthesize(self, title=None, theme=None):
        """Render the case ledger as saved documents (timeline/table/chart/crossref)."""
        title = (title or self.case.name or "case").strip()
        series = getattr(self, "_series", None)
        res = wkspc.synthesize(self.case.ledger, title, series=series,
                               theme=theme, lines=self.case.ledger)
        if not res:
            return ("no dated artifacts yet - run /case analyses, /agent tasks, "
                    "or /chart data first")
        path, doc = res
        return f"saved {path}\n\n" + doc

    def journal(self, arg=None):
        """Run the full journalism suite over the case + library (CaseFile)."""
        from research.journalism import (suite_report, docs_from_library,
                                         claims_from_ledger)
        name = (arg or self.case.name or "case").strip()
        docs = docs_from_library(self.library)
        claims = claims_from_ledger(self.case.ledger)
        events = []
        for line in self.case.ledger:
            m = re.match(r"^(NOTE|EVENT):\s*(\d{4}-\d{1,2}-\d{1,2})\s+(.+)$",
                         line.strip())
            if m:
                events.append({"when": m.group(2), "what": m.group(3),
                               "source_id": "-"})
        md = suite_report(name, docs, claims, events=events)
        return md

    def chart(self, arg):
        """Parse 'title | label:a,b,c | label:x,y,z' and render/save a series chart."""
        title, series = wkspc.parse_series_arg(arg)
        if len(series) < 2:
            return "need at least two series: / chart <title> | <label>:<v1,v2,..> | <label>:<v2,v3,..>"
        self._series = series
        ch = wkspc.render_series(series, title=title)
        path = wkspc.save_md(title, ch)
        return f"saved {path}\n\n" + ch


def parse_cmd(line: str):
    if line.startswith("/"):
        parts = line[1:].split(maxsplit=1)
        return (parts[0].lower(), parts[1].strip() if len(parts) > 1 else "")
    return None