#!/usr/bin/env bash
# codex-yolo — Launch parallel OpenAI Codex CLI agents with auto-approval via tmux
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
source "$SCRIPT_DIR/lib/common.sh"

SESSION_NAME=""
AUDIT_LOG=""

usage() {
    cat <<'EOF'
Usage: codex-yolo [options] ["task one" "task two" ...]

Launch parallel OpenAI Codex CLI agents in tmux with auto-approval.

Options:
  -s, --session NAME    Custom tmux session name (default: codex-yolo-<timestamp>)
  -d, --dir PATH        Working directory for agents (default: current directory)
  -m, --model MODEL     Model to use (e.g., o4-mini, o3, gpt-4.1)
  -p, --poll SECONDS    Approver poll interval (default: 0.3)
  -f, --file FILE       Read a multiline prompt from a text file
  -c, --command STRING  Slash command to run in the control pane after launch
                        (e.g. "/loop 10m /plan ...", "/queue [...]", "/help").
                        Must start with '/'. Multi-line strings are sent as a paste.
                        Compose multiple commands via /queue rather than repeating -c.
                        Works with --resume to inject into an existing session.
  -r, --resume          Re-attach to an existing yolo session
  --permissions PROFILE Set Codex /permissions profile (default: full-access when allowed, else auto-review)
  --no-codex-sandbox    Disable Codex sandboxing (for externally sandboxed containers)
  --force-codex-sandbox Require Codex sandboxing; do not auto-fallback when unsupported
  -h, --help            Show this help

Worktree options:
  -w, --worktree          Run each agent in its own git worktree
  --base-branch BRANCH    Base branch for worktrees (default: current branch)
  --no-merge              Skip auto-merge after agents complete
  --no-cleanup            Keep worktrees after merge
  --conflict-poll SECS    Conflict detection interval (default: 5)

Examples:
  codex-yolo                     # interactive codex with auto-approval
  codex-yolo "fix the login bug" "add unit tests for auth"
  codex-yolo --model o4-mini --dir /my/project "refactor the API layer"
  codex-yolo --file prompt.txt
  codex-yolo --resume
  codex-yolo -c "/loop 10m /plan Draft the next implementation plan"
  codex-yolo -c '/queue ["/plan first", "/loop 1h /plan continue"]'
  codex-yolo --resume -c "/loops"
  codex-yolo -w -s feat -d /repo "add auth" "add tests" "add docs"
EOF
}

# Attach or switch to a tmux session, handling nested tmux correctly.
tmux_attach() {
    local target="$1"
    if [[ -n "${TMUX:-}" ]]; then
        # Already inside tmux — switch the current client to the new session
        tmux switch-client -t "$target"
    else
        tmux attach -t "$target"
    fi
}

codex_yolo_should_reconcile_auto_review() {
    local worktree_mode="$1" task_count="$2" first_task="$3"

    [[ "$worktree_mode" == "0" ]] || return 1
    [[ "$task_count" == "1" ]] || return 1
    [[ -z "$first_task" ]] || return 1
    [[ "${CODEX_YOLO_PERMISSION_PROFILE:-}" == "codex-auto-review" ]]
}

codex_yolo_start_auto_review_reconciliation() {
    local script_dir="$1" session="$2" audit_log="$3"
    local startup_delay="${CODEX_YOLO_AUTO_REVIEW_STARTUP_DELAY:-1}"

    log_info "Reconciling Codex Auto-review permissions in TUI"
    (
        sleep "$startup_delay" 2>/dev/null || true
        bash "$script_dir/lib/control-pane.sh" "$session" "$audit_log" auto-review-once
    ) >/dev/null 2>&1 &
    disown
}

codex_yolo_inject_control_command() {
    local script_dir="$1" session="$2" audit_log="$3" command="$4"
    local await_auto_review="${5:-0}"

    log_info "Injecting control command: ${command:0:60}${command:60:+...}"
    (
        export CODEX_YOLO_INJECT_AWAIT_AUTO_REVIEW="$await_auto_review"
        bash "$script_dir/lib/control-pane.sh" "$session" "$audit_log" inject-once "$command"
    ) >/dev/null 2>&1 &
    disown
}

main() {
    local session_name="" work_dir="" model="" poll_interval="0.3" resume=0 prompt_file=""
    local worktree_mode=0 base_branch="" no_merge=0 no_cleanup=0 conflict_poll="5"
    local codex_sandbox_policy="${CODEX_YOLO_CODEX_SANDBOX:-auto}"
    local codex_permissions_policy="${CODEX_YOLO_PERMISSIONS:-auto}"
    local control_command=""
    local tasks=()

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -s|--session)  session_name="$2"; shift 2 ;;
            -d|--dir)      work_dir="$2"; shift 2 ;;
            -m|--model)    model="$2"; shift 2 ;;
            -p|--poll)     poll_interval="$2"; shift 2 ;;
            -f|--file)     prompt_file="$2"; shift 2 ;;
            -c|--command)
                if [[ -n "$control_command" ]]; then
                    log_error "Only one -c/--command may be given; compose multiple commands via /queue [...]"
                    exit 1
                fi
                control_command="$2"
                shift 2
                ;;
            -w|--worktree)     worktree_mode=1; shift ;;
            --base-branch)     base_branch="$2"; shift 2 ;;
            --no-merge)        no_merge=1; shift ;;
            --no-cleanup)      no_cleanup=1; shift ;;
            --conflict-poll)   conflict_poll="$2"; shift 2 ;;
            --permissions)          codex_permissions_policy="$2"; shift 2 ;;
            --no-codex-sandbox)    codex_sandbox_policy="off"; shift ;;
            --force-codex-sandbox) codex_sandbox_policy="force"; shift ;;
            -r|--resume)   resume=1; shift ;;
            -h|--help)     usage; exit 0 ;;
            -*)            log_error "Unknown option: $1"; usage; exit 1 ;;
            *)          tasks+=("$1"); shift ;;
        esac
    done

    # Validate -c/--command
    if [[ -n "$control_command" ]]; then
        if [[ "${control_command:0:1}" != "/" ]]; then
            log_error "-c/--command must start with '/' (got: ${control_command:0:40})"
            exit 1
        fi
    fi

    # Check prerequisites
    check_prereqs || exit 1

    # Handle --resume
    if (( resume )); then
        local existing
        existing="$(tmux list-sessions -F '#{session_name}' 2>/dev/null | grep '^codex-yolo-' | head -1)" || true
        if [[ -z "$existing" ]]; then
            log_error "No existing yolo session found"
            exit 1
        fi
        log_info "Re-attaching to session: $existing"
        if [[ -n "$control_command" ]]; then
            local resumed_audit_log="$(log_dir)/codex-yolo-${existing}.log"
            codex_yolo_inject_control_command "$SCRIPT_DIR" "$existing" "$resumed_audit_log" "$control_command"
        fi
        exec tmux_attach "$existing"
    fi

    # Read prompt from file if --file was given
    if [[ -n "$prompt_file" ]]; then
        if [[ ! -f "$prompt_file" ]]; then
            log_error "File not found: $prompt_file"
            exit 1
        fi
        local file_content
        file_content="$(cat "$prompt_file")"
        if [[ -z "$file_content" ]]; then
            log_error "File is empty: $prompt_file"
            exit 1
        fi
        tasks+=("$file_content")
    fi

    # No tasks = launch a single interactive codex session
    if [[ ${#tasks[@]} -eq 0 ]]; then
        tasks+=("")
    fi

    # Set defaults
    SESSION_NAME="${session_name:-codex-yolo-$(date '+%H%M%S')}"
    AUDIT_LOG="$(log_dir)/codex-yolo-${SESSION_NAME}.log"
    work_dir="${work_dir:-$(pwd)}"

    # Validate work_dir
    if [[ ! -d "$work_dir" ]]; then
        log_error "Working directory does not exist: $work_dir"
        exit 1
    fi

    # Resolve to absolute path
    work_dir="$(cd "$work_dir" && pwd)"

    # Ensure Codex CLI config directory exists
    ensure_codex_config
    configure_codex_sandbox "$codex_sandbox_policy" || exit 1
    configure_codex_permissions "$codex_permissions_policy" || exit 1

    # ── worktree setup ───────────────────────────────────────────────────────
    if (( worktree_mode )); then
        source "$SCRIPT_DIR/lib/worktree-manager.sh"

        wt_validate_repo "$work_dir" || exit 1

        if ! check_git_merge_tree; then
            log_error "git 2.38+ required for merge-tree --write-tree (found: $(git version))"
            exit 1
        fi

        if [[ ${#tasks[@]} -lt 2 ]]; then
            log_error "Worktree mode requires at least 2 tasks"
            exit 1
        fi

        base_branch="${base_branch:-$(wt_current_branch "$work_dir")}"
        log_info "Creating ${#tasks[@]} worktrees from branch: $base_branch"

        wt_create_all "$work_dir" "$SESSION_NAME" "$base_branch" "${#tasks[@]}" || exit 1
    fi

    # Truncate audit log for this run
    : > "$AUDIT_LOG"

    log_info "Creating tmux session: $SESSION_NAME"
    log_info "Working directory: $work_dir"
    log_info "Tasks: ${#tasks[@]}"

    # Create detached tmux session with first agent window
    local first_task="${tasks[0]}"
    local agent_dir="$work_dir"
    local cmd

    if (( worktree_mode )); then
        agent_dir="$(wt_path_for "$SESSION_NAME" 1)"
    fi

    if (( worktree_mode )); then
        cmd="$(build_exec_agent_cmd "$model" "$first_task")"
    else
        cmd="$(build_agent_cmd "$model" "$first_task")"
    fi
    if (( worktree_mode )); then
        cmd="$cmd ; touch '$(wt_done_marker "$SESSION_NAME" 1)'"
    fi

    tmux new-session -d -s "$SESSION_NAME" -n "agent-1" -c "$agent_dir"
    tmux send-keys -t "$SESSION_NAME:agent-1" "$cmd" C-m

    log_info "Started agent-1: ${first_task:+${first_task:0:60}...}${first_task:-interactive}"

    local auto_review_pending=0
    if codex_yolo_should_reconcile_auto_review "$worktree_mode" "${#tasks[@]}" "$first_task"; then
        codex_yolo_start_auto_review_reconciliation "$SCRIPT_DIR" "$SESSION_NAME" "$AUDIT_LOG"
        auto_review_pending=1
    fi

    # Create additional agent windows
    local i
    for (( i=1; i<${#tasks[@]}; i++ )); do
        local win_name="agent-$((i+1))"
        local task="${tasks[$i]}"
        agent_dir="$work_dir"

        if (( worktree_mode )); then
            agent_dir="$(wt_path_for "$SESSION_NAME" $((i+1)))"
        fi

        if (( worktree_mode )); then
            cmd="$(build_exec_agent_cmd "$model" "$task")"
        else
            cmd="$(build_agent_cmd "$model" "$task")"
        fi
        if (( worktree_mode )); then
            cmd="$cmd ; touch '$(wt_done_marker "$SESSION_NAME" $((i+1)))'"
        fi

        tmux new-window -t "$SESSION_NAME" -n "$win_name" -c "$agent_dir"
        tmux send-keys -t "$SESSION_NAME:$win_name" "$cmd" C-m

        log_info "Started $win_name: ${task:+${task:0:60}...}${task:-interactive}"
    done

    # Start approver daemon in background (fully detached from this process).
    # The daemon self-terminates when the tmux session is gone.
    log_info "Starting approver daemon (poll=${poll_interval}s)"
    nohup bash "$SCRIPT_DIR/lib/approver-daemon.sh" "$SESSION_NAME" "$poll_interval" "$AUDIT_LOG" \
        >/dev/null 2>&1 &
    disown

    # ── worktree daemons ─────────────────────────────────────────────────────
    if (( worktree_mode )); then
        log_info "Starting conflict daemon (poll=${conflict_poll}s)"
        nohup bash "$SCRIPT_DIR/lib/conflict-daemon.sh" "$SESSION_NAME" "$conflict_poll" "$AUDIT_LOG" \
            >/dev/null 2>&1 &
        disown

        if (( ! no_merge )); then
            local merge_opts=""
            (( no_cleanup )) && merge_opts="$merge_opts --no-cleanup"
            [[ -n "$model" ]] && merge_opts="$merge_opts --model $model"
            [[ -n "${CODEX_YOLO_PERMISSION_PROFILE:-}" ]] && merge_opts="$merge_opts --permissions '$CODEX_YOLO_PERMISSION_PROFILE'"
            (( CODEX_YOLO_BYPASS_CODEX_SANDBOX )) && merge_opts="$merge_opts --no-codex-sandbox"
            if [[ -n "${CODEX_YOLO_FAKE_BWRAP_DIR:-}" ]]; then
                local fake_bwrap_dir="${CODEX_YOLO_FAKE_BWRAP_DIR//\'/\'\\\'\'}"
                merge_opts="$merge_opts --fake-bwrap-dir '$fake_bwrap_dir'"
            fi

            tmux new-window -t "$SESSION_NAME" -n "merge" -c "$work_dir"
            tmux send-keys -t "$SESSION_NAME:merge" \
                "bash '$SCRIPT_DIR/lib/merge-resolver.sh' '$SESSION_NAME' '$AUDIT_LOG' $merge_opts" C-m
        fi
    fi

    # Create control window with audit log tailing and slash commands
    tmux new-window -t "$SESSION_NAME" -n "control" -c "$work_dir"
    local control_mode="standard"
    (( worktree_mode )) && control_mode="worktree"
    local control_script="${SCRIPT_DIR//\'/\'\\\'\'}"
    local control_session="${SESSION_NAME//\'/\'\\\'\'}"
    local control_log="${AUDIT_LOG//\'/\'\\\'\'}"
    tmux send-keys -t "$SESSION_NAME:control" \
        "bash '$control_script/lib/control-pane.sh' '$control_session' '$control_log' '$control_mode'" C-m

    # Inject one-shot control command if requested. When auto-review reconciliation
    # is also in flight, gate the injection on its completion so the first /loop or
    # /plan keystrokes don't land on the Codex welcome / permissions screen.
    if [[ -n "$control_command" ]]; then
        codex_yolo_inject_control_command \
            "$SCRIPT_DIR" "$SESSION_NAME" "$AUDIT_LOG" "$control_command" "$auto_review_pending"
    fi

    # Select the first agent window
    tmux select-window -t "$SESSION_NAME:agent-1"

    log_info "Attaching to session (Ctrl-b n/p to switch windows, Ctrl-b d to detach)"
    log_info "Audit log: $AUDIT_LOG"

    # Attach (or switch if already inside tmux)
    tmux_attach "$SESSION_NAME"
}

build_agent_cmd() {
    local model="$1" task="$2"
    local cmd
    local prefix
    local permission_arg

    prefix="$(codex_yolo_command_prefix)"
    permission_arg="$(codex_yolo_permission_config_arg)"

    if (( CODEX_YOLO_FORCE_CODEX_SANDBOX )); then
        cmd="${prefix}codex --full-auto $permission_arg"
    elif (( CODEX_YOLO_BYPASS_CODEX_SANDBOX )); then
        cmd="${prefix}codex --dangerously-bypass-approvals-and-sandbox $permission_arg"
    else
        cmd="${prefix}codex --yolo $permission_arg"
    fi
    cmd="${cmd%" "}"

    if [[ -n "$model" ]]; then
        cmd="$cmd --model $model"
    fi

    # No task = interactive mode
    if [[ -z "$task" ]]; then
        echo "$cmd"
        return
    fi

    # For multiline prompts, write to a temp file and pipe into codex.
    # This avoids shell quoting issues with tmux send-keys.
    if [[ "$task" == *$'\n'* ]]; then
        local tmpfile
        tmpfile="/tmp/codex-yolo-prompt-$RANDOM$RANDOM.txt"
        printf '%s' "$task" > "$tmpfile"
        echo "cat '$tmpfile' | $cmd"
        return
    fi

    # Escape single quotes in the task for shell safety
    local escaped_task="${task//\'/\'\\\'\'}"
    cmd="$cmd '$escaped_task'"

    echo "$cmd"
}

build_exec_agent_cmd() {
    local model="$1" task="$2"
    local cmd
    cmd="$(codex_yolo_command_prefix)codex exec $(codex_yolo_permission_config_arg)"
    cmd="${cmd%" "}"

    if (( CODEX_YOLO_BYPASS_CODEX_SANDBOX )); then
        cmd="$cmd --dangerously-bypass-approvals-and-sandbox"
    fi

    if [[ -n "$model" ]]; then
        cmd="$cmd --model $model"
    fi

    if [[ "$task" == *$'\n'* ]]; then
        local tmpfile
        tmpfile="/tmp/codex-yolo-prompt-$RANDOM$RANDOM.txt"
        printf '%s' "$task" > "$tmpfile"
        echo "cat '$tmpfile' | $cmd -"
        return
    fi

    local escaped_task="${task//\'/\'\\\'\'}"
    cmd="$cmd '$escaped_task'"

    echo "$cmd"
}

main "$@"
