#!/bin/bash
# o — unified CLI for the ollie 9P namespace.
#
# Hierarchical access:
#   o <op> [args]                         Root level operations
#   o session <sname> <op> [args]         Session-level operations
#   o agent <sname> <aname> <op> [args]   Agent-level operations
#
# Environment variables ($session, $agent) can substitute for positional args.
#
# Subcommands:
#   read   <path>         Read a file (-l to loop)
#   write  <path> [data]  Write to a file (args or stdin)
#   ls     [path]         List namespace entries
#   new    [name] [cwd]   Create session or agent (context-dependent)
#   ctl    <command>      Send a control command
#   prompt                Interactive REPL for submitting prompts (agent-level)
#   tui                   Launch tmux TUI (agent-level)
#   generate <prompt>     One-shot generation (root-level only)
#   env                   Print resolved session/agent environment
#
# Examples:
#   o ls                              List root (sessions, models, etc.)
#   o session myproj ls               List session contents
#   o session myproj new coding       Create agent 'coding' in 'myproj'
#   o agent myproj coding prompt      Interactive prompt
#   o agent myproj coding read chat   Read chat output

set -euo pipefail

# ── Globals set by dispatch ─────────────────────────────────────
PREFIX=""        # Path prefix based on context
CTX_SESSION=""   # Session name if in session/agent context
CTX_AGENT=""     # Agent name if in agent context

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

cmd_read() {
    local loop=false
    while [[ "${1:-}" == -* ]]; do
        case "$1" in
            -l|--loop) loop=true; shift ;;
            *) break ;;
        esac
    done
    if [[ $# -lt 1 ]]; then
        echo "usage: o read [-l] <path>" >&2
        exit 1
    fi
    local path="${PREFIX}${PREFIX:+/}$1"
    if $loop; then
        while true; do
            ollie-9p read "$path" || break
        done
    else
        ollie-9p read "$path"
    fi
}

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

cmd_ls() {
    local path="${PREFIX}"
    if [[ $# -ge 1 ]]; then
        path="${PREFIX}${PREFIX:+/}$1"
    fi
    ollie-9p ls "$path"
}

cmd_generate() {
    # One-shot generation via /generate (root-level only)
    if [[ -n "$PREFIX" ]]; then
        echo "error: generate is a root-level command" >&2
        echo "  usage: o generate <prompt>" >&2
        exit 1
    fi
    local prompt=""
    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_ctl() {
    if [[ $# -lt 1 ]]; then
        echo "usage: o ctl <command> [args...]" >&2
        if [[ -n "$CTX_AGENT" ]]; then
            echo "  Agent:   stop, compact, clear, inject, agent, model, cwd, name" >&2
        elif [[ -n "$CTX_SESSION" ]]; then
            echo "  Session: kill, pause, resume" >&2
        else
            echo "  (requires session or agent context)" >&2
        fi
        exit 1
    fi

    if [[ -z "$PREFIX" ]]; then
        echo "error: ctl requires session or agent context" >&2
        echo "  usage: o session <name> ctl <cmd>" >&2
        echo "     or: o agent <session> <agent> ctl <cmd>" >&2
        exit 1
    fi

    if [[ -n "$CTX_AGENT" ]]; then
        # Agent-level: rdwr to get response
        printf '%s' "$*" | ollie-9p rdwr "${PREFIX}/ctl"
    else
        # Session-level: write only
        echo "$*" | ollie-9p write "${PREFIX}/ctl"
    fi
}

cmd_prompt() {
    if [[ -z "$CTX_AGENT" ]]; then
        echo "error: prompt requires agent context" >&2
        echo "  usage: o agent <session> <agent> prompt" >&2
        exit 1
    fi

    # If stdin is not a terminal, read it all and send as a single prompt
    if [[ ! -t 0 ]]; then
        cat | ollie-9p write "${PREFIX}/prompt"
        return
    fi

    # Interactive REPL mode
    echo "[$CTX_SESSION/$CTX_AGENT]"
    echo "Type . on a blank line to send, Ctrl+D to exit."
    echo "/ prefix runs ctl: /stop, /compact, /model X, /inject X"
    echo "!q exits the TUI."
    echo ""

    local buf="" line
    
    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 "${PREFIX}/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"
            cmd_ctl "${line#/}" || true
            continue
        fi
        if [[ -z "$buf" && "$line" == !* ]]; then
            history -s "$line"
            local bang="${line#!}"
            case "$bang" in
                q|quit) tmux kill-session -t "o-${CTX_SESSION}" 2>/dev/null; exit 0 ;;
            esac
            continue
        fi
        if [[ "$line" == "." ]]; then
            if [[ -n "$buf" ]]; then
                history -s "$buf"
                printf '%s\n' "$buf" | ollie-9p write "${PREFIX}/prompt"
                buf=""
            fi
            continue
        fi
        if [[ -z "$line" && -n "$buf" ]]; then
            buf="$buf"$'\n'
        elif [[ -z "$line" ]]; then
            continue
        else
            if [[ -n "$buf" ]]; then
                buf="$buf"$'\n'"$line"
            else
                buf="$line"
            fi
        fi
    done
}

cmd_log() {
    local lines=1000
    while [[ $# -gt 0 && "$1" == -* ]]; do
        case "$1" in
            -n) shift; lines="${1:-1000}"; shift ;;
            -n*) lines="${1#-n}"; shift ;;
            *) break ;;
        esac
    done

    if [[ -z "$CTX_AGENT" ]]; then
        echo "error: log requires agent context" >&2
        echo "  usage: o agent <session> <agent> log [-n N]" >&2
        exit 1
    fi
    ollie-9p read "${PREFIX}/log" | tail -n "$lines"
}

cmd_new() {
    if [[ -n "$CTX_AGENT" ]]; then
        echo "error: cannot create nested agents" >&2
        exit 1
    fi

    if [[ -n "$CTX_SESSION" ]]; then
        # Create agent in current session
        if [[ $# -lt 1 ]]; then
            echo "usage: o session <sname> new <agent> [cwd]" >&2
            exit 1
        fi
        local agent_name="$1"
        local agent_cwd="${2:-$PWD}"
        printf 'name=%s\ncwd=%s\n' "$agent_name" "$agent_cwd" | ollie-9p write "${PREFIX}/agent/new"
        echo "created agent: $CTX_SESSION/$agent_name"
    else
        # Create session at root
        if [[ $# -lt 1 ]]; then
            echo "usage: o new <session> [cwd]" >&2
            exit 1
        fi
        local session_name="$1"
        local session_cwd="${2:-$PWD}"
        printf 'name=%s\ncwd=%s\n' "$session_name" "$session_cwd" | ollie-9p write "session/new"
        echo "created session: $session_name"
    fi
}

cmd_tui() {
    if [[ -z "$CTX_AGENT" ]]; then
        echo "error: tui requires agent context" >&2
        echo "  usage: o agent <session> <agent> tui" >&2
        exit 1
    fi

    local s="$CTX_SESSION"
    local a="$CTX_AGENT"

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

    local tmux_session="o-${s}"
    tmux kill-session -t "$tmux_session" 2>/dev/null || true

    # Layout: chat (top 80%), prompt (bottom 20%)
    tmux -CC new-session -d -s "$tmux_session" -c "$HOME" \
        -e "session=$s" -e "agent=$a"
    tmux send-keys -t "$tmux_session" "o agent $s $a log -n 1000 | grep -v '\\[\\[\\[.*\\]\\]\\]' | grep -v '^\`\`\`'; o agent $s $a read chat" Enter

    tmux split-window -t "$tmux_session" -v -l 20% \
        -e "session=$s" -e "agent=$a"
    tmux send-keys -t "$tmux_session" "clear; o agent $s $a prompt" Enter

    tmux set-option -t "$tmux_session" status-right "starting... | %H:%M"
    tmux set-option -t "$tmux_session" status-right-length 64

    while true; do
        ollie-9p read "${PREFIX}/statewait" 2>/dev/null | while read -r state; do
            bg=""
            case "$state" in
                idle)    bg="colour70" ;;
                calling*) bg="colour214" ;;
                thinking) bg="colour33" ;;
                paused)  bg="colour245" ;;
                *)       bg="default" ;;
            esac
            tmux set-option -t "$tmux_session" status-right-style "fg=default,bg=$bg" 2>/dev/null
            tmux set-option -t "$tmux_session" status-right "$state | %H:%M" 2>/dev/null
        done
    done &
    local bg_pid=$!
    trap "kill $bg_pid 2>/dev/null" EXIT

    tmux select-pane -t "$tmux_session:0.1"
    exec tmux attach -t "$tmux_session"
}

cmd_env() {
    # Print export statements for current context
    [[ -n "$CTX_SESSION" ]] && echo "export session=$CTX_SESSION"
    [[ -n "$CTX_AGENT" ]] && echo "export agent=$CTX_AGENT"
}

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

Context selectors:
  o <cmd>                           Root level
  o session <sname> <cmd>           Session level
  o agent <sname> <aname> <cmd>     Agent level

Environment variables $session and $agent can substitute for positional args:
  export session=foo agent=bar
  o agent prompt                    # equivalent to: o agent foo bar prompt

Commands (availability depends on context):

  Root level:
    ls                    List sessions, models, etc.
    new <session> [cwd]   Create a new session
    generate <prompt>     One-shot generation (no session needed)

  Session level (o session <name> ...):
    ls                    List session contents (agents, config)
    new <agent> [cwd]     Create a new agent
    ctl <cmd>             Session control (kill, pause, resume)
    read/write <path>     Access session files

  Agent level (o agent <session> <agent> ...):
    ls                    List agent files
    ctl <cmd>             Agent control (stop, compact, model, etc.)
    prompt                Interactive multi-line prompt REPL
    log [-n N]            Show last N lines of agent log
    tui                   Launch tmux layout
    read/write <path>     Access agent files (chat, cfg, etc.)

  Common:
    read [-l] <path>      Read a file (-l: loop)
    write <path> [data]   Write to a file
    env                   Print export statements for context
    help                  Show this help

Examples:
  o ls                              # list root (sessions, models)
  o new myproj                      # create session 'myproj'
  o session myproj ls               # list session contents
  o session myproj new coding       # create agent 'coding'
  o agent myproj coding prompt      # interactive prompt
  o agent myproj coding read chat   # read chat stream
  o agent myproj coding tui         # launch tmux TUI
  o generate "explain monads"       # one-shot, no session
  
  # With environment:
  export session=myproj agent=coding
  o agent prompt                    # uses env vars
EOF
}

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

dispatch_cmd() {
    local cmd="$1"; shift
    case "$cmd" in
        read)     cmd_read "$@" ;;
        write)    cmd_write "$@" ;;
        ls)       cmd_ls "$@" ;;
        generate) cmd_generate "$@" ;;
        ctl)      cmd_ctl "$@" ;;
        prompt)   cmd_prompt "$@" ;;
        log)      cmd_log "$@" ;;
        new)      cmd_new "$@" ;;
        tui)      cmd_tui "$@" ;;
        env)      cmd_env "$@" ;;
        help|-h|--help) cmd_help ;;
        *)
            echo "unknown command: $cmd" >&2
            echo "run 'o help' for usage" >&2
            exit 1
            ;;
    esac
}

# Main dispatch: check for context selectors
if [[ $# -lt 1 ]]; then
    cmd_help
    exit 0
fi

case "$1" in
    session)
        shift
        # Get session name from arg or env
        CTX_SESSION="${1:-${session:-}}"
        if [[ -n "${1:-}" ]]; then shift; fi
        if [[ -z "$CTX_SESSION" ]]; then
            echo "error: session name required" >&2
            echo "  usage: o session <name> <command>" >&2
            exit 1
        fi
        PREFIX="session/$CTX_SESSION"
        if [[ $# -lt 1 ]]; then
            echo "error: command required after 'session <name>'" >&2
            exit 1
        fi
        dispatch_cmd "$@"
        ;;
    agent)
        shift
        # Get session and agent from args or env
        CTX_SESSION="${1:-${session:-}}"
        if [[ -n "${1:-}" ]]; then shift; fi
        CTX_AGENT="${1:-${agent:-}}"
        if [[ -n "${1:-}" ]]; then shift; fi
        if [[ -z "$CTX_SESSION" || -z "$CTX_AGENT" ]]; then
            echo "error: session and agent names required" >&2
            echo "  usage: o agent <session> <agent> <command>" >&2
            exit 1
        fi
        PREFIX="session/$CTX_SESSION/agent/$CTX_AGENT"
        if [[ $# -lt 1 ]]; then
            echo "error: command required after 'agent <session> <agent>'" >&2
            exit 1
        fi
        dispatch_cmd "$@"
        ;;
    *)
        # Root level - check env vars for backwards compat
        if [[ -n "${agent:-}" && -n "${session:-}" ]]; then
            # Has agent context from env
            CTX_SESSION="$session"
            CTX_AGENT="$agent"
            PREFIX="session/$CTX_SESSION/agent/$CTX_AGENT"
        elif [[ -n "${session:-}" ]]; then
            # Has session context from env
            CTX_SESSION="$session"
            PREFIX="session/$CTX_SESSION"
        fi
        dispatch_cmd "$@"
        ;;
esac