#!/usr/bin/env python3
"""METATRON — lightweight Ollama terminal harness.

Talks to Ollama on nightmare (hardcoded). Launched by typing `metatron`:
  1. model selection screen (fetches live models + sizes from nightmare)
  2. chat REPL with streaming
  3. system access: `!command` + native run_command tool
  4. internet: `?query` (search) / `?fetch <url>` (read a page) + web_search tool

Stdlib-only (urllib + subprocess). Hardcoded: Ollama URL, SearXNG URL, system
prompt, tools. Only the model list is dynamic.
"""
import json
import os
import re
import subprocess
import sys
import time
import urllib.parse
import urllib.request
from html import unescape

OLLAMA = "http://10.30.20.29:11434"   # nightmare
SEARXNG = "http://10.30.20.35:6969"   # self-hosted meta-search (CT 516)
HERE = os.path.dirname(os.path.abspath(__file__))
LAST_MODEL = os.path.join(HERE, ".last_model")

SYSTEM_PROMPT = (
    "You are METATRON, a system-access AI harness running on a Windows Commando "
    "VM. You have full system access through the run_command tool — execute shell "
    "commands when asked and report their output. Use web_search for any up-to-date "
    "or factual internet information. Be concise and direct."
)

TOOLS = [{
    "type": "function",
    "function": {
        "name": "run_command",
        "description": "Execute a system command (cmd/PowerShell) and return its output. "
                       "Use for file ops, process management, recon, anything on the box.",
        "parameters": {
            "type": "object",
            "properties": {
                "command": {"type": "string", "description": "The command to run"},
            },
            "required": ["command"],
        },
    },
}, {
    "type": "function",
    "function": {
        "name": "web_search",
        "description": "Search the internet and return the top results (title, URL, snippet). "
                       "Use for current events, facts, anything you don't already know.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The search query"},
            },
            "required": ["query"],
        },
    },
}]

C = {"r": "\033[0m", "b": "\033[1m", "c": "\033[36m", "g": "\033[32m",
     "y": "\033[33m", "m": "\033[35m"}


def get(path):
    req = urllib.request.Request(OLLAMA + path)
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.load(r)


def post_stream(path, payload):
    data = json.dumps(payload).encode()
    req = urllib.request.Request(OLLAMA + path, data=data,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as r:
        for line in r:
            line = line.strip()
            if line:
                yield json.loads(line)


def list_models():
    d = get("/api/tags")
    out = []
    for m in d.get("models", []):
        ps = (m.get("details") or {}).get("parameter_size", "") or ""
        out.append((m["name"], ps))
    return out


def load_last():
    try:
        return open(LAST_MODEL).read().strip()
    except Exception:
        return None


def save_last(model):
    try:
        open(LAST_MODEL, "w").write(model)
    except Exception:
        pass


def pick_model(models):
    last = load_last()
    names = [n for n, _ in models]
    print(f"\n{C['b']}{C['m']}  METATRON{C['r']} — select a model "
          f"{C['y']}(smaller = faster){C['r']}\n")
    for i, (name, ps) in enumerate(models, 1):
        mark = f" {C['m']}*{C['r']}" if name == last else ""
        print(f"  {C['c']}[{i}]{C['r']} {name}  {C['g']}{ps}{C['r']}{mark}")
    tail = f"  {C['c']}[0]{C['r']} {C['y']}quit{C['r']}"
    if last and last in names:
        tail += f"  {C['y']}(Enter = {last}){C['r']}"
    print(tail + "\n")
    while True:
        try:
            sel = input(f"{C['g']}model>{C['r']} ").strip()
            if sel == "" and last and last in names:
                return last
            if sel in ("0", "q", "quit", "exit"):
                sys.exit(0)
            idx = int(sel) - 1
            if 0 <= idx < len(models):
                return names[idx]
        except ValueError:
            pass
        print(f"{C['y']}  pick a number 1-{len(models)}{C['r']}")


def warm_model(model):
    sys.stdout.write(f"{C['y']}  warming {model} ...{C['r']}")
    sys.stdout.flush()
    payload = {"model": model, "messages": [{"role": "user", "content": "hi"}],
               "stream": False, "think": False, "keep_alive": "30m",
               "options": {"num_predict": 1}}
    try:
        data = json.dumps(payload).encode()
        req = urllib.request.Request(OLLAMA + "/api/chat", data=data,
                                     headers={"Content-Type": "application/json"})
        urllib.request.urlopen(req, timeout=120)
        sys.stdout.write(f"\r{C['g']}  ready.{C['r']}      \n")
    except Exception as e:
        sys.stdout.write(f"\r{C['y']}  warm failed ({e}){C['r']}      \n")


def supports_tools(model):
    try:
        data = json.dumps({"model": model}).encode()
        req = urllib.request.Request(OLLAMA + "/api/show", data=data,
                                     headers={"Content-Type": "application/json"})
        d = json.load(urllib.request.urlopen(req, timeout=15))
        return "tools" in (d.get("capabilities") or [])
    except Exception:
        return False


def run_command(cmd):
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True,
                           timeout=120)
        out = (r.stdout or "") + (r.stderr or "")
        return out.strip() or "(no output)"
    except Exception as e:
        return f"error: {e}"


def web_search(query):
    url = f"{SEARXNG}/search?q={urllib.parse.quote(query)}&format=json"
    try:
        d = json.load(urllib.request.urlopen(url, timeout=25))
    except Exception as e:
        return f"search error: {e}"
    results = d.get("results", [])
    if not results:
        return "no results"
    lines = []
    for r in results[:8]:
        lines.append(f"- {r.get('title', '')}\n  {r.get('url', '')}\n"
                     f"  {(r.get('content') or '')[:220]}")
    return "\n".join(lines)


def fetch_page(url):
    if not re.match(r"^https?://", url):
        url = "https://" + url
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        raw = urllib.request.urlopen(req, timeout=25).read().decode("utf-8", "ignore")
    except Exception as e:
        return f"fetch error: {e}"
    t = re.sub(r"<script.*?</script>", " ", raw, flags=re.S)
    t = re.sub(r"<style.*?</style>", " ", t, flags=re.S)
    t = re.sub(r"<[^>]+>", " ", t)
    t = unescape(t)
    t = re.sub(r"\s+", " ", t)
    return t.strip()[:2000]


def truncate(text, max_lines=50):
    lines = text.split("\n")
    if len(lines) > max_lines:
        return "\n".join(lines[:max_lines]) + f"\n  [... {len(lines) - max_lines} more lines]"
    return text


def print_help():
    print(f"""
  {C['b']}commands{C['r']}
    {C['c']}!<cmd>{C['r']}        run a system command
    {C['c']}?<query>{C['r']}      search the web (SearXNG)
    {C['c']}?fetch <url>{C['r']}  read a page's text
    {C['c']}/save{C['r']}         save this session to markdown
    {C['c']}/help{C['r']}         this list
    {C['c']}exit{C['r']}          quit
""")


def save_transcript(messages, model):
    fn = os.path.join(os.getcwd(),
                      f"metatron-{time.strftime('%Y%m%d-%H%M%S')}.md")
    try:
        with open(fn, "w", encoding="utf-8") as f:
            f.write(f"# METATRON session — {model}\n\n")
            for m in messages:
                role, content = m.get("role"), m.get("content", "")
                if role == "user":
                    f.write(f"**you:** {content}\n\n")
                elif role == "assistant" and content:
                    f.write(f"**metatron:** {content}\n\n")
                elif role == "tool":
                    f.write(f"```\n{content}\n```\n\n")
        return fn
    except Exception as e:
        return f"save failed: {e}"


def chat(model):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    use_tools = supports_tools(model)
    print(f"\n{C['b']}  METATRON {C['c']}:: {model}{C['r']}  "
          f"{C['y']}(!cmd, ?search, ?fetch, /save, /help, exit){C['r']}\n")
    while True:
        try:
            user = input(f"{C['g']}you>{C['r']} ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not user:
            continue
        if user.lower() in ("exit", "quit", "/q", "/exit"):
            break
        if user in ("/help", "help", "?"):
            print_help()
            continue
        if user == "/save":
            fn = save_transcript(messages, model)
            print(f"  {C['g']}saved: {fn}{C['r']}\n")
            continue
        if user.startswith("!"):
            cmd = user[1:].strip()
            print(f"{C['y']}  $ {cmd}{C['r']}")
            print(f"  {truncate(run_command(cmd))}\n")
            continue
        if user.startswith("?"):
            q = user[1:].strip()
            if q.startswith("fetch "):
                url = q[6:].strip()
                print(f"{C['c']}  [fetch] {url}{C['r']}")
                print(f"  {fetch_page(url)}\n")
            elif re.match(r"^https?://|^www\.", q):
                print(f"{C['c']}  [fetch] {q}{C['r']}")
                print(f"  {fetch_page(q)}\n")
            else:
                print(f"{C['c']}  [search] {q}{C['r']}")
                print(f"  {web_search(q)}\n")
            continue

        messages.append({"role": "user", "content": user})
        for _ in range(6):
            payload = {"model": model, "messages": messages, "stream": True,
                       "think": False, "keep_alive": "30m"}
            if use_tools:
                payload["tools"] = TOOLS
            buf = ""
            tool_calls = []
            got_content = False
            sys.stdout.write("  ...")
            sys.stdout.flush()
            try:
                for obj in post_stream("/api/chat", payload):
                    msg = obj.get("message", {})
                    piece = msg.get("content") or ""
                    if piece:
                        if not got_content:
                            sys.stdout.write("\r  ")
                            got_content = True
                        buf += piece
                        sys.stdout.write(piece)
                        sys.stdout.flush()
                    if obj.get("done") and msg.get("tool_calls"):
                        tool_calls = msg["tool_calls"]
            except Exception as e:
                print(f"\n{C['y']}  [ollama error: {e}]{C['r']}")
                break
            if not got_content:
                sys.stdout.write("\r")
                sys.stdout.flush()

            if tool_calls:
                messages.append({"role": "assistant", "content": buf,
                                 "tool_calls": tool_calls})
                for tc in tool_calls:
                    fn = tc.get("function", {})
                    name = fn.get("name")
                    try:
                        args = json.loads(fn.get("arguments", "{}"))
                    except Exception:
                        args = {}
                    if name == "run_command":
                        cmd = args.get("command", "")
                        print(f"\n{C['y']}  $ {cmd}{C['r']}")
                        out = run_command(cmd)
                        print(f"  {truncate(out)}")
                        messages.append({"role": "tool", "content": out})
                    elif name == "web_search":
                        q = args.get("query", "")
                        print(f"\n{C['c']}  [search] {q}{C['r']}")
                        out = web_search(q)
                        print(f"  {out[:600]}")
                        messages.append({"role": "tool", "content": out})
                continue
            else:
                messages.append({"role": "assistant", "content": buf})
                break
        print("\n")


def main():
    try:
        models = list_models()
    except Exception as e:
        print(f"{C['y']}  can't reach Ollama at {OLLAMA}: {e}{C['r']}")
        sys.exit(1)
    if not models:
        print(f"{C['y']}  no models on nightmare{C['r']}")
        sys.exit(1)
    model = pick_model(models)
    save_last(model)
    warm_model(model)
    chat(model)


if __name__ == "__main__":
    main()