#!/usr/bin/env bash
# bg-model-call — Submit a prompt to a third-party model via the agent-runtimes
# control plane, wait for completion, and write the result to a file.
#
# Designed for use with the ScheduleWakeup polling pattern: run this script in
# the background; your wakeup handler checks for the <output>.done sentinel.
#
# See specs/bg-model-call.spec.md for the full spec.

set -uo pipefail

# ---------------------------------------------------------------------------
# Model → template mapping (override with --template)
# ---------------------------------------------------------------------------
declare -A MODEL_TEMPLATES=(
    [minimax]="minimax-planning"
    [sonnet]="sonnet-code-repo"
    [haiku]="haiku-code-repo"
    [opus]="opus-code-repo"
)

# ---------------------------------------------------------------------------
# Colour output helpers
# ---------------------------------------------------------------------------
RED='\033[0;31m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
NC='\033[0m'

info()  { printf "${CYAN}[bg-model-call]${NC} %s\n" "$*" >&2; }
warn()  { printf "${YELLOW}[bg-model-call] WARNING:${NC} %s\n" "$*" >&2; }
error() { printf "${RED}[bg-model-call] ERROR:${NC} %s\n" "$*" >&2; }

# ---------------------------------------------------------------------------
# Usage
# ---------------------------------------------------------------------------
usage() {
    cat <<'EOF'
bg-model-call — Submit a prompt to a model via the agent-runtimes control plane

Usage:
  bg-model-call --model <alias> --prompt <file> --output <file> [OPTIONS]

Required:
  --model <alias>       Model alias: minimax, sonnet, haiku, opus
  --prompt <file>       Path to file containing the prompt text
  --output <file>       Path where the result will be written

Optional:
  --template <name>     Override the task template (without .yaml)
  --project <id>        Project ID for CP grouping (default: bg-model-call)
  --task-name <name>    Task name shown in agent-monitor
                        (default: bg-model-call-<model>)
  --url <cp-url>        Control plane URL (default: $CP_URL or http://localhost:8100)
  --timeout <seconds>   Max seconds to wait for task (default: 1800)
  --poll-interval <s>   Seconds between polls (default: 10)
  --dryrun, -n          Print what would be submitted without calling the CP
  --help, -h            Show this help and exit

Model-to-template mapping:
  minimax  → minimax-planning
  sonnet   → sonnet-code-repo
  haiku    → haiku-code-repo
  opus     → opus-code-repo

The .done sentinel (<output>.done) is written on success — poll for this file
from a ScheduleWakeup handler.
EOF
}

# ---------------------------------------------------------------------------
# Argument defaults
# ---------------------------------------------------------------------------
ARG_MODEL=""
ARG_TEMPLATE=""
ARG_PROMPT=""
ARG_OUTPUT=""
ARG_PROJECT="bg-model-call"
ARG_TASK_NAME=""
ARG_URL=""
ARG_TIMEOUT="1800"
ARG_POLL_INTERVAL="10"
DRYRUN=0

# ---------------------------------------------------------------------------
# Parse arguments
# ---------------------------------------------------------------------------
while [[ $# -gt 0 ]]; do
    case "$1" in
        --model)
            ARG_MODEL="${2:?--model requires a value}"
            shift 2
            ;;
        --template)
            ARG_TEMPLATE="${2:?--template requires a value}"
            shift 2
            ;;
        --prompt)
            ARG_PROMPT="${2:?--prompt requires a value}"
            shift 2
            ;;
        --output)
            ARG_OUTPUT="${2:?--output requires a value}"
            shift 2
            ;;
        --project)
            ARG_PROJECT="${2:?--project requires a value}"
            shift 2
            ;;
        --task-name)
            ARG_TASK_NAME="${2:?--task-name requires a value}"
            shift 2
            ;;
        --url)
            ARG_URL="${2:?--url requires a value}"
            shift 2
            ;;
        --timeout)
            ARG_TIMEOUT="${2:?--timeout requires a value}"
            shift 2
            ;;
        --poll-interval)
            ARG_POLL_INTERVAL="${2:?--poll-interval requires a value}"
            shift 2
            ;;
        --dryrun|-n)
            DRYRUN=1
            shift
            ;;
        --help|-h)
            usage
            exit 0
            ;;
        *)
            error "Unknown flag: $1"
            usage
            exit 1
            ;;
    esac
done

# ---------------------------------------------------------------------------
# Validate required args
# ---------------------------------------------------------------------------
if [[ -z "$ARG_PROMPT" ]]; then
    error "--prompt is required"
    exit 1
fi
if [[ -z "$ARG_OUTPUT" ]]; then
    error "--output is required"
    exit 1
fi
if [[ -z "$ARG_MODEL" && -z "$ARG_TEMPLATE" ]]; then
    error "--model is required (or use --template to specify a template directly)"
    exit 1
fi

# Validate numeric args
if ! [[ "$ARG_TIMEOUT" =~ ^[1-9][0-9]*$ ]]; then
    error "--timeout must be a positive integer, got: $ARG_TIMEOUT"
    exit 1
fi
if ! [[ "$ARG_POLL_INTERVAL" =~ ^[1-9][0-9]*$ ]]; then
    error "--poll-interval must be a positive integer, got: $ARG_POLL_INTERVAL"
    exit 1
fi

# ---------------------------------------------------------------------------
# Resolve values
# ---------------------------------------------------------------------------

# Resolve CP URL: --url flag → $CP_URL → default
if [[ -n "$ARG_URL" ]]; then
    CP_URL_RESOLVED="$ARG_URL"
elif [[ -n "${CP_URL:-}" ]]; then
    CP_URL_RESOLVED="$CP_URL"
else
    CP_URL_RESOLVED="http://localhost:8100"
fi

# Resolve template
if [[ -n "$ARG_TEMPLATE" ]]; then
    RESOLVED_TEMPLATE="$ARG_TEMPLATE"
elif [[ -n "$ARG_MODEL" ]]; then
    if [[ -z "${MODEL_TEMPLATES[$ARG_MODEL]:-}" ]]; then
        error "Unknown model alias '$ARG_MODEL'. Valid values: ${!MODEL_TEMPLATES[*]}"
        exit 1
    fi
    RESOLVED_TEMPLATE="${MODEL_TEMPLATES[$ARG_MODEL]}"
fi

# Resolve task name
if [[ -n "$ARG_TASK_NAME" ]]; then
    RESOLVED_TASK_NAME="$ARG_TASK_NAME"
elif [[ -n "$ARG_MODEL" ]]; then
    RESOLVED_TASK_NAME="bg-model-call-${ARG_MODEL}"
else
    RESOLVED_TASK_NAME="bg-model-call-${RESOLVED_TEMPLATE}"
fi

# Validate prompt file (must be readable even in dryrun)
if [[ ! -f "$ARG_PROMPT" ]]; then
    error "Cannot read prompt file: $ARG_PROMPT"
    exit 1
fi
PROMPT_SIZE=$(wc -c < "$ARG_PROMPT")

# Validate output directory exists
OUTPUT_DIR="$(dirname "$ARG_OUTPUT")"
if [[ ! -d "$OUTPUT_DIR" ]]; then
    error "Output directory does not exist: $OUTPUT_DIR"
    exit 1
fi

# ---------------------------------------------------------------------------
# Locate dispatch-task and wait-for-tasks CLIs
# ---------------------------------------------------------------------------
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
AR_SCRIPTS="${SCRIPT_DIR}/../../../projects/agent-runtimes/scripts"

find_cli() {
    local name="$1"
    if command -v "$name" &>/dev/null; then
        command -v "$name"
    elif [[ -x "${AR_SCRIPTS}/${name}" ]]; then
        echo "${AR_SCRIPTS}/${name}"
    else
        echo ""
    fi
}

DISPATCH_TASK_BIN="$(find_cli dispatch-task)"
WAIT_FOR_TASKS_BIN="$(find_cli wait-for-tasks)"

# ---------------------------------------------------------------------------
# Dryrun: print resolved values and exit
# ---------------------------------------------------------------------------
if [[ "$DRYRUN" -eq 1 ]]; then
    echo "[dryrun] bg-model-call — no changes will be made"
    echo "  CP URL:        ${CP_URL_RESOLVED}"
    if [[ -n "$ARG_MODEL" ]]; then
        echo "  Model:         ${ARG_MODEL}"
    fi
    echo "  Template:      ${RESOLVED_TEMPLATE}"
    echo "  Task name:     ${RESOLVED_TASK_NAME}"
    echo "  Project ID:    ${ARG_PROJECT}"
    echo "  Prompt file:   ${ARG_PROMPT}"
    echo "  Prompt size:   ${PROMPT_SIZE} bytes"
    echo "  Output file:   ${ARG_OUTPUT}"
    echo "  Sentinel:      ${ARG_OUTPUT}.done"
    echo "  Timeout:       ${ARG_TIMEOUT}s"
    echo "  Poll interval: ${ARG_POLL_INTERVAL}s"
    if [[ -n "$DISPATCH_TASK_BIN" ]]; then
        echo "  dispatch-task: ${DISPATCH_TASK_BIN}"
    else
        echo "  dispatch-task: not found (will use curl fallback)"
    fi
    if [[ -n "$WAIT_FOR_TASKS_BIN" ]]; then
        echo "  wait-for-tasks: ${WAIT_FOR_TASKS_BIN}"
    else
        echo "  wait-for-tasks: not found (will use curl polling fallback)"
    fi
    exit 0
fi

# ---------------------------------------------------------------------------
# Helper: write failure sentinel and exit 1
# ---------------------------------------------------------------------------
write_failure() {
    local task_id="${1:-unknown}" state="${2:-failed}" detail="${3:-}"
    {
        echo "task_id=${task_id}"
        echo "state=${state}"
        if [[ -n "$detail" ]]; then
            echo "error=${detail}"
        fi
    } > "${ARG_OUTPUT}.failed"

    # Remove tracking file if it exists
    if [[ -n "${task_id}" && "${task_id}" != "unknown" ]]; then
        rm -f "${OUTPUT_DIR}/.bg-model-call-${task_id}"
    fi

    error "Task ${task_id} ${state}. See ${ARG_OUTPUT}.failed"
    exit 1
}

# ---------------------------------------------------------------------------
# Submit task
# ---------------------------------------------------------------------------
info "Submitting task (model=${ARG_MODEL:-via-template}, template=${RESOLVED_TEMPLATE})"

submit_via_dispatch_task() {
    "$DISPATCH_TASK_BIN" \
        --name "$RESOLVED_TASK_NAME" \
        --project-id "$ARG_PROJECT" \
        --prompt-file "$ARG_PROMPT" \
        --template "$RESOLVED_TEMPLATE" \
        --cp-url "$CP_URL_RESOLVED" \
        --format json \
        --skip-tag-check 2>/dev/null
}

submit_via_curl() {
    local prompt_text
    prompt_text="$(cat "$ARG_PROMPT")"
    local payload
    payload="$(printf '{"name":"%s","project_id":"%s","prompt":%s,"runtime":{"model":"%s","timeout":%s}}' \
        "$RESOLVED_TASK_NAME" \
        "$ARG_PROJECT" \
        "$(printf '%s' "$prompt_text" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))')" \
        "${ARG_MODEL:-${RESOLVED_TEMPLATE}}" \
        "$ARG_TIMEOUT")"
    curl -s -X POST "${CP_URL_RESOLVED}/tasks" \
        -H "Content-Type: application/json" \
        -d "$payload"
}

SUBMIT_RESPONSE=""
if [[ -n "$DISPATCH_TASK_BIN" ]]; then
    SUBMIT_RESPONSE="$(submit_via_dispatch_task 2>&1)" || true
else
    warn "dispatch-task not found; using curl fallback"
    SUBMIT_RESPONSE="$(submit_via_curl 2>&1)" || true
fi

# Extract task_id from JSON response
TASK_ID="$(printf '%s' "$SUBMIT_RESPONSE" | python3 -c 'import json,sys; d=json.load(sys.stdin); print(d.get("task_id",""))' 2>/dev/null || true)"

if [[ -z "$TASK_ID" ]]; then
    error "Failed to submit task. Response: ${SUBMIT_RESPONSE}"
    write_failure "unknown" "failed" "submission failed: ${SUBMIT_RESPONSE}"
fi

info "Task submitted: ${TASK_ID}"

# Write tracking file
{
    echo "task_id=${TASK_ID}"
    echo "model=${ARG_MODEL:-}"
    echo "template=${RESOLVED_TEMPLATE}"
    echo "submitted_at=$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
} > "${OUTPUT_DIR}/.bg-model-call-${TASK_ID}"

# ---------------------------------------------------------------------------
# Wait for task completion
# ---------------------------------------------------------------------------
info "Waiting for task ${TASK_ID} (timeout=${ARG_TIMEOUT}s, poll=${ARG_POLL_INTERVAL}s)"

FINAL_STATE=""

wait_via_wait_for_tasks() {
    "$WAIT_FOR_TASKS_BIN" \
        --url "$CP_URL_RESOLVED" \
        --timeout "$ARG_TIMEOUT" \
        --interval "$ARG_POLL_INTERVAL" \
        --format text \
        "$TASK_ID" >/dev/null 2>&1
    # After wait-for-tasks exits, fetch the final state
    local state_json
    state_json="$(curl -s "${CP_URL_RESOLVED}/tasks/${TASK_ID}")"
    printf '%s' "$state_json" | python3 -c 'import json,sys; print(json.load(sys.stdin).get("state","unknown"))' 2>/dev/null || echo "unknown"
}

wait_via_curl_poll() {
    local elapsed=0
    local state="pending"
    local terminal_states="succeeded failed timed_out cancelled"

    while true; do
        local response
        response="$(curl -s "${CP_URL_RESOLVED}/tasks/${TASK_ID}" 2>/dev/null)" || true
        state="$(printf '%s' "$response" | python3 -c 'import json,sys; print(json.load(sys.stdin).get("state","unknown"))' 2>/dev/null || echo "unknown")"

        info "  state=${state} (${elapsed}s elapsed)"

        # Check if terminal
        for ts in $terminal_states; do
            if [[ "$state" == "$ts" ]]; then
                echo "$state"
                return 0
            fi
        done

        # Check timeout
        if [[ "$elapsed" -ge "$ARG_TIMEOUT" ]]; then
            echo "timed_out"
            return 0
        fi

        sleep "$ARG_POLL_INTERVAL"
        elapsed=$((elapsed + ARG_POLL_INTERVAL))
    done
}

if [[ -n "$WAIT_FOR_TASKS_BIN" ]]; then
    FINAL_STATE="$(wait_via_wait_for_tasks)"
else
    warn "wait-for-tasks not found; using curl polling fallback"
    FINAL_STATE="$(wait_via_curl_poll)"
fi

info "Task ${TASK_ID} reached state: ${FINAL_STATE}"

# ---------------------------------------------------------------------------
# Handle terminal state
# ---------------------------------------------------------------------------
if [[ "$FINAL_STATE" != "succeeded" ]]; then
    write_failure "$TASK_ID" "$FINAL_STATE" "task did not succeed"
fi

# ---------------------------------------------------------------------------
# Fetch logs and write output
# ---------------------------------------------------------------------------
info "Fetching task logs..."
LOGS_RESPONSE="$(curl -s "${CP_URL_RESOLVED}/tasks/${TASK_ID}/logs" 2>/dev/null)" || true

if [[ -z "$LOGS_RESPONSE" ]]; then
    warn "Empty logs response from CP; output file will be empty"
fi

printf '%s\n' "$LOGS_RESPONSE" > "$ARG_OUTPUT"

# Write .done sentinel
{
    echo "task_id=${TASK_ID}"
    echo "completed_at=$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
} > "${ARG_OUTPUT}.done"

# Remove tracking file
rm -f "${OUTPUT_DIR}/.bg-model-call-${TASK_ID}"

info "Done. Result written to ${ARG_OUTPUT}"
info "Sentinel written to ${ARG_OUTPUT}.done"
exit 0