#!/usr/bin/env python3
"""
cli.py — Interactive agent: streaming, tool use, personality.

  python3 cli.py              # interactive
  python3 cli.py "question"   # one-shot
"""

import sys, os, re, json, time, signal, subprocess, urllib.request, readline
from pathlib import Path

ROOT = Path(__file__).parent
HISTORY_FILE = ROOT / ".cli_history"
MEMORY_FILE  = ROOT / "memory.jsonl"

# ── ANSI ──────────────────────────────────────────────
B = "\033[1m"; D = "\033[2m"
R = "\033[31m"; G = "\033[32m"; Y = "\033[33m"
C = "\033[36m"; M = "\033[35m"; X = "\033[0m"
def ansi(s, c): return f"{c}{s}{X}"

# ── System prompt ─────────────────────────────────────
SYSTEM = """You are a Linux AI agent. You MUST use tools for any real information — NEVER guess or fabricate.

Tool format (exact):
  ACTION: tool_name | parameters

Tools:
  terminal(command)  — run any shell command
  read_file(path)    — read file contents
  write_file(path | content) — create/overwrite file
  search_files(pattern) — grep/rg search
  web_search(query)  — search the web
  memory(add=text | query=text | list) — persistent memory

RULES:
1. For listing files: ACTION: terminal | ls
2. For reading files: ACTION: read_file | path
3. For searching: ACTION: search_files | pattern
4. After getting tool results, report ONLY what the tool returned.
5. If a tool returns nothing, say "no results" — do NOT make up data.
6. Be concise. One tool at a time."""

# ── Server ────────────────────────────────────────────
PORT = 8081

def health():
    try:
        req = urllib.request.Request(f"http://localhost:{PORT}/health")
        with urllib.request.urlopen(req, timeout=2) as r:
            return json.loads(r.read()).get("status") == "ok"
    except: return False

# ── Streaming chat ────────────────────────────────────
def chat_stream(msgs, max_tok=1024, temp=0.7):
    body = json.dumps({"messages": msgs, "max_tokens": max_tok,
                       "temperature": temp, "stream": True}).encode()
    try:
        req = urllib.request.Request(f"http://localhost:{PORT}/v1/chat/completions",
            body, {"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=180) as r:
            for line in r:
                line = line.decode().strip()
                if line.startswith("data: ") and line != "data: [DONE]":
                    try:
                        d = json.loads(line[6:])
                        c = d["choices"][0].get("delta", {}).get("content")
                        if c: yield c
                    except: pass
    except GeneratorExit: raise
    except Exception as e:
        yield f"\n{ansi('[error]',R)} {e}"

def chat_sync(msgs, max_tok=256, temp=0.3, grammar=None):
    body = {"messages": msgs, "max_tokens": max_tok,
            "temperature": temp, "stream": False}
    if grammar: body["grammar"] = grammar
    try:
        req = urllib.request.Request(f"http://localhost:{PORT}/v1/chat/completions",
            json.dumps(body).encode(), {"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=60) as r:
            return json.loads(r.read())["choices"][0]["message"]["content"].strip()
    except Exception as e: return f"[error: {e}]"

# ── Tool execution ────────────────────────────────────
def execute(name, params):
    try:
        p = params.strip().strip("'\"")

        if name == "terminal":
            r = subprocess.run(p, shell=True, capture_output=True, text=True,
                             timeout=30, cwd=str(ROOT))
            out = r.stdout.strip()
            if r.stderr.strip(): out += f"\n[stderr]: {r.stderr.strip()[:300]}"
            if r.returncode: out += f"\n[exit {r.returncode}]"
            return (out or "(no output)")[:4000]

        elif name == "read_file":
            fp = Path(p)
            if not fp.is_absolute(): fp = ROOT / fp
            if not fp.exists(): return f"Not found: {fp}"
            if fp.stat().st_size > 500_000: return f"Too large ({fp.stat().st_size} bytes)"
            lines = fp.read_text().splitlines()
            out = "\n".join(f"{i+1:4d}|{l}" for i, l in enumerate(lines[:300]))
            if len(lines) > 300: out += f"\n... ({len(lines)-300} more lines)"
            return out

        elif name == "write_file":
            parts = p.split("|", 2)
            if len(parts) < 2:
                m = re.match(r"^['\"]?(.+?)['\"]?\s+(.+)", p, re.DOTALL)
                parts = [m.group(1), m.group(2)] if m else [p, ""]
            fp = Path(parts[0].strip().strip("'\""))
            if not fp.is_absolute(): fp = ROOT / fp
            fp.parent.mkdir(parents=True, exist_ok=True)
            content = parts[1].strip() if len(parts) > 1 else ""
            fp.write_text(content)
            return f"Wrote {len(content)}B → {fp}"

        elif name == "search_files":
            r = subprocess.run(["rg", "--no-heading", "-n", "--max-count=5",
                p, str(ROOT)], capture_output=True, text=True, timeout=10)
            return (r.stdout.strip() or "No matches")[:3000]

        elif name == "web_search":
            try:
                r = subprocess.run(["ddg", p, "-n", "3"],
                    capture_output=True, text=True, timeout=10)
                return r.stdout.strip()[:2000] or "No results"
            except:
                return "web_search unavailable"

        elif name == "memory":
            if p.startswith("add="):
                return mem_add(p[4:].strip().strip("'\""))
            elif p.startswith("query="):
                return mem_search(p[6:].strip().strip("'\""))
            return mem_list()

        return f"Unknown tool: {name}"
    except subprocess.TimeoutExpired: return "Timeout"
    except Exception as e: return f"Error: {e}"

# ── Memory ────────────────────────────────────────────
def mem_load():
    if not MEMORY_FILE.exists(): return []
    es = []
    for l in MEMORY_FILE.read_text().splitlines():
        if l.strip():
            try: es.append(json.loads(l))
            except: pass
    return es

def mem_save(es):
    MEMORY_FILE.write_text("\n".join(json.dumps(e) for e in es[:100]) + "\n")

def mem_add(text):
    es = mem_load(); eid = str(hash(text))[-8:]
    for e in es:
        if e.get("id") == eid: e["hits"] += 1; mem_save(es); return f"Hit: {text[:60]}"
    es.append({"id": eid, "text": text, "ts": time.time(), "hits": 1})
    mem_save(es); return f"Saved: {text[:80]}"

def mem_search(q):
    es = mem_load(); ws = set(q.lower().split())
    sc = [(sum(1 for w in ws if w in e["text"].lower()) + e.get("hits",0)*0.1, e["text"]) for e in es]
    sc.sort(reverse=True)
    return "\n".join(f"  [{s:.1f}] {t[:100]}" for s,t in sc[:5]) or "(none)"

def mem_list():
    es = sorted(mem_load(), key=lambda e: e.get("hits",0), reverse=True)[:10]
    return "\n".join(f"  [{e['hits']}×] {e['text'][:80]}" for e in es) or "(empty)"

# ── Agent loop ─────────────────────────────────────────
MAX_ROUNDS = 3

def run_turn(user_msg, history):
    msgs = [{"role": "system", "content": SYSTEM}]
    mems = mem_search(user_msg)
    if mems and mems != "(none)":
        msgs[0]["content"] += f"\n\nMEMORIES:\n{mems}"
    for h in history[-10:]: msgs.append(h)
    msgs.append({"role": "user", "content": user_msg})

    seen_tools = set()

    for round_n in range(MAX_ROUNDS):
        # Stream response
        sys.stdout.write(ansi("●", C) + " "); sys.stdout.flush()
        text = ""
        for chunk in chat_stream(msgs, max_tok=800, temp=0.7):
            text += chunk; sys.stdout.write(chunk); sys.stdout.flush()
        sys.stdout.write("\n")

        # Detect tool call
        m = re.search(r'ACTION\s*:\s*(\w+)', text, re.IGNORECASE)
        if not m:
            m = re.search(r'(?:TOOL|CALL)\s*:\s*(\w+)', text, re.IGNORECASE)
        if not m:
            m = re.search(r'```(?:tool|action)\s*\n\s*(\w+)', text, re.IGNORECASE)

        if not m:
            # No tool — final response
            history.append({"role": "user", "content": user_msg})
            history.append({"role": "assistant", "content": text})
            return text

        tool = m.group(1).strip().lower()
        # Extract params: text after tool name. Could be same line or next.
        rest = text[m.end():]
        # Try same-line first, then next line
        same_line = rest.split('\n')[0].strip()
        if same_line and not same_line.startswith('*'):
            params = re.sub(r'^[\(\|\=]\s*', '', same_line).strip().rstrip(')')
        else:
            # Params on next line
            lines = rest.split('\n')
            params = lines[1].strip() if len(lines) > 1 else ""

        # Fix common mistakes
        if tool == "terminal" and params:
            if params[:4].isupper() and len(params) < 10:
                params = params.lower()

        sys.stdout.write(f"  {ansi('▸',Y)} {ansi(tool,B)}: {params[:100]}\n")

        # Dedup: skip if we already ran this exact tool+params
        dedup_key = f"{tool}|{params}"
        if dedup_key in seen_tools:
            # Feed a "no new info" result and continue
            msgs.append({"role": "assistant", "content": text})
            msgs.append({"role": "user", "content": f"Already ran: {tool} | {params}. Try a different approach."})
            continue
        seen_tools.add(dedup_key)

        result = execute(tool, params)
        preview = result[:400].replace("\n", "\n  ")
        sys.stdout.write(f"  {ansi('▹',D)} {preview}\n")
        if len(result) > 400:
            sys.stdout.write(f"  {ansi(f'({len(result)} chars)',D)}\n")

        # Feed result
        msgs.append({"role": "assistant", "content": text})
        msgs.append({"role": "user", "content": f"Result of {tool}:\n{result}\n\nRespond directly."})

    # Max rounds — final
    sys.stdout.write(ansi("●", C) + " "); sys.stdout.flush()
    text = ""
    for chunk in chat_stream(msgs, max_tok=800, temp=0.7):
        text += chunk; sys.stdout.write(chunk); sys.stdout.flush()
    sys.stdout.write("\n")
    history.append({"role": "user", "content": user_msg})
    history.append({"role": "assistant", "content": text})
    return text

# ── Context viewer ─────────────────────────────────────
def show_ctx(history):
    print(f"\n{ansi('── CONTEXT ──',B)}")
    print(f"  System: {len(SYSTEM)} chars | History: {len(history)} msgs")
    print(f"  Memories: {len(mem_load())} | Dir: {ROOT}")
    for h in history[-6:]:
        clr = C if h["role"] == "user" else G
        print(f"  {ansi(h['role'][:4].upper(), clr)} {h['content'][:100].replace(chr(10),' ')}")

# ── Main ───────────────────────────────────────────────
def interactive():
    print(f"\n  {ansi('◆',M)} {ansi('Emerging Systems Agent',B)}  {ansi('/help /context /clear /quit',D)}\n")

    if not health():
        print(ansi("  llama-server not running on port 8081. Start with:", R))
        print(ansi("  python3 run.py --server start", D))
        return

    print(f"  {ansi('●',G)} Hermes 3B on port {PORT}\n")

    history = []
    try: readline.read_history_file(str(HISTORY_FILE))
    except: pass

    while True:
        try:
            line = input(ansi("▸ ", C)).strip()
        except (EOFError, KeyboardInterrupt):
            print(f"\n{ansi('Goodbye.', D)}"); break
        if not line: continue

        if line.startswith("/"):
            c = line[1:].strip().lower()
            if c in ("q","quit","exit"): print(f"{ansi('Goodbye.',D)}"); break
            elif c in ("h","help"):
                print(f"  /help /clear /context /mem /mem add <t> /mem search <q> /model /quit")
            elif c == "clear": history = []; print(ansi("  Cleared.", D))
            elif c == "context": show_ctx(history)
            elif c == "model":
                print(f"  {ansi('Hermes 3B Q4_K',B)} ctx=4096 port={PORT}")
            elif c.startswith("mem "):
                sub = c[4:].strip()
                if sub.startswith("add "): print(f"  {mem_add(sub[4:])}")
                elif sub.startswith("search "): print(mem_search(sub[7:]))
                else: print(mem_list())
            else: print(ansi(f"  Unknown: /{c}", R))
            continue

        run_turn(line, history)
        print()
        try: readline.write_history_file(str(HISTORY_FILE))
        except: pass

if __name__ == "__main__":
    if len(sys.argv) > 1 and not sys.argv[1].startswith("/"):
        if not health():
            print("llama-server not running. Use: python3 run.py --server start"); sys.exit(1)
        run_turn(" ".join(sys.argv[1:]), [])
        print()
    else:
        interactive()