#!/bin/bash

# CMtec CMDR Joystick 25 Install Script
# Supports building, testing, and deploying RP2040 firmware

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Show usage information
show_usage() {
    echo -e "${GREEN}CMtec CMDR Joystick 25 Install Script${NC}"
    echo "========================================"
    echo
    echo "Usage: $0 [flash|test|check|clean] [options]"
    echo
    echo "Commands:"
    echo "  flash        Build and flash RP2040 firmware"
    echo "  test         Run comprehensive test suite"
    echo "  check        Quick compilation and lint checks"
    echo "  clean        Clean build artifacts and temp files"
    echo
    echo "flash options:"
    echo "  --local                    Build + flash locally (RP2040 mass storage)"
    echo "  --ssh                      Build + transfer UF2 via SSH"
    echo "  --mount PATH               Remote RP2040 mount (default /Volumes/RPI-RP2)"
    echo
    echo "SSH options (when using --ssh):"
    echo "  --target user@host        Combined user and host"
    echo "  --user USER               SSH username"
    echo "  --host HOST               SSH hostname or IP"
    echo "  --port PORT               SSH port (default 22)"
    echo "  --key PATH                SSH private key path"
    echo
    echo "Examples:"
    echo "  $0                        # Show this help"
    echo "  $0 test                   # Run all tests"
    echo "  $0 check                  # Quick compilation check"
    echo "  $0 flash --local          # Build and flash firmware locally"
    echo "  $0 flash --ssh --target user@host --mount /Volumes/RPI-RP2"
    echo "  $0 clean                  # Clean build artifacts"
    exit "${1:-0}"
}

# Parse command line arguments
COMMAND="${1:-}"
shift 1 || true

# If no command provided, show usage
if [ -z "$COMMAND" ]; then
    show_usage 0
fi

case "$COMMAND" in
flash | test | check | clean) ;;
-h | --help | help)
    show_usage 0
    ;;
*)
    echo -e "${RED}Error: Unknown command '$COMMAND'${NC}"
    echo
    show_usage 1
    ;;
esac

echo -e "${GREEN}CMtec CMDR Joystick 25 Install Script${NC}"
echo "========================================"
echo -e "${BLUE}Mode: $COMMAND${NC}"
echo

# Get script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Ensure cargo-binutils is available for `cargo objcopy`
ensure_cargo_binutils() {
    if ! command -v cargo-objcopy >/dev/null 2>&1; then
        echo "Installing cargo-binutils and llvm-tools-preview (required for UF2 conversion)..."
        rustup component add llvm-tools-preview || true
        if ! cargo install cargo-binutils; then
            echo -e "${RED}Error: Failed to install cargo-binutils${NC}"
            echo "Install manually: rustup component add llvm-tools-preview && cargo install cargo-binutils"
            exit 1
        fi
    fi
}

# Function to check prerequisites
check_prerequisites() {
    echo "Checking prerequisites..."

    if ! command -v cargo &>/dev/null; then
        echo -e "${RED}Error: cargo (Rust toolchain) not found${NC}"
        echo "Install Rust: https://rustup.rs/"
        exit 1
    fi

    if ! command -v python3 &>/dev/null; then
        echo -e "${RED}Error: python3 not found${NC}"
        echo "Install Python 3"
        exit 1
    fi

    # Check for required Rust target
    if ! rustup target list --installed | grep -q "thumbv6m-none-eabi"; then
        echo "Installing Rust target thumbv6m-none-eabi..."
        rustup target add thumbv6m-none-eabi
    fi

    # Check for x86_64 target for testing
    if ! rustup target list --installed | grep -q "x86_64-unknown-linux-gnu"; then
        echo "Installing Rust target x86_64-unknown-linux-gnu (for testing)..."
        rustup target add x86_64-unknown-linux-gnu
    fi

    echo -e "${GREEN}✓${NC} Prerequisites check complete"
    echo
}

# Function to run comprehensive test suite
run_tests() {
    echo "Running comprehensive test suite..."
    echo

    check_prerequisites

    # Change to RP2040 directory
    RP2040_DIR="$SCRIPT_DIR/rp2040"
    if [ ! -d "$RP2040_DIR" ]; then
        echo -e "${RED}Error: RP2040 directory not found at $RP2040_DIR${NC}"
        exit 1
    fi

    cd "$RP2040_DIR"

    # Test 1: Cargo check for embedded target
    echo "Testing embedded target compilation..."
    if ! cargo check --target thumbv6m-none-eabi; then
        echo -e "${RED}✗${NC} Embedded target compilation check failed"
        exit 1
    fi
    echo -e "${GREEN}✓${NC} Embedded target compilation check passed"

    # Test 2: Host target tests (17 comprehensive tests)
    echo "Running host target test suite..."
    if ! cargo test --lib --target x86_64-unknown-linux-gnu --features std; then
        echo -e "${RED}✗${NC} Host target tests failed"
        exit 1
    fi
    echo -e "${GREEN}✓${NC} Host target tests passed (17 tests)"

    # Test 3: Clippy for code quality (warnings allowed for now)
    echo "Running clippy code quality checks..."
    if ! cargo clippy --target thumbv6m-none-eabi; then
        echo -e "${RED}✗${NC} Clippy code quality check failed"
        exit 1
    fi
    echo -e "${GREEN}✓${NC} Clippy code quality check passed"

    # Test 4: Release build test
    echo "Testing release build..."
    if ! cargo build --release --target thumbv6m-none-eabi; then
        echo -e "${RED}✗${NC} Release build failed"
        exit 1
    fi
    echo -e "${GREEN}✓${NC} Release build passed"

    cd "$SCRIPT_DIR"

    echo
    echo -e "${GREEN}All tests completed successfully!${NC}"
    echo
    echo "Test summary:"
    echo -e "${GREEN}✓${NC} Embedded target compilation check"
    echo -e "${GREEN}✓${NC} Host target test suite (17 tests: expo + storage modules)"
    echo -e "${GREEN}✓${NC} Clippy code quality checks"
    echo -e "${GREEN}✓${NC} Release build verification"
    echo
    echo -e "${GREEN}Codebase is ready for deployment.${NC}"
}

# Function to run quick checks
run_check() {
    echo "Running quick compilation and lint checks..."
    echo

    check_prerequisites

    # Change to RP2040 directory
    RP2040_DIR="$SCRIPT_DIR/rp2040"
    if [ ! -d "$RP2040_DIR" ]; then
        echo -e "${RED}Error: RP2040 directory not found at $RP2040_DIR${NC}"
        exit 1
    fi

    cd "$RP2040_DIR"

    # Quick cargo check
    echo "Checking embedded target compilation..."
    if ! cargo check --target thumbv6m-none-eabi; then
        echo -e "${RED}✗${NC} Compilation check failed"
        exit 1
    fi
    echo -e "${GREEN}✓${NC} Compilation check passed"

    # Quick clippy check (warnings allowed for now)
    echo "Running clippy..."
    if ! cargo clippy --target thumbv6m-none-eabi; then
        echo -e "${RED}✗${NC} Clippy check failed"
        exit 1
    fi
    echo -e "${GREEN}✓${NC} Clippy check passed"

    cd "$SCRIPT_DIR"

    echo
    echo -e "${GREEN}Quick checks completed successfully!${NC}"
}

# Function to clean build artifacts
clean_build() {
    echo "Cleaning build artifacts and temporary files..."
    echo

    RP2040_DIR="$SCRIPT_DIR/rp2040"
    if [ ! -d "$RP2040_DIR" ]; then
        echo -e "${RED}Error: RP2040 directory not found at $RP2040_DIR${NC}"
        exit 1
    fi

    cd "$RP2040_DIR"

    # Clean cargo artifacts
    echo "Cleaning cargo build artifacts..."
    cargo clean
    echo -e "${GREEN}✓${NC} Cargo artifacts cleaned"

    # Remove UF2 and binary files
    if [ -f "firmware.uf2" ]; then
        rm -f firmware.uf2
        echo -e "${GREEN}✓${NC} Removed firmware.uf2"
    fi

    if [ -f "target/thumbv6m-none-eabi/release/cmdr-joystick-25.bin" ]; then
        rm -f target/thumbv6m-none-eabi/release/cmdr-joystick-25.bin
        echo -e "${GREEN}✓${NC} Removed binary files"
    fi

    # Clean any temp files
    rm -f /tmp/cmdr_*.tmp 2>/dev/null || true

    cd "$SCRIPT_DIR"

    echo
    echo -e "${GREEN}Cleanup completed!${NC}"
}

# Function to build firmware locally
build_firmware() {
    RP2040_DIR="$SCRIPT_DIR/rp2040"

    if [ ! -d "$RP2040_DIR" ]; then
        echo -e "${RED}Error: RP2040 directory not found at $RP2040_DIR${NC}"
        exit 1
    fi

    check_prerequisites

    # Build firmware
    echo "Building RP2040 firmware..."
    cd "$RP2040_DIR"

    if ! cargo build --release --target thumbv6m-none-eabi; then
        echo -e "${RED}Error: Failed to build firmware${NC}"
        exit 1
    fi

    echo -e "${GREEN}✓${NC} Firmware build complete"
    echo

    # Create UF2 file
    echo "Converting to UF2 format..."

    ensure_cargo_binutils
    if ! cargo objcopy --release --target thumbv6m-none-eabi -- -O binary target/thumbv6m-none-eabi/release/cmdr-joystick-25.bin; then
        echo -e "${RED}Error: Failed to create binary file${NC}"
        exit 1
    fi

    # Check for uf2conv.py script
    if [ ! -f "uf2conv.py" ]; then
        echo "Downloading uf2conv.py..."
        if ! curl -L -o uf2conv.py "https://raw.githubusercontent.com/microsoft/uf2/master/utils/uf2conv.py"; then
            echo -e "${RED}Error: Failed to download uf2conv.py${NC}"
            exit 1
        fi
        chmod +x uf2conv.py
    fi

    if ! python3 uf2conv.py -b 0x10000000 -f 0xe48bff56 -c -o firmware.uf2 target/thumbv6m-none-eabi/release/cmdr-joystick-25.bin; then
        echo -e "${RED}Error: Failed to convert to UF2 format${NC}"
        exit 1
    fi

    echo -e "${GREEN}✓${NC} UF2 conversion complete"
    echo
}

# Function to build and transfer firmware via SSH
flash_firmware_ssh() {
    RP2040_DIR="$SCRIPT_DIR/rp2040"

    if [ ! -d "$RP2040_DIR" ]; then
        echo -e "${RED}Error: RP2040 directory not found at $RP2040_DIR${NC}"
        exit 1
    fi

    # Parse CLI args for SSH
    SSH_HOST=""; SSH_USER=""; SSH_PORT="22"; SSH_KEY=""; REMOTE_MOUNT="/Volumes/RPI-RP2"
    while [ $# -gt 0 ]; do
        case "$1" in
            --target) shift; TARGET="$1"; SSH_USER="${TARGET%@*}"; SSH_HOST="${TARGET#*@}" ;;
            --host) shift; SSH_HOST="$1" ;;
            --user) shift; SSH_USER="$1" ;;
            --port) shift; SSH_PORT="$1" ;;
            --key) shift; SSH_KEY="$1" ;;
            --mount) shift; REMOTE_MOUNT="$1" ;;
            *) echo -e "${YELLOW}Warning: Unknown SSH option: $1${NC}" ;;
        esac
        shift || true
    done

    if [ -z "$SSH_HOST" ] || [ -z "$SSH_USER" ]; then
        echo -e "${RED}Error: provide SSH credentials via --target user@host or --user/--host${NC}"
        exit 1
    fi

    # Build firmware locally first
    build_firmware

    echo "Testing SSH connection to $SSH_USER@$SSH_HOST:$SSH_PORT..."

    mkdir -p "${HOME}/.ssh"
    SSH_CONTROL_PATH="${HOME}/.ssh/cmdr_%C"
    local SSH_OPTS=(
        -p "$SSH_PORT"
        -o "StrictHostKeyChecking=accept-new"
        -o "ControlMaster=auto"
        -o "ControlPersist=60s"
        -o "ControlPath=$SSH_CONTROL_PATH"
        -o "PreferredAuthentications=publickey,password,keyboard-interactive"
    )
    if [ -n "$SSH_KEY" ]; then
        SSH_OPTS+=( -i "$SSH_KEY" )
    fi

    # Test SSH connection
    if ! ssh "${SSH_OPTS[@]}" -o ConnectTimeout=10 -fN "$SSH_USER@$SSH_HOST" 2>/dev/null; then
        echo -e "${RED}Error: Cannot connect to $SSH_USER@$SSH_HOST:$SSH_PORT${NC}"
        echo "Tips:"
        echo "- Verify network reachability and SSH port"
        echo "- Ensure SSH keys are properly configured"
        exit 1
    fi

    echo -e "${GREEN}✓${NC} SSH connection successful"

    # Check if remote mount exists
    echo "Checking remote RP2040 mount path: $REMOTE_MOUNT"
    if ! ssh "${SSH_OPTS[@]}" "$SSH_USER@$SSH_HOST" "test -d \"$REMOTE_MOUNT\""; then
        echo -e "${RED}Remote mount not found: $REMOTE_MOUNT${NC}"
        echo -e "${YELLOW}Ensure the RP2040 is in BOOTSEL mode and mounted on the remote host.${NC}"
        echo "Common paths: /Volumes/RPI-RP2 (macOS), /media/\$USER/RPI-RP2 (Linux)"
        exit 1
    fi

    echo "Transferring firmware.uf2 to remote host..."
    local SCP_OPTS=( -P "$SSH_PORT" -o "StrictHostKeyChecking=accept-new" -o "ControlPath=$SSH_CONTROL_PATH" )
    if [ -n "$SSH_KEY" ]; then
        SCP_OPTS+=( -i "$SSH_KEY" )
    fi

    cd "$RP2040_DIR"
    if ! scp "${SCP_OPTS[@]}" firmware.uf2 "$SSH_USER@$SSH_HOST:$REMOTE_MOUNT/firmware.uf2"; then
        echo -e "${RED}✗${NC} Failed to transfer firmware to remote RP2040"
        exit 1
    fi

    echo -e "${GREEN}✓${NC} Transferred firmware.uf2 to $SSH_USER@$SSH_HOST:$REMOTE_MOUNT"

    cd "$SCRIPT_DIR"

    echo
    echo -e "${GREEN}SSH firmware deployment completed!${NC}"
    echo "Remote: $SSH_USER@$SSH_HOST:$REMOTE_MOUNT/firmware.uf2"
    echo -e "${YELLOW}The remote RP2040 should now restart with new firmware.${NC}"
}

# Function to flash firmware locally
flash_firmware_local() {
    build_firmware

    RP2040_DIR="$SCRIPT_DIR/rp2040"
    cd "$RP2040_DIR"

    # Detect OS and set up mount detection
    OS="$(uname -s)"

    # Wait for and flash to RP2040
    echo "Waiting for RP2040 in bootloader mode..."
    echo -e "${YELLOW}Put your RP2040 into bootloader mode (hold BOOTSEL button while plugging in USB)${NC}"

    case "$OS" in
    Linux*)
        # Wait for RPI-RP2 mount point with timeout
        MAX_WAIT=120
        waited=0
        while [ ! -d "/media/RPI-RP2" ] && [ ! -d "/media/$USER/RPI-RP2" ] && [ ! -d "/run/media/$USER/RPI-RP2" ]; do
            sleep 1
            waited=$((waited + 1))
            if [ "$waited" -ge "$MAX_WAIT" ]; then
                echo -e "${RED}Timeout waiting for RP2040 mount (RPI-RP2)${NC}"
                echo "Ensure the device is in BOOTSEL mode and mounted (try replugging with BOOTSEL held)."
                echo "Expected mount at one of: /media/RPI-RP2, /media/$USER/RPI-RP2, /run/media/$USER/RPI-RP2"
                exit 1
            fi
        done

        # Find the actual mount point
        if [ -d "/media/RPI-RP2" ]; then
            MOUNT_POINT="/media/RPI-RP2"
        elif [ -d "/media/$USER/RPI-RP2" ]; then
            MOUNT_POINT="/media/$USER/RPI-RP2"
        elif [ -d "/run/media/$USER/RPI-RP2" ]; then
            MOUNT_POINT="/run/media/$USER/RPI-RP2"
        else
            echo -e "${RED}Error: Could not find RPI-RP2 mount point${NC}"
            exit 1
        fi

        echo -e "${GREEN}✓${NC} Found RP2040 at $MOUNT_POINT"

        # Flash firmware
        echo "Flashing firmware..."
        cp firmware.uf2 "$MOUNT_POINT/"
        echo -e "${GREEN}✓${NC} Firmware flashed successfully!"
        ;;
    Darwin*)
        # Wait for RPI-RP2 volume with timeout
        MAX_WAIT=120
        waited=0
        while [ ! -d "/Volumes/RPI-RP2" ]; do
            sleep 1
            waited=$((waited + 1))
            if [ "$waited" -ge "$MAX_WAIT" ]; then
                echo -e "${RED}Timeout waiting for RP2040 mount at /Volumes/RPI-RP2${NC}"
                echo "Ensure the device is in BOOTSEL mode and mounted (try replugging with BOOTSEL held)."
                exit 1
            fi
        done

        echo -e "${GREEN}✓${NC} Found RP2040 at /Volumes/RPI-RP2"

        # Flash firmware
        echo "Flashing firmware..."
        cp firmware.uf2 "/Volumes/RPI-RP2/"
        echo -e "${GREEN}✓${NC} Firmware flashed successfully!"
        ;;
    *)
        echo -e "${RED}Error: Unsupported operating system for local flashing${NC}"
        echo "This script supports Linux and macOS for local flashing."
        exit 1
        ;;
    esac

    cd "$SCRIPT_DIR"

    echo
    echo -e "${GREEN}Firmware installation completed!${NC}"
    echo
    echo "Firmware file: $RP2040_DIR/firmware.uf2"
    echo -e "${GREEN}RP2040 firmware deployed successfully!${NC}"
    echo -e "${YELLOW}The device should now restart with new firmware.${NC}"
}

# Detect OS
case "$(uname -s)" in
Linux*)
    OS="Linux"
    ;;
Darwin*)
    OS="macOS"
    ;;
*)
    echo -e "${RED}Error: Unsupported operating system${NC}"
    echo "This script supports Linux and macOS only."
    exit 1
    ;;
esac

echo "Detected OS: $OS"
if [ "$COMMAND" = "flash" ]; then
    echo "RP2040 firmware directory: $SCRIPT_DIR/rp2040"
elif [ "$COMMAND" = "test" ]; then
    echo "Testing Rust embedded firmware with comprehensive test suite"
fi
echo

# Execute command
case "$COMMAND" in
flash)
    MODE=""
    REM_ARGS=()
    for arg in "$@"; do
        case "$arg" in
        --local) MODE="local" ;;
        --ssh) MODE="ssh" ;;
        *) REM_ARGS+=("$arg") ;;
        esac
    done
    if [ "$MODE" = "local" ]; then
        flash_firmware_local
    elif [ "$MODE" = "ssh" ]; then
        flash_firmware_ssh "${REM_ARGS[@]}"
    else
        echo -e "${RED}Error: specify --local or --ssh for flash${NC}"
        show_usage 1
    fi
    ;;
test)
    run_tests
    ;;
check)
    run_check
    ;;
clean)
    clean_build
    ;;
esac