#!/bin/bash
# o — unified CLI for the ollie 9P namespace.
#
# Path resolution:
#   Paths are resolved based on context:
#   - Absolute paths (starting with session/, agents, models, etc.) → used as-is
#   - Relative paths with $session set but no $agent → session/$session/<path>
#   - Relative paths with both $session and $agent → session/$session/agent/$agent/<path>
#   - No context set → path used as-is (root-relative)
#
# Subcommands:
#   read   <path>         Read a file
#   write  <path> [data]  Write to a file (args or stdin)
#   readloop <path>       Read in a loop (re-reads after each return)
#   ls     [path]         List namespace entries
#   ctl    <command>      Send a control command (stop, compact, kill, ...)
#   prompt                Interactive REPL for submitting prompts (requires agent)
#   tail                  Tail chat output (requires agent)
#   tui                   Launch tmux TUI (requires agent)
#   env                   Print resolved session/agent environment
#   stop                  Interrupt running agent (requires agent)
#   kill                  Kill session (requires session)
#
# Environment:
#   session               Session name (set via 'eval $(o env)')
#   agent                 Agent name

set -euo pipefail

# ── Path classification ─────────────────────────────────────────

# Root-level entries (no context needed)
ROOT_ENTRIES="session elevate agents backends models help ctl eventwait complete generate route aliases tools"

# Session-level files (exist directly under session/{name}/)
SESSION_FILES="env paused connected ctl name id elevate"

# Agent-level files (exist under session/{name}/agent/{name}/)
AGENT_FILES="prompt prompt.prev fifo.in fifo.out chat connection state statewait log offset plan cfg ctl cwd usage cost ctxsz models systemprompt context tail tools name id proc"

is_root_path() {
    local path="$1"
    local first="${path%%/*}"
    for entry in $ROOT_ENTRIES; do
        [[ "$first" == "$entry" ]] && return 0
    done
    return 1
}

is_session_file() {
    local name="$1"
    for f in $SESSION_FILES; do
        [[ "$name" == "$f" ]] && return 0
    done
    [[ "$name" == "agent" || "$name" == agent/* ]] && return 0
    return 1
}

is_agent_file() {
    local name="$1"
    local base="${name%%/*}"
    for f in $AGENT_FILES; do
        [[ "$base" == "$f" ]] && return 0
    done
    return 1
}

# ── Path resolution with validation ─────────────────────────────

# Resolve a path based on current context, with meaningful errors.
# Returns 0 on success (path on stdout), 1 on error (message on stderr).
resolve_path() {
    local path="$1"

    # Root-level paths used as-is
    if is_root_path "$path"; then
        echo "$path"
        return 0
    fi

    # Get context
    local s="${session:-}"
    local a="${agent:-}"

    # Determine what level this path requires
    local base="${path%%/*}"

    # Check if it's an agent-level file
    if is_agent_file "$base"; then
        if [[ -z "$s" ]]; then
            echo "error: '$path' requires session context" >&2
            echo "  export session=<name>  or  use full path: session/<name>/agent/<agent>/$path" >&2
            return 1
        fi
        if [[ -z "$a" ]]; then
            echo "error: '$path' requires agent context" >&2
            echo "  export agent=<name>  or  use full path: session/$s/agent/<name>/$path" >&2
            return 1
        fi
        echo "session/$s/agent/$a/$path"
        return 0
    fi

    # Check if it's a session-level file/dir
    if is_session_file "$base"; then
        if [[ -z "$s" ]]; then
            echo "error: '$path' requires session context" >&2
            echo "  export session=<name>  or  use full path: session/<name>/$path" >&2
            return 1
        fi
        echo "session/$s/$path"
        return 0
    fi

    # Unknown relative path — try with whatever context we have
    if [[ -n "$s" && -n "$a" ]]; then
        echo "session/$s/agent/$a/$path"
    elif [[ -n "$s" ]]; then
        echo "session/$s/$path"
    else
        echo "$path"
    fi
    return 0
}

# ── Context requirements for specific commands ──────────────────

# Resolve session and agent from positional args or env vars.
# Usage: resolve_context "$@"  →  returns (s, a) by modifying caller's vars.
# Pass "$@" from the calling command; it shifts positional args.
resolve_context() {
    local _s _a
    _s="${1:-${session:-}}"
    if [[ -n "${1+set}" ]]; then
        shift
        _a="${1:-${agent:-}}"
    else
        _a="${agent:-}"
    fi
    echo "$_s|$_a"
}

require_session() {
    local s="${session:-}"
    if [[ -z "$s" ]]; then
        echo "error: no session set" >&2
        echo "  export session=<name>" >&2
        return 1
    fi
    echo "$s"
}

require_agent() {
    local s="$1"
    local a="${agent:-}"
    if [[ -z "$a" ]]; then
        echo "error: no agent set" >&2
        echo "  export agent=<name>" >&2
        return 1
    fi
    echo "$a"
}

agent_path() {
    local s a
    s=$(require_session) || return 1
    a=$(require_agent "$s") || return 1
    echo "session/$s/agent/$a"
}

# ── Subcommands ──────────────────────────────────────────────────

cmd_read() {
    if [[ $# -lt 1 ]]; then
        echo "usage: o read <path>" >&2
        exit 1
    fi
    local path
    path=$(resolve_path "$1") || exit 1
    ollie-9p read "$path"
}

cmd_write() {
    if [[ $# -lt 1 ]]; then
        echo "usage: o write <path> [data...]" >&2
        exit 1
    fi
    local file="$1"; shift
    local path
    path=$(resolve_path "$file") || exit 1
    if [[ $# -gt 0 ]]; then
        printf '%s\n' "$*" | ollie-9p write "$path"
    else
        ollie-9p write "$path"
    fi
}

cmd_ls() {
    if [[ $# -ge 1 ]]; then
        local path
        path=$(resolve_path "$1") || exit 1
        ollie-9p ls "$path"
    else
        # No arg: list based on context level
        local s="${session:-}"
        local a="${agent:-}"

        if [[ -n "$s" && -n "$a" ]]; then
            ollie-9p ls "session/$s/agent/$a"
        elif [[ -n "$s" ]]; then
            ollie-9p ls "session/$s"
        else
            ollie-9p ls ""
        fi
    fi
}

cmd_generate() {
    # One-shot generation via /generate (no session/agent needed)
    # Usage: o generate "prompt"
    #        echo "prompt" | o generate
    #        o generate -1 "prompt"  (same as above)
    local prompt=""

    # Handle -1 flag (ignored, for compatibility)
    [[ "${1:-}" == "-1" ]] && shift

    if [[ $# -gt 0 ]]; then
        prompt="$*"
    else
        prompt=$(cat)
    fi

    [[ -z "$prompt" ]] && { echo "usage: o generate <prompt>" >&2; exit 1; }

    printf '%s' "$prompt" | ollie-9p rdwr generate
}

cmd_complete() {
    # Code completion via /complete (no session/agent needed)
    # Usage: o complete '{"file":"...","prefix":"...","suffix":"..."}'
    #        echo '{"file":"..."}' | o complete
    local json=""

    if [[ $# -gt 0 ]]; then
        json="$*"
    else
        json=$(cat)
    fi

    [[ -z "$json" ]] && { echo "usage: o complete <json>" >&2; exit 1; }

    printf '%s' "$json" | ollie-9p rdwr complete
}

cmd_route() {
    # Task routing via /route (no session/agent needed)
    # Usage: o route "task description"
    #        echo "task" | o route
    local task=""

    if [[ $# -gt 0 ]]; then
        task="$*"
    else
        task=$(cat)
    fi

    [[ -z "$task" ]] && { echo "usage: o route <task>" >&2; exit 1; }

    printf '%s' "$task" | ollie-9p rdwr route
}

cmd_readloop() {
    if [[ $# -lt 1 ]]; then
        echo "usage: o readloop <path>" >&2
        exit 1
    fi
    local file="$1"; shift
    local path
    path=$(resolve_path "$file") || exit 1
    while true; do
        ollie-9p read "$path" || break
    done
}

cmd_ctl() {
    if [[ $# -lt 1 ]]; then
        echo "usage: o ctl <command> [args...]" >&2
        echo "  commands: stop, compact, kill, clear, model <m>, backend <b>, cwd <d>" >&2
        exit 1
    fi
    local path
    path=$(agent_path) || exit 1
    echo "$*" | ollie-9p write "$path/ctl"
}

cmd_prompt() {
    local ctx s a
    ctx=$(resolve_context "$@") || exit 1
    s="${ctx%%|*}"
    a="${ctx##*|}"
    if [[ -z "$s" || -z "$a" ]]; then
        echo "error: session and agent are required" >&2
        echo "  usage: o prompt <session> <agent>" >&2
        echo "  or:    export session=<name> agent=<name>" >&2
        exit 1
    fi
    local path="session/$s/agent/$a"

    echo "[$s/$a]"
    echo "Type . on a blank line to send, Ctrl+D to exit."
    echo "Slash commands: /stop /compact /kill /clear /model /backend /agent /cwd /i /help"
    echo ""

    local buf="" line
    
    # Enable bracketed paste mode - terminal sends \e[200~ before and \e[201~ after paste
    bind 'set enable-bracketed-paste on' 2>/dev/null
    
    while true; do
        local prompt_str
        if [[ -z "$buf" ]]; then
            prompt_str=$'\033[1m>\033[0m '
        else
            prompt_str="  "
        fi
        printf '%s' "$prompt_str"
        if ! IFS= read -r line; then
            if [[ -n "$buf" ]]; then
                history -s "$buf"
                printf '%s\n' "$buf" | ollie-9p write "$path/prompt"
            fi
            printf "\n"
            break
        fi
        
        # Drain any additional lines that arrived quickly (multi-line paste)
        while IFS= read -r -t 0.1 extra_line; do
            line="$line"$'\n'"$extra_line"
        done
        if [[ -z "$buf" && "$line" == /* ]]; then
            history -s "$line"
            local cmd="${line#/}"
            case "$cmd" in
                stop)    echo "stop" | ollie-9p write "$path/ctl"; echo "[stopped]" ;;
                compact) echo "compact" | ollie-9p write "$path/ctl"; echo "[compacted]" ;;
                kill)    echo "kill" | ollie-9p write "session/$s/ctl"; echo "[session killed]"; break ;;
                clear)   echo "clear" | ollie-9p write "$path/ctl"; echo "[cleared]" ;;
                model*)  echo "$cmd" | ollie-9p write "$path/ctl"; echo "[${cmd}]" ;;
                backend*) echo "$cmd" | ollie-9p write "$path/ctl"; echo "[${cmd}]" ;;
                cwd*)    echo "$cmd" | ollie-9p write "$path/ctl"; echo "[${cmd}]" ;;
                quit|q)  tmux kill-session -t "o-${s}" 2>/dev/null; exit 0 ;;
                agent*)  echo "$cmd" | ollie-9p write "$path/ctl"; exit 0 ;;
                i*)    echo "/$line" | ollie-9p write "$path/prompt"; echo "[injected]" ;;
                help)
                    echo "  /stop      — interrupt running agent"
                    echo "  /compact   — compact context"
                    echo "  /kill      — kill the session"
                    echo "  /clear     — clear history"
                    echo "  /model X   — switch model"
                    echo "  /backend X — switch backend"
                    echo "  /cwd X     — change working directory"
                    echo "  /quit /q   — kill the tmux TUI"
                    echo "  /i         — interrupt running agent"
                    ;;
                *)  echo "unknown command: /$cmd (try /help)" ;;
            esac
            continue
        fi
        if [[ "$line" == "." ]]; then
            # Send prompt on "." alone
            if [[ -n "$buf" ]]; then
                history -s "$buf"
                printf '%s\n' "$buf" | ollie-9p write "$path/prompt"
                buf=""
            fi
            continue
        fi
        if [[ -z "$line" && -n "$buf" ]]; then
            # Empty line during multi-line input: append newline
            buf="$buf"$'\n'
        elif [[ -z "$line" ]]; then
            continue
        else
            if [[ -n "$buf" ]]; then
                buf="$buf"$'\n'"$line"
            else
                buf="$line"
            fi
        fi
    done
}

cmd_chatstream() {
    local ctx s a
    ctx=$(resolve_context "$@") || exit 1
    s="${ctx%%|*}"
    a="${ctx##*|}"
    if [[ -z "$s" || -z "$a" ]]; then
        echo "error: session and agent are required" >&2
        echo "  usage: o chatstream <session> <agent>" >&2
        echo "  or:    export session=<name> agent=<name>" >&2
        exit 1
    fi
    exec ollie-9p read "session/$s/agent/$a/chat"
}

cmd_log() {
    local lines=1000
    local ctx s a

    # Parse -n flag for line count
    while [[ $# -gt 0 && "$1" == -* ]]; do
        case "$1" in
            -n) shift; lines="${1:-1000}"; shift ;;
            -n*) lines="${1#-n}"; shift ;;
            *) break ;;
        esac
    done

    ctx=$(resolve_context "$@") || exit 1
    s="${ctx%%|*}"
    a="${ctx##*|}"
    if [[ -z "$s" || -z "$a" ]]; then
        echo "error: session and agent are required" >&2
        echo "  usage: o log [-n lines] [session] [agent]" >&2
        echo "  or:    export session=<name> agent=<name>" >&2
        exit 1
    fi
    ollie-9p read "session/$s/agent/$a/log" | tail -n "$lines"
}

cmd_env() {
    # Accept optional positional args to set context
    if [[ $# -ge 1 ]]; then
        echo "export session=$1"
        shift
    elif [[ -n "${session:-}" ]]; then
        echo "export session=$session"
    fi
    if [[ $# -ge 1 ]]; then
        echo "export agent=$1"
    elif [[ -n "${agent:-}" ]]; then
        echo "export agent=$agent"
    fi
}

cmd_stop() {
    local ctx s a
    ctx=$(resolve_context "$@") || exit 1
    s="${ctx%%|*}"
    a="${ctx##*|}"
    if [[ -z "$s" || -z "$a" ]]; then
        echo "error: session and agent are required" >&2
        echo "  usage: o stop <session> <agent>" >&2
        echo "  or:    export session=<name> agent=<name>" >&2
        exit 1
    fi
    echo "stop" | ollie-9p write "session/$s/agent/$a/ctl"
}

cmd_kill() {
    local s="${1:-${session:-}}"
    if [[ -z "$s" ]]; then
        echo "error: no session set" >&2
        echo "  usage: o kill <session>" >&2
        echo "  or:    export session=<name>" >&2
        exit 1
    fi
    echo "kill" | ollie-9p write "session/$s/ctl"
}

cmd_tui() {
    local tmux_session="${session:-}"
    local s a

    # Check for subcommands first
    if [[ "${1:-}" == "exit" ]]; then
        if [[ -z "$tmux_session" ]]; then
            echo "error: no session context" >&2
            exit 1
        fi
        tmux kill-session -t "o-${tmux_session}" 2>/dev/null && echo "[killed o-${tmux_session}]" || echo "no tui session: o-${tmux_session}"
        return
    fi

    # Accept optional positional args: o tui [session] [agent]
    local ctx
    ctx=$(resolve_context "$@")
    s="${ctx%%|*}"
    a="${ctx##*|}"

    if [[ -z "$s" ]]; then
        echo "error: no session set" >&2
        echo "  usage: o tui <session> [agent]" >&2
        echo "  or:    export session=<name>" >&2
        exit 1
    fi
    if [[ -z "$a" ]]; then
        echo "error: no agent set" >&2
        echo "  usage: o tui <session> <agent>" >&2
        echo "  or:    export agent=<name>" >&2
        exit 1
    fi

    tmux_session="o-${s}"

    tmux kill-session -t "$tmux_session" 2>/dev/null || true

    # Layout:
    # ┌──────────────────────────────────┐
    # │  chat                            │
    # │                                  │
    # ├──────────────────────────────────┤
    # │  prompt (20%)                    │
    # └──────────────────────────────────┘
    # State shown in tmux status-right.

    # Pane 0: chat (top, 80%) - show last 1000 lines of log, then tail chat
    tmux -CC new-session -d -s "$tmux_session" -c "$HOME" \
        -e "session=$s" -e "agent=$a"
    tmux send-keys -t "$tmux_session" "o log -n 1000 | grep -v '\\[\\[\\[.*\\]\\]\\]'; o chatstream | grep -v '\\[\\[\\[.*\\]\\]\\]'" Enter

    # Pane 1: prompt REPL (bottom, 20% height)
    tmux split-window -t "$tmux_session" -v -l 20% \
        -e "session=$s" -e "agent=$a"
    tmux send-keys -t "$tmux_session" "clear; o prompt" Enter

    # Show agent state in tmux status bar (live updates via statewait)
    tmux set-option -t "$tmux_session" status-right "starting... | %H:%M"
    tmux set-option -t "$tmux_session" status-right-length 64
    # Background: loops on statewait, updates status bar instantly on each change
    while true; do
        ollie-9p read "session/$s/agent/$a/statewait" 2>/dev/null | while read -r state; do
            bg=""
            case "$state" in
                idle)    bg="colour70" ;;    # green
                calling*) bg="colour214" ;;  # orange
                thinking) bg="colour33" ;;   # blue
                paused)  bg="colour245" ;;   # gray
                *)       bg="default" ;;
            esac
            tmux set-option -t "$tmux_session" status-right-style "fg=default,bg=$bg"
            tmux set-option -t "$tmux_session" status-right "$state | %H:%M"
        done
    done &

    # Focus the prompt pane
    tmux select-pane -t "$tmux_session:0.1"

    exec tmux attach -t "$tmux_session"
}

cmd_help() {
    cat <<'EOF'
usage: o <command> [args...]

Path resolution:
  - Root paths (session/*, agents, models, etc.) → used as-is
  - Agent files (chat, state, prompt, etc.)      → require $session + $agent
  - Session files (env, ctl, name, etc.)         → require $session

Commands:
  read <path>          Read a file
  write <path> [data]  Write to a file (args or stdin)
  readloop <path>      Read in a loop (re-reads after return)
  ls [path]            List namespace entries

  generate <prompt>    One-shot generation (no session needed)
  complete <json>      Code completion (no session needed)
  route <task>         Route task to best model (no session needed)

  ctl <cmd>            Send control command to agent
  stop                 Interrupt running agent
  kill                 Kill active session

  prompt               Interactive multi-line prompt REPL
  chatstream           Stream chat output
  log [-n N]           Show last N lines of agent log (default 1000)
  tui                  Launch tmux layout
  tui exit             Kill tmux TUI session

  env [session] [agent]  Print export statements for context
  help                   Show this help

Environment:
  session              Session name
  agent                Agent name

Examples:
  o ls                           # list at current context level
  o ls session                   # list all sessions
  o read models                  # read root-level file
  o generate "explain monads"    # one-shot, no session
  o complete '{"file":"main.go","prefix":"func "}'  # code completion
  o route "implement login"      # route to best model
  export session=foo agent=bar   # set context
  o read chat                    # read agent's chat
  o readloop statewait           # blocks until agent state changes
EOF
}

# ── Dispatch ─────────────────────────────────────────────────────

if [[ $# -lt 1 ]]; then
    cmd_help
    exit 0
fi

subcmd="$1"; shift
case "$subcmd" in
    read)     cmd_read "$@" ;;
    write)    cmd_write "$@" ;;
    readloop) cmd_readloop "$@" ;;
    ls)       cmd_ls "$@" ;;
    generate) cmd_generate "$@" ;;
    complete) cmd_complete "$@" ;;
    route)    cmd_route "$@" ;;
    ctl)      cmd_ctl "$@" ;;
    stop)     cmd_stop "$@" ;;
    kill)     cmd_kill "$@" ;;
    prompt)   cmd_prompt "$@" ;;
    chatstream) cmd_chatstream "$@" ;;
    log)      cmd_log "$@" ;;
    tui)      cmd_tui "$@" ;;
    env)      cmd_env "$@" ;;
    help|-h|--help) cmd_help ;;
    *)
        echo "unknown command: $subcmd" >&2
        echo "run 'o help' for usage" >&2
        exit 1
        ;;
esac