#!/usr/bin/env python3
"""
MIND CLI (v0) — supports Ops 1–8 via REST calls to the Go backend.

Assumptions:
  - API base: http://localhost:8080 (change with \\set api ...)
  - llama.cpp server: http://localhost:8081 (change with \\set llama ...)
  - Auth: dev token ("Bearer dev"). Adjust as needed.

Core usage model:
  - Free text (no leading '\\') calls /completion (Op 2 helper) on the active conversation/branch.
  - Backslash commands perform explicit ops.

Operations → Commands mapping
  1) Start conversation            → \\new <title> <owner_id>
  2) Commit (prompt/answer or node)→ free text → /completion; or \\node <user|assistant> <content> [parent_id]
  3) Fork branch                   → \\branch <name> [head_node_id]
  4) Delete branch or node         → \\delbranch <name>   |  \\delnode <node_id>
  5) Extract subset as new tree    → \\extract <title> <node_id ...>
  6) Merge two branches            → \\merge <left> <right> <left-first|right-first>
  7) Detach branch to new convo    → \\detach <branch>
  8) Append/copy tree to another   → \\append-tree <src_conv> <src_branch> [dst_conv [dst_branch]]

Helpful extras:
  - \\listconvs <owner_id>
  - \\listbranches <conversation_id>
  - \\linearize (active conv/branch)
  - \\llama <prompt> (debug llama-server directly)
  - \\useconv <id> ; \\usebranch <name> ; \\who ; \\quit
"""

import json
import os
import shlex
import sys
from typing import Dict, Any, List

import requests

DEFAULT_API = os.environ.get("MIND_API", "http://localhost:8080")
DEFAULT_LLAMA = os.environ.get("LLAMA_API", "http://localhost:8081")

auth_headers = lambda: {"Authorization": "Bearer dev", "Content-Type": "application/json"}

class State:
    def __init__(self):
        self.api = DEFAULT_API.rstrip('/')
        self.llama = DEFAULT_LLAMA.rstrip('/')
        self.conversation_id = None  # int
        self.branch = None           # str

    def as_dict(self):
        return {
            "api": self.api,
            "llama": self.llama,
            "conversation_id": self.conversation_id,
            "branch": self.branch,
        }


def pretty(obj: Any) -> str:
    try:
        return json.dumps(obj, indent=2, ensure_ascii=False)
    except Exception:
        return str(obj)


def post_json(url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
    r = requests.post(url, headers=auth_headers(), data=json.dumps(payload), timeout=60)
    r.raise_for_status()
    if r.text.strip() == "":
        return {}
    return r.json()


def delete_json(url: str, payload: Dict[str, Any] | None = None) -> Dict[str, Any]:
    r = requests.delete(url, headers=auth_headers(), data=json.dumps(payload) if payload else None, timeout=60)
    r.raise_for_status()
    return r.json() if r.text.strip() else {}


def get_json(url: str, params: Dict[str, Any]) -> Dict[str, Any] | List[Any]:
    r = requests.get(url, headers=auth_headers(), params=params, timeout=60)
    r.raise_for_status()
    ct = r.headers.get('Content-Type','')
    if 'application/json' in ct:
        return r.json()
    return {"raw": r.text}


# ---------- Commands ----------

def cmd_help():
    print(__doc__)


def cmd_set(st: State, args):
    if len(args) < 2:
        print("usage: \\set <api|llama> <URL>")
        return
    key, val = args[0], args[1]
    if key == 'api':
        st.api = val.rstrip('/')
    elif key == 'llama':
        st.llama = val.rstrip('/')
    else:
        print("unknown key; use 'api' or 'llama'")
        return
    print("ok:", pretty(st.as_dict()))


# (1) Start conversation

def cmd_new(st: State, args):
    if len(args) < 2:
        print("usage: \\new <title> <owner_id>")
        return
    title = args[0]
    try:
        owner_id = int(args[1])
    except ValueError:
        print("owner_id must be an integer")
        return
    out = post_json(f"{st.api}/conversations", {"title": title, "owner_id": owner_id})
    cid = out.get("id")
    st.conversation_id = cid
    print("created conversation:", cid)


def cmd_listconvs(st: State, args):
    if not args:
        print("usage: \\listconvs <owner_id>")
        return
    owner_id = int(args[0])
    out = get_json(f"{st.api}/conversations", {"owner_id": owner_id})
    print(pretty(out))


# (3) Fork branch

def cmd_branch(st: State, args):
    if st.conversation_id is None:
        print("set a conversation first: \\useconv <id> or \\new <title> <owner_id>")
        return
    if len(args) < 1:
        print("usage: \\branch <name> [head_node_id]")
        return
    name = args[0]
    head = int(args[1]) if len(args) > 1 else 0
    payload = {
        "conversation_id": st.conversation_id,
        "name": name,
        "head_node_id": head
    }
    out = post_json(f"{st.api}/branches", payload)
    st.branch = out.get("name", name)
    print("active branch:", st.branch, pretty(out))


def cmd_usebranch(st: State, args):
    if not args:
        print("usage: \\usebranch <name>")
        return
    st.branch = args[0]
    print("active branch:", st.branch)

    
def cmd_listbranches(st: State, args):
    if not args:
        if st.conversation_id is None:
            print("usage: \\listbranches <conversation_id> (or set active with \\useconv)")
            return
        conv = st.conversation_id
    else:
        conv = int(args[0])
    out = get_json(f"{st.api}/branches", {"conversation_id": conv})
    print(pretty(out))


def cmd_useconv(st: State, args):
    if not args:
        print("usage: \\useconv <conversation_id>")
        return
    st.conversation_id = int(args[0])
    print("active conversation:", st.conversation_id)


# (2) Commit helpers

def cmd_node(st: State, args):
    """Create a single node (explicit Op 2). usage: \\node <user|assistant> <content> [parent_id]"""
    if st.conversation_id is None:
        print("set active conversation first")
        return
    if len(args) < 2:
        print("usage: \\node <user|assistant> <content> [parent_id]")
        return
    author = args[0]
    content = args[1]
    parent_id = int(args[2]) if len(args) > 2 else None
    payload = {
        "conversation_id": st.conversation_id,
        "author_kind": author,
        "content": content
    }
    if parent_id is not None:
        payload["parent_id"] = parent_id
    out = post_json(f"{st.api}/nodes", payload)
    print(pretty(out))


def send_prompt(st: State, line: str):
    # Free text → /completion (with auto-branch-create retry)
    if st.conversation_id is None or not st.branch:
        print("need active conversation and branch (\\new/\\useconv and \\branch/\\usebranch)")
        return
    payload = {
        "conversation_id": st.conversation_id,
        "branch": st.branch,
        "prompt": line.strip(),
    }
    try:
        out = post_json(f"{st.api}/completion", payload)
    except requests.HTTPError as e:
        body = e.response.text if e.response is not None else ""
        # If branch doesn't exist yet, create it and retry once
        if e.response is not None and e.response.status_code in (400, 404, 500) and "branch" in body.lower():
            _ = post_json(f"{st.api}/branches", {
                "conversation_id": st.conversation_id,
                "name": st.branch,
                "head_node_id": 0
            })
            out = post_json(f"{st.api}/completion", payload)
        else:
            raise
    print("assistant:\n", out.get("answer"))

    
# (4) Delete branch / node
def cmd_delbranch(st: State, args):
    if st.conversation_id is None:
        print("set active conversation first")
        return
    if not args:
        print("usage: \\delbranch <name>")
        return
    name = args[0]
    out = delete_json(f"{st.api}/branches", {"conversation_id": st.conversation_id, "name": name})
    if st.branch == name:
        st.branch = None
    print("deleted branch:", name, pretty(out))


def cmd_delnode(st: State, args):
    if not args:
        print("usage: \\delnode <node_id>")
        return
    node_id = int(args[0])
    out = delete_json(f"{st.api}/nodes/{node_id}")
    print(pretty(out))


# (5) Extract selected nodes to new conversation

def cmd_extract(st: State, args):
    if len(args) < 2:
        print("usage: \\extract <title> <node_id ...>")
        return
    title = args[0]
    node_ids = [int(x) for x in args[1:]]
    payload = {"conversation_id": st.conversation_id, "title": title, "node_ids": node_ids}
    out = post_json(f"{st.api}/extract", payload)
    new_cid = out.get("new_conversation_id")
    print("extracted to conversation:", new_cid, pretty(out))


# (6) Merge two branches (concat order)

def cmd_merge(st: State, args):
    if st.conversation_id is None:
        print("set active conversation first")
        return
    if len(args) < 3:
        print("usage: \\merge <left_branch> <right_branch> <left-first|right-first>")
        return
    left, right, order = args[0], args[1], args[2]
    payload = {"conversation_id": st.conversation_id, "left_branch": left, "right_branch": right, "order": order}
    out = post_json(f"{st.api}/merge", payload)
    print("merge result:", pretty(out))


# (7) Detach branch → new conversation

def cmd_detach(st: State, args):
    if st.conversation_id is None:
        print("set active conversation first")
        return
    if not args:
        print("usage: \\detach <branch>")
        return
    payload = {"conversation_id": st.conversation_id, "branch": args[0]}
    out = post_json(f"{st.api}/detach", payload)
    print("detached to:", pretty(out))


# (8) Append/copy tree

def cmd_append_tree(st: State, args):
    if len(args) < 2:
        print("usage: \\append-tree <src_conv> <src_branch> [dst_conv [dst_branch]]")
        return
    src_conv = int(args[0])
    src_branch = args[1]
    dst_conv = int(args[2]) if len(args) > 2 else st.conversation_id
    dst_branch = args[3] if len(args) > 3 else st.branch
    payload = {"src_conversation_id": src_conv, "src_branch": src_branch, "dst_conversation_id": dst_conv, "dst_branch": dst_branch}
    out = post_json(f"{st.api}/append-tree", payload)
    print("append-tree result:", pretty(out))


# Linearize & llama & misc

def cmd_linearize(st: State):
    if st.conversation_id is None or not st.branch:
        print("need active conversation and branch (\\useconv / \\usebranch)")
        return
    out = get_json(f"{st.api}/linearize", {"conversation_id": st.conversation_id, "branch": st.branch})
    if isinstance(out, dict):
        print("==== Transcript ====")
        print(out.get("text", ""))
    else:
        print(pretty(out))


def cmd_llama(st: State, args):
    if not args:
        print("usage: \\llama <prompt>")
        return
    prompt = " ".join(args)
    payload = {"prompt": prompt}
    try:
        r = requests.post(f"{st.llama}/completion", headers={"Content-Type": "application/json"}, data=json.dumps(payload), timeout=60)
        r.raise_for_status()
        data = r.json()
        text = None
        if isinstance(data, dict):
            if 'completion' in data:
                text = data['completion']
            elif 'content' in data and isinstance(data['content'], list) and data['content']:
                seg = data['content'][0]
                if isinstance(seg, dict) and 'text' in seg:
                    text = seg['text']
        print("llama says:", text or pretty(data))
    except Exception as e:
        print("llama error:", e)


def cmd_who(st: State):
    print(pretty(st.as_dict()))


# ---------- REPL ----------

def repl():
    st = State()
    print("MIND CLI — type \\help for commands. Free text sends /completion.")
    while True:
        try:
            line = input("> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not line:
            continue
        if line.startswith("\\"):
            parts = shlex.split(line[1:])
            if not parts:
                continue
            cmd, *args = parts
            cmd = cmd.lower()
            try:
                if cmd == 'help':
                    cmd_help()
                elif cmd == 'set':
                    cmd_set(st, args)
                elif cmd == 'new':
                    cmd_new(st, args)
                elif cmd == 'listconvs':
                    cmd_listconvs(st, args)
                elif cmd == 'listbranches':
                    cmd_listbranches(st, args)
                elif cmd == 'useconv':
                    cmd_useconv(st, args)
                elif cmd == 'branch':
                    cmd_branch(st, args)
                elif cmd == 'usebranch':
                    cmd_usebranch(st, args)
                elif cmd == 'node':
                    cmd_node(st, args)
                elif cmd == 'delbranch':
                    cmd_delbranch(st, args)
                elif cmd == 'delnode':
                    cmd_delnode(st, args)
                elif cmd == 'extract':
                    cmd_extract(st, args)
                elif cmd == 'merge':
                    cmd_merge(st, args)
                elif cmd == 'detach':
                    cmd_detach(st, args)
                elif cmd in ('append-tree', 'appendtree', 'append'):
                    cmd_append_tree(st, args)
                elif cmd == 'linearize':
                    cmd_linearize(st)
                elif cmd == 'llama':
                    cmd_llama(st, args)
                elif cmd == 'who':
                    cmd_who(st)
                elif cmd in ('quit', 'exit', 'q'):
                    break
                else:
                    print("unknown command, try \\help")
            except requests.HTTPError as e:
                try:
                    print("HTTP", e.response.status_code, e.response.text)
                except Exception:
                    print("HTTP error:", e)
            except Exception as e:
                print("error:", e)
        else:
            send_prompt(st, line)


if __name__ == '__main__':
    repl()