#!/usr/bin/env python3
"""
muse-chat-repl.py: Fast interactive terminal chat session for a NetVM muse agent.
Uses direct gateway calls via muse-cli-node in the agent's netns.
"""
import sys
import os
import json
import time
import subprocess
from pathlib import Path

NETVM_BIN = Path(__file__).resolve().parent

# Color formatting helpers
def c_bold(s): return f"\033[1m{s}\033[0m"
def c_cyan(s): return f"\033[36m{s}\033[0m"
def c_green(s): return f"\033[32m{s}\033[0m"
def c_yellow(s): return f"\033[33m{s}\033[0m"
def c_red(s): return f"\033[31m{s}\033[0m"
def c_dim(s): return f"\033[2m{s}\033[0m"

def run_muse(node, *args):
    cmd = [str(NETVM_BIN / "muse-cli-node"), node] + list(args)
    res = subprocess.run(cmd, capture_output=True, text=True)
    return res.returncode, res.stdout, res.stderr

def get_threads(node):
    rc, out, _ = run_muse(node, "threads")
    if rc != 0:
        return []
    try:
        data = json.loads(out)
        if isinstance(data, list):
            return data
    except Exception:
        pass
    return []

def select_thread(node):
    print(c_bold(f"\n=== SELECT THREAD FOR {node.upper()} ===") + c_dim(" (Prefer sidechats per CHAT_POLICY.md)\n"))
    threads = get_threads(node)
    
    # Filter and prioritize sidechats
    top_threads = threads[:15]
    for idx, t in enumerate(top_threads, 1):
        title = t.get("title") or "[Untitled]"
        sid = t.get("session_id")
        updated = t.get("updated", "")
        print(f"  [{idx:2d}] {c_bold(title[:40]):<42} {c_dim(sid[:8])} {c_dim(updated)}")
        
    print(f"\n  [n]  {c_green('+ Create a new sidechat')}")
    print(f"  [c]  {c_cyan('Custom thread UUID')}")
    print(f"  [q]  Quit\n")
    
    while True:
        try:
            choice = input(f"Choose option [1-{len(top_threads)}, n, c, q]: ").strip().lower()
            if not choice:
                continue
            if choice == "q":
                sys.exit(0)
            if choice == "n":
                title = input("Enter title for new sidechat: ").strip()
                if not title:
                    title = "operator-session"
                print(c_dim(f"Creating sidechat '{title}'..."))
                rc, out, err = run_muse(node, "session-start", "--title", title)
                if rc == 0:
                    try:
                        res = json.loads(out)
                        tid = res.get("session_id")
                        if tid:
                            print(c_green(f"✔ Created session {tid}"))
                            return tid, title
                    except Exception:
                        pass
                print(c_red(f"Failed to create session: {err or out}"))
                continue
            if choice == "c":
                tid = input("Enter thread UUID: ").strip()
                if tid:
                    return tid, tid[:8]
                continue
            if choice.isdigit():
                num = int(choice)
                if 1 <= num <= len(top_threads):
                    t = top_threads[num - 1]
                    return t["session_id"], t.get("title") or t["session_id"][:8]
        except (KeyboardInterrupt, EOFError):
            print()
            sys.exit(0)

def render_history(node, thread_id, limit=5):
    rc, out, _ = run_muse(node, "history", "--thread", thread_id, "--limit", str(limit))
    if rc != 0:
        return []
    try:
        msgs = json.loads(out)
        if isinstance(msgs, list):
            print(c_dim(f"\n--- Recent Messages ({len(msgs)}) ---"))
            for m in msgs:
                role = m.get("role", "")
                text = m.get("text", "")
                if role == "assistant":
                    print(f"{c_yellow(node.upper())}: {text}")
                else:
                    print(f"{c_green('YOU')}: {text}")
            print(c_dim("----------------------------\n"))
            return msgs
    except Exception:
        pass
    return []

def main():
    if len(sys.argv) < 2:
        print("Usage: muse-chat-repl.py <account> [--thread <uuid>]", file=sys.stderr)
        sys.exit(1)
        
    node = sys.argv[1]
    thread_id = None
    thread_title = None
    
    if len(sys.argv) >= 4 and sys.argv[2] in ("--thread", "-t"):
        thread_id = sys.argv[3]
        thread_title = thread_id[:8]
        
    if not thread_id:
        thread_id, thread_title = select_thread(node)
        
    print(c_bold(f"\nEntering chat with {c_cyan(node.upper())} in thread '{thread_title}' ({thread_id[:8]}...)"))
    print(c_dim("Commands: /exit (or /quit), /refresh, /history, /switch\n"))
    
    last_msgs = render_history(node, thread_id, limit=5)
    last_seq = last_msgs[-1].get("seq", 0) if last_msgs else 0
    
    while True:
        try:
            prompt_str = f"{c_green('super')} ({c_cyan(thread_title[:16])}) > "
            line = input(prompt_str).strip()
            if not line:
                continue
                
            if line in ("/exit", "/quit", ":q"):
                print(c_dim("\nExited chat session.\n"))
                break
            elif line in ("/refresh", "/history"):
                last_msgs = render_history(node, thread_id, limit=10)
                continue
            elif line == "/switch":
                thread_id, thread_title = select_thread(node)
                print(c_bold(f"\nSwitched to thread '{thread_title}' ({thread_id[:8]}...)\n"))
                render_history(node, thread_id, limit=5)
                continue
                
            # Send message
            print(c_dim("Sending..."), end="\r", flush=True)
            rc, out, err = run_muse(node, "send", "--thread", thread_id, line)
            if rc != 0:
                print(c_red(f"Error sending message: {err or out}"))
                continue
                
            print(c_green("Sent. Waiting for response..."), end="\r", flush=True)
            
            # Poll for assistant reply up to 30 seconds
            start_time = time.time()
            got_reply = False
            while time.time() - start_time < 30:
                time.sleep(2)
                rc, out, _ = run_muse(node, "history", "--thread", thread_id, "--limit", "3")
                if rc == 0:
                    try:
                        msgs = json.loads(out)
                        if msgs and msgs[-1].get("role") == "assistant":
                            # Check if newer than our sent context
                            if msgs[-1].get("seq", 0) > last_seq:
                                print(" " * 50, end="\r")  # clear line
                                print(f"{c_yellow(node.upper())}: {msgs[-1].get('text')}\n")
                                last_seq = msgs[-1].get("seq", 0)
                                got_reply = True
                                break
                    except Exception:
                        pass
                        
            if not got_reply:
                print(" " * 50, end="\r")
                print(c_dim("(Agent is still processing in background. Use /refresh to inspect)\n"))
                
        except (KeyboardInterrupt, EOFError):
            print(c_dim("\nExited chat session.\n"))
            break

if __name__ == "__main__":
    main()