tui/analyst.py

"""TinyLiquid Analyst - neon terminal UI (stdlib curses, zero deps).

An OpenCode-style agent terminal themed like Parrot OS meets Matrix:
near-black background, matrix-green neon, pink/blue/purple accents.
Chat with the analyst/skeptic personas, run SOP-driven cases, search and
read the library, manage case files.

Usage:
  .venv/bin/python tui/analyst.py --ckpt ckpt/dpo
  .venv/bin/python tui/analyst.py --ckpt ckpt/dpo --case mycase.json
"""

import argparse
import curses
import random
import textwrap

from tui.engine import AnalystEngine
from tui.cli import handle_cmd

WRAP = 78

# --- neon palette (256-color) ---
C_BG = 0          # terminal default (black)
C_GREEN = 82      # matrix green
C_PINK = 205
C_BLUE = 39
C_PURPLE = 141
C_CYAN = 51
C_WHITE = 255
C_DIM = 240

PAIR = {
    "bg": 1, "green": 2, "pink": 3, "blue": 4, "purple": 5,
    "cyan": 6, "white": 7, "dim": 8, "user": 9, "model": 10,
    "skeptic": 11, "system": 12, "error": 13, "header": 14,
}


def setup_colors(stdscr):
    curses.start_color()
    curses.use_default_colors()
    try:
        if curses.COLORS >= 256:
            curses.init_pair(PAIR["bg"], C_WHITE, C_BG)
            curses.init_pair(PAIR["green"], C_GREEN, C_BG)
            curses.init_pair(PAIR["pink"], C_PINK, C_BG)
            curses.init_pair(PAIR["blue"], C_BLUE, C_BG)
            curses.init_pair(PAIR["purple"], C_PURPLE, C_BG)
            curses.init_pair(PAIR["cyan"], C_CYAN, C_BG)
            curses.init_pair(PAIR["white"], C_WHITE, C_BG)
            curses.init_pair(PAIR["dim"], C_DIM, C_BG)
            curses.init_pair(PAIR["user"], C_CYAN, C_BG)
            curses.init_pair(PAIR["model"], C_GREEN, C_BG)
            curses.init_pair(PAIR["skeptic"], C_PINK, C_BG)
            curses.init_pair(PAIR["system"], C_PURPLE, C_BG)
            curses.init_pair(PAIR["error"], C_PINK, C_BG)
            curses.init_pair(PAIR["header"], C_GREEN, C_BG)
        else:
            for k in PAIR:
                curses.init_pair(PAIR[k], curses.COLOR_WHITE, curses.COLOR_BLACK)
            curses.init_pair(PAIR["green"], curses.COLOR_GREEN, curses.COLOR_BLACK)
            curses.init_pair(PAIR["pink"], curses.COLOR_RED, curses.COLOR_BLACK)
            curses.init_pair(PAIR["blue"], curses.COLOR_CYAN, curses.COLOR_BLACK)
            curses.init_pair(PAIR["purple"], curses.COLOR_MAGENTA, curses.COLOR_BLACK)
            curses.init_pair(PAIR["cyan"], curses.COLOR_CYAN, curses.COLOR_BLACK)
            curses.init_pair(PAIR["user"], curses.COLOR_CYAN, curses.COLOR_BLACK)
            curses.init_pair(PAIR["model"], curses.COLOR_GREEN, curses.COLOR_BLACK)
            curses.init_pair(PAIR["skeptic"], curses.COLOR_RED, curses.COLOR_BLACK)
            curses.init_pair(PAIR["system"], curses.COLOR_MAGENTA, curses.COLOR_BLACK)
            curses.init_pair(PAIR["error"], curses.COLOR_RED, curses.COLOR_BLACK)
            curses.init_pair(PAIR["header"], curses.COLOR_GREEN, curses.COLOR_BLACK)
    except curses.error:
        pass


def wrap(text, width):
    return "\n".join(textwrap.fill(p, width) for p in text.splitlines()) or " "


class Rain:
    """Minimal matrix-rain strip for the header (lightweight, idle-animated)."""

    def __init__(self, width, rows=1):
        self.width = max(1, width)
        self.rows = rows
        self.cols = [random.randint(0, self.width - 1) for _ in range(self.rows)]

    def frame(self):
        chars = [" "] * self.width
        for i, c in enumerate(self.cols):
            if random.random() < 0.75:
                chars[c] = random.choice("01$#&@!?ABCDEF")
            self.cols[i] = (c + random.choice([-1, 0, 1])) % self.width
        return "".join(chars)


class TerminalUI:
    def __init__(self, stdscr, eng):
        self.stdscr = stdscr
        self.eng = eng
        self.lines = []          # list of (kind, text)
        self.history = []
        self.hist_idx = -1
        self.scroll = 0
        self.input_text = ""
        self.cursor = 0
        self.rain = Rain(10)
        self.status = f"persona: analyst | case: {eng.case.name}"

    def add(self, kind, text):
        for ln in wrap(text, WRAP).splitlines() or [""]:
            self.lines.append((kind, ln))
        self.scroll = 0

    def header(self, wd):
        title = " TINYLIQUID ANALYST "
        left = (wd - len(title)) // 2
        top = "═" * wd
        line = " " * left + title
        return top, line[:wd], self.rain.frame()[:wd]

    def draw(self):
        h, wd = self.stdscr.getmaxyx()
        top, title, rain = self.header(wd)
        body_h = max(1, h - 6)

        self.stdscr.erase()
        self.stdscr.addnstr(0, 0, top[:wd], wd, curses.color_pair(PAIR["purple"]))
        self.stdscr.addnstr(1, 0, title[:wd], wd, curses.color_pair(PAIR["header"]))
        self.stdscr.addnstr(2, 0, rain[:wd], wd, curses.color_pair(PAIR["green"]))

        start = len(self.lines) - body_h - self.scroll
        start = max(0, min(start, max(0, len(self.lines) - body_h)))
        for i in range(body_h):
            idx = start + i
            if idx < len(self.lines):
                kind, ln = self.lines[idx]
                pair = {"user": PAIR["user"], "model": PAIR["model"],
                        "skeptic": PAIR["skeptic"], "system": PAIR["system"],
                        "error": PAIR["error"]}.get(kind, PAIR["white"])
                self.stdscr.addnstr(3 + i, 0, ln[:wd], wd, curses.color_pair(pair))

        self.stdscr.addnstr(h - 3, 0, "─" * wd, wd, curses.color_pair(PAIR["purple"]))
        self.stdscr.addnstr(h - 2, 0, self.status[:wd], wd,
                            curses.color_pair(PAIR["system"]) | curses.A_REVERSE)
        prompt = "❯ "
        self.stdscr.addnstr(h - 1, 0, prompt + self.input_text[:wd - len(prompt) - 1], wd,
                            curses.color_pair(PAIR["green"]))
        try:
            self.stdscr.move(h - 1, min(self.cursor + len(prompt), wd - 1))
        except curses.error:
            pass
        self.stdscr.refresh()

    def run_cmd(self, cmd):
        if cmd in ("/quit", "/exit", "/q"):
            return False
        if cmd == "/clear":
            self.lines = []
            self.add("system", "cleared.")
            return True
        self.add("user", cmd if cmd.startswith("/") else "you> " + cmd)
        if not cmd.startswith("/"):
            # streamed chat: redraw as tokens arrive
            buf = []
            self.lines.append(("model", "…"))
            try:
                def on_token(tok_id):
                    buf.append(tok_id)
                    self.lines[-1] = ("model", self.eng.tok.decode(buf))
                    self.draw()
                self.eng.chat(cmd, on_token=on_token)
            except Exception as e:
                self.lines[-1] = ("error", f"error: {e}")
                self.draw()
                return True
            self.lines[-1] = ("model", self.eng.tok.decode(buf))
            self.draw()
            return True
        try:
            out = handle_cmd(self.eng, cmd)
        except Exception as e:
            out = f"error: {e}"
            self.add("error", out)
            return True
        kind = "skeptic" if cmd.startswith(("/skeptic",)) else "model"
        # render JSON reports nicely
        if out.strip().startswith("{") or out.strip().startswith("["):
            import json as _json
            try:
                data = _json.loads(out)
                out = _json.dumps(data, indent=2, ensure_ascii=False)
            except Exception:
                pass
        for ln in out.splitlines():
            self.add(kind, ln)
        return True

    def loop(self):
        setup_colors(self.stdscr)
        curses.curs_set(1)
        self.stdscr.timeout(80)
        self.add("system", "TinyLiquid Analyst — on-device forensic research terminal")
        self.add("dim", "Authorized research/OSINT only. Outputs are decision support, never a verdict.")
        self.add("dim", "type /help for commands | /quit to exit")
        while True:
            self.draw()
            try:
                ch = self.stdscr.get_wch()
            except curses.error:
                continue
            if ch in ("\n", "\r", "KEY_ENTER"):
                cmd = self.input_text.strip()
                self.input_text, self.cursor = "", 0
                self.history.append(cmd)
                self.hist_idx = len(self.history)
                if cmd and not self.run_cmd(cmd):
                    break
                continue
            if ch in ("\x7f", "KEY_BACKSPACE"):
                if self.cursor > 0:
                    self.input_text = self.input_text[:self.cursor - 1] + self.input_text[self.cursor:]
                    self.cursor -= 1
                continue
            if ch == "KEY_DC":
                self.input_text = self.input_text[:self.cursor] + self.input_text[self.cursor + 1:]
                continue
            if ch == "KEY_LEFT":
                self.cursor = max(0, self.cursor - 1)
                continue
            if ch == "KEY_RIGHT":
                self.cursor = min(len(self.input_text), self.cursor + 1)
                continue
            if ch == "KEY_HOME":
                self.cursor = 0
                continue
            if ch == "KEY_END":
                self.cursor = len(self.input_text)
                continue
            if ch in ("KEY_PPAGE",):
                self.scroll = min(self.scroll + 5, max(0, len(self.lines) - 1))
                continue
            if ch in ("KEY_NPAGE",):
                self.scroll = max(0, self.scroll - 5)
                continue
            if ch == "KEY_UP":
                if self.history and self.hist_idx > 0:
                    self.hist_idx -= 1
                    self.input_text = self.history[self.hist_idx]
                    self.cursor = len(self.input_text)
                continue
            if ch == "KEY_DOWN":
                if self.history and self.hist_idx < len(self.history) - 1:
                    self.hist_idx += 1
                    self.input_text = self.history[self.hist_idx]
                    self.cursor = len(self.input_text)
                else:
                    self.input_text, self.hist_idx = "", len(self.history)
                continue
            if isinstance(ch, str) and ch.isprintable():
                self.input_text = self.input_text[:self.cursor] + ch + self.input_text[self.cursor:]
                self.cursor += 1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="ckpt/v8_lora/best.pt")
    ap.add_argument("--tok", default="data/tokenizer.json")
    ap.add_argument("--case", default=None)
    ap.add_argument("--threads", type=int, default=4)
    args = ap.parse_args()

    eng = AnalystEngine(ckpt=args.ckpt, tok_path=args.tok, threads=args.threads)
    if args.case:
        eng.load_case(args.case)

    def go(stdscr):
        ui = TerminalUI(stdscr, eng)
        ui.status = f"persona: {eng.persona} | case: {eng.case.name} | {eng.ckpt}"
        ui.loop()

    try:
        curses.wrapper(go)
    except KeyboardInterrupt:
        pass
    print("session ended.")


if __name__ == "__main__":
    main()