#!/usr/bin/env bash
# ╔═══════════════════════════════════════════════════════════════════════════╗
# ║  Antigravity MITM LS Wrapper                                            ║
# ║                                                                         ║
# ║  This script replaces the real Antigravity language server binary.       ║
# ║  It injects HTTPS_PROXY and NODE_EXTRA_CA_CERTS environment variables   ║
# ║  so the MITM proxy can intercept LS<->API traffic.                      ║
# ║                                                                         ║
# ║  Install:   ./mitm-wrapper.sh install                                   ║
# ║  Uninstall: ./mitm-wrapper.sh uninstall                                 ║
# ║  (No args = act as wrapper, exec the real binary with injected env)     ║
# ╚═══════════════════════════════════════════════════════════════════════════╝
set -euo pipefail

# ── Config ────────────────────────────────────────────────────────────────────
# Resolve the real user's home (not /root when running under sudo)
if [[ -n "${SUDO_USER:-}" ]]; then
    REAL_HOME="$(getent passwd "$SUDO_USER" | cut -d: -f6)"
else
    REAL_HOME="$HOME"
fi
MITM_PORT_FILE="${REAL_HOME}/.config/antigravity-proxy/mitm-port"
if [[ -n "${ANTIGRAVITY_MITM_PORT:-}" ]]; then
    MITM_PORT="$ANTIGRAVITY_MITM_PORT"
elif [[ -f "$MITM_PORT_FILE" ]]; then
    MITM_PORT="$(cat "$MITM_PORT_FILE" 2>/dev/null || echo 8742)"
else
    MITM_PORT="8742"
fi
CA_PATH="${REAL_HOME}/.config/antigravity-proxy/mitm-ca.pem"

# Antigravity LS — discovered dynamically from running processes.
# Hardcoded paths are only used as a fallback if no LS process is running.
LS_FALLBACK_DIRS=(
    "/usr/share/antigravity/resources/app/extensions/antigravity/bin"
    "${REAL_HOME}/.antigravity/extensions/antigravity.antigravity-*/dist/bundled/language-server/bin"
    "${REAL_HOME}/.cursor/extensions/antigravity.antigravity-*/dist/bundled/language-server/bin"
    "${REAL_HOME}/.vscode/extensions/antigravity.antigravity-*/dist/bundled/language-server/bin"
    "/opt/antigravity/language-server/bin"
)

BACKUP_SUFFIX=".real"

# ── Colors ────────────────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

# ── Find LS binary ───────────────────────────────────────────────────────────
find_ls_binary() {
    # Method 1: Find from running process via /proc
    if [[ -d /proc ]]; then
        for pid_dir in /proc/[0-9]*; do
            local exe_target
            exe_target="$(readlink "${pid_dir}/exe" 2>/dev/null)" || continue
            # Strip " (deleted)" suffix that appears when the binary was unlinked
            exe_target="${exe_target% (deleted)}"
            if [[ "$exe_target" == *language_server_linux* ]] || \
               [[ "$exe_target" == *antigravity-language-server* ]]; then
                # FIX: If the running process is the backup (.real), strip the suffix
                # so we return the path to the base binary name.
                echo "${exe_target%$BACKUP_SUFFIX}"
                return 0
            fi
        done
    fi

    # Method 2: Fallback — scan known directories for common binary names
    local bin_names=("language_server_linux_x64" "language_server_linux_arm64" "antigravity-language-server")
    for dir_pattern in "${LS_FALLBACK_DIRS[@]}"; do
        for dir in $dir_pattern; do
            [[ -d "$dir" ]] || continue
            for name in "${bin_names[@]}"; do
                local path="${dir}/${name}"
                if [[ -f "$path" || -f "${path}${BACKUP_SUFFIX}" ]]; then
                    echo "$path"
                    return 0
                fi
            done
        done
    done
    return 1
}

# ── Install ──────────────────────────────────────────────────────────────────
cmd_install() {
    # Find the LS binary first (quiet, just to check permissions)
    local ls_path
    ls_path=$(find_ls_binary) || ls_path="${1:-}"

    # Allow override
    if [[ -n "${1:-}" ]]; then
        ls_path="$1"
    fi

    # Check permissions upfront — re-exec with sudo before doing anything
    if [[ -n "$ls_path" ]]; then
        local ls_dir
        ls_dir="$(dirname "$ls_path")"
        if [[ ! -w "$ls_dir" ]] && [[ "$EUID" -ne 0 ]]; then
            echo -e "  ${RED}✗${NC} ${ls_dir} requires elevated permissions"
            echo -e "  run: sudo $0 install ${1:-}"
            exit 1
        fi
    fi

    echo -e "${BOLD}${CYAN}Antigravity MITM Wrapper Installer${NC}"
    echo -e "───────────────────────────────────"
    echo ""

    # Find the LS binary (for real this time, with output)
    if [[ -z "$ls_path" ]]; then
        echo -e "  ${RED}✗${NC} Could not find Antigravity language server binary."
        echo -e "    No LS process found in /proc, and fallback paths didn't match."
        echo ""
        echo -e "  Set the path manually:"
        echo -e "    $0 install /path/to/language_server_linux_x64"
        exit 1
    fi
    echo -e "  ${GREEN}✓${NC} Found LS: ${ls_path}"

    local real_path="${ls_path}${BACKUP_SUFFIX}"
    local wrapper_dir
    wrapper_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    local wrapper_src="${wrapper_dir}/mitm-wrapper.sh"

    # Verify the binary exists and is not already wrapped
    if [[ -f "$real_path" ]]; then
        echo -e "  ${YELLOW}!${NC} Already installed (backup exists at ${real_path})"
        echo -e "  Run '$0 uninstall' first to reinstall."
        exit 0
    fi

    if [[ ! -f "$ls_path" ]]; then
        echo -e "  ${RED}✗${NC} Binary not found: ${ls_path}"
        exit 1
    fi

    # Verify it's a real binary, not already our wrapper
    if head -c 100 "$ls_path" | grep -q 'ANTIGRAVITY_MITM_PORT'; then
        echo -e "  ${YELLOW}!${NC} Already wrapped (script detected). Run '$0 uninstall' first."
        exit 0
    fi

    # Check CA cert
    if [[ ! -f "$CA_PATH" ]]; then
        echo -e "  ${YELLOW}!${NC}  CA cert not found at ${CA_PATH}"
        echo -e "    Start the proxy first to generate it."
        echo -e "    Continuing install anyway..."
    else
        echo -e "  ${GREEN}✓${NC} CA cert: ${CA_PATH}"
    fi

    # Back up the real binary
    cp -p "$ls_path" "$real_path"
    echo -e "  ${GREEN}✓${NC} Backed up real binary to ${real_path}"

    # Remove the original before writing (avoids "Text file busy" if LS is running)
    rm -f "$ls_path"

    # Create the wrapper script in-place
    tee "$ls_path" > /dev/null << 'WRAPPER_EOF'
#!/usr/bin/env bash
# Antigravity MITM LS Wrapper — auto-generated, do not edit.
# The LS is a Go binary — it reads HTTPS_PROXY and SSL_CERT_FILE (not NODE_EXTRA_CA_CERTS).
# Go's gRPC library also reads GRPC_DEFAULT_SSL_ROOTS_FILE_PATH for root certs.
# We build a combined CA bundle (system CAs + MITM CA) and inject it.

REAL_BINARY="${BASH_SOURCE[0]}.real"

if [[ ! -f "$REAL_BINARY" ]]; then
    echo "ERROR: Real LS binary not found at $REAL_BINARY" >&2
    echo "Run 'mitm-wrapper.sh uninstall' and reinstall." >&2
    exit 1
fi

# Inject MITM proxy (don't override if already set)
export HTTPS_PROXY="${HTTPS_PROXY:-http://127.0.0.1:__MITM_PORT__}"

# Build combined CA bundle: system CAs + MITM CA
MITM_CA="__CA_PATH__"
COMBINED_CA="/tmp/antigravity-mitm-combined-ca.pem"
if [[ -f "$MITM_CA" ]]; then
    # Find system CA bundle
    SYS_CA=""
    for candidate in /etc/ssl/certs/ca-certificates.crt /etc/pki/tls/certs/ca-bundle.crt /etc/ssl/cert.pem; do
        if [[ -f "$candidate" ]]; then
            SYS_CA="$candidate"
            break
        fi
    done
    if [[ -n "$SYS_CA" ]]; then
        cat "$SYS_CA" "$MITM_CA" > "$COMBINED_CA" 2>/dev/null
        export SSL_CERT_FILE="$COMBINED_CA"
        # Go's gRPC library may use this instead of SSL_CERT_FILE
        export GRPC_DEFAULT_SSL_ROOTS_FILE_PATH="$COMBINED_CA"
    fi
fi

exec "$REAL_BINARY" "$@"
WRAPPER_EOF

    # Substitute actual values
    sed -i "s|__MITM_PORT__|${MITM_PORT}|g" "$ls_path"
    sed -i "s|__CA_PATH__|${CA_PATH}|g" "$ls_path"

    # Make executable
    chmod +x "$ls_path"

    echo -e "  ${GREEN}✓${NC} Wrapper installed at ${ls_path}"
    echo ""
    echo -e "  ${BOLD}How it works:${NC}"
    echo -e "    When Antigravity starts the LS, the wrapper will:"
    echo -e "    1. Set ${CYAN}HTTPS_PROXY${NC}=http://127.0.0.1:${MITM_PORT}"
    echo -e "    2. Build combined CA bundle (system + MITM) at /tmp/antigravity-mitm-combined-ca.pem"
    echo -e "    3. Set ${CYAN}SSL_CERT_FILE${NC} to the combined bundle"
    echo -e "    4. Exec the real LS binary with all original args"
    echo ""
    echo -e "  ${YELLOW}Note:${NC} Restart Antigravity for the wrapper to take effect."
    echo ""
}

# ── Uninstall ────────────────────────────────────────────────────────────────
cmd_uninstall() {
    # Check permissions upfront
    local ls_path
    ls_path=$(find_ls_binary) || true
    if [[ -n "$ls_path" ]] && [[ ! -w "$(dirname "$ls_path")" ]] && [[ "$EUID" -ne 0 ]]; then
        echo -e "  ${RED}✗${NC} $(dirname "$ls_path") requires elevated permissions"
        echo -e "  run: sudo $0 uninstall"
        exit 1
    fi

    echo -e "${BOLD}${CYAN}Antigravity MITM Wrapper Uninstaller${NC}"
    echo -e "─────────────────────────────────────"
    echo ""

    if [[ -n "$ls_path" ]]; then
        local real_path="${ls_path}${BACKUP_SUFFIX}"
        if [[ -f "$real_path" ]]; then
            mv -f "$real_path" "$ls_path"
            echo -e "  ${GREEN}✓${NC} Restored real binary at ${ls_path}"
        else
            echo -e "  ${YELLOW}!${NC} No backup found at ${real_path}"
            echo -e "    The LS binary may not be wrapped."
        fi
    else
        echo -e "  ${RED}✗${NC} Could not find Antigravity language server binary."
    fi

    echo ""
    echo -e "  ${YELLOW}Note:${NC} Restart Antigravity for the change to take effect."
    echo ""
}

# ── Status ───────────────────────────────────────────────────────────────────
cmd_status() {
    echo -e "${BOLD}${CYAN}Antigravity MITM Wrapper Status${NC}"
    echo -e "────────────────────────────────"
    echo ""

    local ls_path
    if ls_path=$(find_ls_binary); then
        echo -e "  ${GREEN}✓${NC} LS binary: ${ls_path}"

        local real_path="${ls_path}${BACKUP_SUFFIX}"
        if [[ -f "$real_path" ]]; then
            echo -e "  ${GREEN}✓${NC} Wrapper: ${BOLD}installed${NC}"
            echo -e "  ${GREEN}✓${NC} Real binary: ${real_path}"

            # Check if wrapper is valid
            if head -c 200 "$ls_path" | grep -q 'MITM LS Wrapper'; then
                echo -e "  ${GREEN}✓${NC} Wrapper script: valid"
            else
                echo -e "  ${RED}✗${NC} Wrapper script: ${BOLD}corrupted or replaced${NC}"
            fi
        else
            echo -e "  ${YELLOW}○${NC} Wrapper: ${BOLD}not installed${NC}"
        fi
    else
        echo -e "  ${RED}✗${NC} LS binary: not found"
    fi

    # Check CA cert
    if [[ -f "$CA_PATH" ]]; then
        echo -e "  ${GREEN}✓${NC} CA cert: ${CA_PATH}"
    else
        echo -e "  ${RED}✗${NC} CA cert: not found (start proxy first)"
    fi

    # Check MITM port
    if ss -tlnp 2>/dev/null | grep -q ":${MITM_PORT} "; then
        echo -e "  ${GREEN}✓${NC} MITM proxy: listening on :${MITM_PORT}"
    else
        echo -e "  ${YELLOW}○${NC} MITM proxy: not running on :${MITM_PORT}"
    fi

    echo ""
}

# ── Main ─────────────────────────────────────────────────────────────────────
case "${1:-}" in
    install)
        shift
        cmd_install "${1:-}"
        ;;
    uninstall)
        cmd_uninstall
        ;;
    status)
        cmd_status
        ;;
    -h|--help)
        echo "Usage: $0 {install|uninstall|status}"
        echo ""
        echo "Commands:"
        echo "  install [path]   Install MITM wrapper (auto-detect or specify path)"
        echo "  uninstall        Restore original LS binary"
        echo "  status           Show wrapper installation status"
        echo ""
        echo "Environment:"
        echo "  ANTIGRAVITY_MITM_PORT  MITM proxy port (default: 8742)"
        ;;
    *)
        echo "Usage: $0 {install|uninstall|status}"
        exit 1
        ;;
esac