#!/usr/bin/env bash

set -e

CONFIG_DIR_DEFAULT="${HOME}/.config/auggie-container"
DATA_DIR_DEFAULT="${HOME}/.local/share/auggie"
WORKSPACE_DIR="${PWD}"
CONFIG_DIR="$CONFIG_DIR_DEFAULT"
DATA_DIR="$DATA_DIR_DEFAULT"
PULL_LATEST=false
SHELL_MODE=false
FORWARD_AUTH=false
AUTH_PORT=1455
ENABLE_PROXY=false
ENABLE_DATASETTE=false
PROXY_MODE="reverse"
PROXY_PORT="${AUGGIE_PROXY_PORT:-8080}"
DATASETTE_PORT="${AUGGIE_DATASETTE_PORT:-8001}"
TARGET_API_URL="${AUGGIE_TARGET_API_URL:-}"
PROXY_DATA_DIR=""
NETWORK_NAME="auggie-network"
PROXY_CONTAINER_NAME="auggie-proxy"
DATASETTE_CONTAINER_NAME="auggie-datasette"

show_version() {
    echo "auggie-container ${AUGGIE_CONTAINER_VERSION:-dev}"
}

show_usage() {
    cat << 'EOH'
Usage: auggie-container [OPTIONS] [COMMAND]

Run auggie CLI in a containerized environment.

OPTIONS:
    -w, --workspace <path>      Set workspace directory (default: current directory)
    -c, --config <path>         Set config directory (default: ~/.config/auggie-container)
    -d, --data <path>           Set data directory (default: ~/.local/share/auggie)
    --proxy                     Enable API logging via reverse proxy
    --proxy-all                 Intercept all CLI traffic via forward proxy (mitmproxy)
    --target-api-url <url>      Target API URL for reverse proxy (e.g. https://api.augmentcode.com)
    --datasette                 Start Datasette web UI for logs (implies --proxy)
    --pull                      Pull latest images before running
    --shell                     Drop into container shell instead of starting auggie
    --auth-forward              Publish auth server port to host for auggie login
    --auth-port <port>          Override auth server port (default: 1455)
    --version                   Show version
    -h, --help                  Show this help message

COMMAND:
    Any command to pass to auggie (default: interactive mode)

EXAMPLES:
    auggie-container
    auggie-container --pull
    auggie-container -w /path/to/project
    auggie-container --shell
    auggie-container --auth-forward
    auggie-container --proxy --target-api-url https://api.augmentcode.com
    auggie-container --proxy-all
    auggie-container --proxy --datasette --target-api-url https://api.augmentcode.com

EOH
}

while [[ $# -gt 0 ]]; do
    case $1 in
        -w|--workspace)
            WORKSPACE_DIR="$2"
            shift 2
            ;;
        -c|--config)
            CONFIG_DIR="$2"
            shift 2
            ;;
        -d|--data)
            DATA_DIR="$2"
            shift 2
            ;;
        --proxy)
            ENABLE_PROXY=true
            shift
            ;;
        --proxy-all)
            ENABLE_PROXY=true
            PROXY_MODE="forward"
            shift
            ;;
        --target-api-url)
            TARGET_API_URL="$2"
            shift 2
            ;;
        --datasette)
            ENABLE_PROXY=true
            ENABLE_DATASETTE=true
            shift
            ;;
        --pull)
            PULL_LATEST=true
            shift
            ;;
        --shell)
            SHELL_MODE=true
            shift
            ;;
        --auth-forward)
            FORWARD_AUTH=true
            shift
            ;;
        --auth-port)
            AUTH_PORT="$2"
            shift 2
            ;;
        --version)
            show_version
            exit 0
            ;;
        -h|--help)
            show_usage
            exit 0
            ;;
        *)
            break
            ;;
    esac
done

if [[ -z "$PROXY_DATA_DIR" ]]; then
    PROXY_DATA_DIR="${CONFIG_DIR}/proxy"
fi
PROXY_CERTS_DIR="${PROXY_DATA_DIR}/certs"

mkdir -p "${CONFIG_DIR}"
mkdir -p "${DATA_DIR}"

VERSION="${AUGGIE_CONTAINER_VERSION:-latest}"
IMAGE="${AUGGIE_IMAGE:-nezhar/auggie-cli:${VERSION}}"
PROXY_IMAGE="${AUGGIE_PROXY_IMAGE:-nezhar/auggie-proxy:${VERSION}}"
DATASETTE_IMAGE="${AUGGIE_DATASETTE_IMAGE:-nezhar/auggie-datasette:${VERSION}}"

if [[ "$PULL_LATEST" == true ]]; then
    docker pull "$IMAGE"
fi

DOCKER_ARGS=(
    --rm -it
    -e "USER_UID=$(id -u)"
    -e "USER_GID=$(id -g)"
    -e "HOME=/config"
    -e "XDG_DATA_HOME=/data"
    -v "${WORKSPACE_DIR}:/workspace"
    -v "${CONFIG_DIR}:/config"
    -v "${DATA_DIR}:/data"
)

# Pass through AUGMENT_SESSION_AUTH if set
if [[ -n "${AUGMENT_SESSION_AUTH:-}" ]]; then
    DOCKER_ARGS+=(
        -e "AUGMENT_SESSION_AUTH=${AUGMENT_SESSION_AUTH}"
    )
fi

if [[ "$ENABLE_PROXY" == true ]]; then
    if [[ "$PROXY_MODE" != "forward" ]]; then
        if [[ -z "$TARGET_API_URL" ]]; then
            echo "Error: --target-api-url is required when using --proxy" >&2
            exit 1
        fi
        echo "Starting with API logging enabled (reverse proxy to ${TARGET_API_URL})"
    else
        echo "Starting with API logging enabled (forward proxy mode)"
    fi

    mkdir -p "${PROXY_DATA_DIR}" "${PROXY_CERTS_DIR}"

    docker network create "$NETWORK_NAME" 2>/dev/null || true

    docker run -d --rm \
        --name "$PROXY_CONTAINER_NAME" \
        --network "$NETWORK_NAME" \
        -p "${PROXY_PORT}:${PROXY_PORT}" \
        -v "${PROXY_DATA_DIR}:/proxy/logs" \
        -v "${PROXY_CERTS_DIR}:/proxy/certs" \
        -e "TARGET_API_URL=${TARGET_API_URL}" \
        -e "PROXY_PORT=${PROXY_PORT}" \
        -e "PROXY_MODE=${PROXY_MODE}" \
        "$PROXY_IMAGE" 2>/dev/null || true

    if [[ "$ENABLE_DATASETTE" == true ]]; then
        docker run -d --rm \
            --name "$DATASETTE_CONTAINER_NAME" \
            --network "$NETWORK_NAME" \
            -p "${DATASETTE_PORT}:8001" \
            -v "${PROXY_DATA_DIR}:/data" \
            "$DATASETTE_IMAGE" 2>/dev/null || true

        echo "Datasette available at http://localhost:${DATASETTE_PORT}"
    fi

    sleep 2
fi

if [[ "$FORWARD_AUTH" == true ]]; then
    DOCKER_ARGS+=(
        -p "${AUTH_PORT}:${AUTH_PORT}"
    )
fi

if [[ "$ENABLE_PROXY" == true ]]; then
    CA_CERT_PATH=""
    if [[ -f "${PROXY_CERTS_DIR}/mitmproxy-ca-cert.pem" ]]; then
        CA_CERT_PATH="/mitmproxy/mitmproxy-ca-cert.pem"
    elif [[ -f "${PROXY_CERTS_DIR}/mitmproxy-ca.pem" ]]; then
        CA_CERT_PATH="/mitmproxy/mitmproxy-ca.pem"
    else
        echo "Warning: mitmproxy CA cert not found yet in ${PROXY_CERTS_DIR}"
        echo "         HTTPS requests may fail until the cert is generated."
    fi

    DOCKER_ARGS+=(
        --network "$NETWORK_NAME"
        -v "${PROXY_CERTS_DIR}:/mitmproxy:ro"
    )
    if [[ -n "$CA_CERT_PATH" ]]; then
        DOCKER_ARGS+=(
            -e "NODE_EXTRA_CA_CERTS=${CA_CERT_PATH}"
        )
    fi
    if [[ "$PROXY_MODE" != "forward" ]]; then
        # Reverse proxy mode: no extra env needed, auggie talks to proxy directly
        :
    else
        DOCKER_ARGS+=(
            -e "HTTP_PROXY=http://${PROXY_CONTAINER_NAME}:${PROXY_PORT}"
            -e "HTTPS_PROXY=http://${PROXY_CONTAINER_NAME}:${PROXY_PORT}"
        )
    fi
fi

if [[ "$SHELL_MODE" == true ]]; then
    docker run "${DOCKER_ARGS[@]}" "$IMAGE" /bin/bash
else
    docker run "${DOCKER_ARGS[@]}" "$IMAGE" "$@"
fi

if [[ "$ENABLE_PROXY" == true ]]; then
    echo "Proxy services are still running. To stop them, run:"
    echo "  docker stop $PROXY_CONTAINER_NAME"
    if [[ "$ENABLE_DATASETTE" == true ]]; then
        echo "  docker stop $DATASETTE_CONTAINER_NAME"
    fi
fi