#!/usr/bin/env bash
#
# install.sh - Install Docker, Docker Compose, and local Python runtime on Ubuntu 24.04 WSL2
# Idempotent: safe to run multiple times
#

set -euo pipefail

readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly NC='\033[0m'

log_info() { echo -e "${GREEN}[INFO]${NC} $*"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*"; }

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WSL_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
AUTOMATION_DIR="$(cd "$WSL_DIR/.." && pwd)"
PROJECT_ROOT="$(cd "$AUTOMATION_DIR/.." && pwd)"
RUNTIME_DIR="$AUTOMATION_DIR/wsl_runtime"
VENV_DIR="$RUNTIME_DIR/venv"

check_sudo() {
    if [[ $EUID -eq 0 ]]; then
        log_error "This script should not be run as root. It will use sudo when needed."
        exit 1
    fi
}

detect_ubuntu() {
    if [[ ! -f /etc/os-release ]]; then
        log_error "Cannot detect OS version. /etc/os-release not found."
        exit 1
    fi

    # shellcheck disable=SC1091
    source /etc/os-release
    if [[ "${ID:-}" != "ubuntu" ]]; then
        log_warn "This script is designed for Ubuntu. Detected: ${ID:-unknown}"
    fi

    log_info "Detected Ubuntu ${VERSION_ID:-unknown}"
}

check_wsl2() {
    if [[ ! -f /proc/version ]]; then
        log_warn "Cannot verify WSL environment"
        return
    fi

    if grep -qi microsoft /proc/version; then
        log_info "Running in WSL environment"
    else
        log_warn "Not running in WSL. This script is designed for WSL2."
    fi
}

install_docker() {
    log_info "Checking Docker installation..."

    if command -v docker >/dev/null 2>&1; then
        log_info "Docker already installed: $(docker --version)"
    else
        log_info "Installing Docker..."
        sudo apt-get update -q
        sudo apt-get install -y \
            ca-certificates \
            curl \
            gnupg \
            lsb-release \
            software-properties-common

        sudo install -m 0755 -d /etc/apt/keyrings
        if [[ ! -f /etc/apt/keyrings/docker.gpg ]]; then
            curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
            sudo chmod a+r /etc/apt/keyrings/docker.gpg
        fi

        local codename
        codename=$(. /etc/os-release && echo "$VERSION_CODENAME")
        sudo tee /etc/apt/sources.list.d/docker.list >/dev/null <<EOF
deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu $codename stable
EOF

        sudo apt-get update -q
        sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
    fi

    if ! groups "$USER" | grep -q '\bdocker\b'; then
        log_info "Adding user $USER to docker group..."
        sudo usermod -aG docker "$USER"
        log_warn "A new login session may be needed for docker group membership."
    fi

    sudo systemctl enable docker
    sudo systemctl start docker
}

install_python() {
    log_info "Checking Python installation..."

    if command -v python3 >/dev/null 2>&1; then
        log_info "Python already installed: $(python3 --version)"
    else
        sudo apt-get update -q
        sudo apt-get install -y python3 python3-pip python3-venv python3-full
    fi
}

install_utilities() {
    log_info "Installing system utilities..."

    sudo apt-get update -q
    sudo apt-get install -y \
        jq \
        git \
        curl \
        wget \
        rsync \
        net-tools \
        dnsutils \
        htop \
        ncdu \
        tree \
        unzip \
        zip \
        httpie \
        python3-rich \
        pipx
}

configure_docker_wsl2() {
    log_info "Configuring Docker for WSL..."

    local docker_config_dir="/etc/docker"
    local docker_config_file="$docker_config_dir/daemon.json"

    if [[ ! -f "$docker_config_file" ]]; then
        sudo mkdir -p "$docker_config_dir"
        sudo tee "$docker_config_file" >/dev/null <<'EOF'
{
  "log-driver": "json-file",
  "log-opts": {
    "max-size": "10m",
    "max-file": "3"
  },
  "features": {
    "containerd-snapshotter": true
  },
  "iptables": false
}
EOF
        sudo systemctl restart docker
    fi

    local bashrc_file="$HOME/.bashrc"
    if ! grep -q 'WSL Docker helpers' "$bashrc_file" 2>/dev/null; then
        cat >> "$bashrc_file" <<'EOF'

# WSL Docker helpers
export DOCKER_HOST=unix:///var/run/docker.sock
EOF
    fi
}

handle_windows_paths() {
    log_info "Ensuring project symlink exists..."
    if [[ ! -L "$HOME/ableton-mcp-ai" ]]; then
        ln -sfn "$PROJECT_ROOT" "$HOME/ableton-mcp-ai"
    fi
}

install_python_dependencies() {
    log_info "Preparing local virtual environment..."
    mkdir -p "$RUNTIME_DIR"

    if [[ ! -d "$VENV_DIR" ]]; then
        python3 -m venv "$VENV_DIR"
    fi

    # shellcheck disable=SC1091
    source "$VENV_DIR/bin/activate"
    python -m pip install --upgrade pip

    local found_req=false
    local requirements_files=(
        "$PROJECT_ROOT/MCP_Server/requirements.txt"
        "$PROJECT_ROOT/requirements.txt"
    )

    for req_file in "${requirements_files[@]}"; do
        if [[ -f "$req_file" ]]; then
            log_info "Installing dependencies from: $req_file"
            python -m pip install -r "$req_file"
            found_req=true
        fi
    done

    if [[ "$found_req" == "false" ]]; then
        log_warn "No requirements.txt files found"
    fi

    deactivate
}

verify_installation() {
    log_info "Verifying installation..."

    local all_good=true

    if command -v docker >/dev/null 2>&1; then
        log_info "OK Docker: $(docker --version)"
    else
        log_error "FAIL Docker not found"
        all_good=false
    fi

    if docker compose version >/dev/null 2>&1; then
        log_info "OK Docker Compose: $(docker compose version)"
    else
        log_error "FAIL Docker Compose not found"
        all_good=false
    fi

    if command -v python3 >/dev/null 2>&1; then
        log_info "OK Python: $(python3 --version)"
    else
        log_error "FAIL Python3 not found"
        all_good=false
    fi

    if [[ -x "$VENV_DIR/bin/python" ]]; then
        log_info "OK Venv: $VENV_DIR"
    else
        log_error "FAIL Venv not found at $VENV_DIR"
        all_good=false
    fi

    if command -v jq >/dev/null 2>&1; then
        log_info "OK jq installed"
    else
        log_error "FAIL jq not found"
        all_good=false
    fi

    if [[ "$all_good" == "true" ]]; then
        log_info "All dependencies installed successfully"
        return 0
    fi

    log_error "Some dependencies failed to install"
    return 1
}

main() {
    log_info "Starting AbletonMCP-AI WSL installation..."
    echo

    check_sudo
    detect_ubuntu
    check_wsl2
    echo

    install_docker
    install_python
    install_utilities
    configure_docker_wsl2
    handle_windows_paths
    install_python_dependencies
    echo

    verify_installation
    echo

    log_info "Installation complete"
    log_info "Next step: run ./setup.sh and then ./start.sh"
}

main "$@"