#!/bin/bash
# RuvLLM ESP32 - Cross-Platform Installer
# Supports: Linux, macOS, WSL

set -e

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

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

echo -e "${BLUE}"
echo "╔══════════════════════════════════════════════════════════╗"
echo "║          RuvLLM ESP32 - Universal Installer              ║"
echo "║     Tiny LLM + RAG + Federation for Microcontrollers     ║"
echo "╚══════════════════════════════════════════════════════════╝"
echo -e "${NC}"

# Detect OS
detect_os() {
    case "$(uname -s)" in
        Linux*)     OS=linux;;
        Darwin*)    OS=macos;;
        MINGW*|MSYS*|CYGWIN*) OS=windows;;
        *)          OS=unknown;;
    esac
    echo -e "${GREEN}Detected OS: $OS${NC}"
}

# Check dependencies
check_deps() {
    echo -e "\n${YELLOW}Checking dependencies...${NC}"

    # Rust
    if command -v rustc &> /dev/null; then
        RUST_VERSION=$(rustc --version)
        echo -e "${GREEN}✓ Rust: $RUST_VERSION${NC}"
    else
        echo -e "${RED}✗ Rust not found${NC}"
        install_rust
    fi

    # Cargo
    if command -v cargo &> /dev/null; then
        echo -e "${GREEN}✓ Cargo available${NC}"
    else
        echo -e "${RED}✗ Cargo not found${NC}"
        exit 1
    fi
}

# Install Rust
install_rust() {
    echo -e "${YELLOW}Installing Rust...${NC}"
    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
    source "$HOME/.cargo/env"
}

# Install ESP32 toolchain
install_esp_toolchain() {
    echo -e "\n${YELLOW}Installing ESP32 toolchain...${NC}"

    # Install espup
    if ! command -v espup &> /dev/null; then
        echo "Installing espup..."
        cargo install espup
    else
        echo -e "${GREEN}✓ espup already installed${NC}"
    fi

    # Install ESP toolchain
    echo "Running espup install (this may take a few minutes)..."
    espup install

    # Source the export file
    if [ -f "$HOME/export-esp.sh" ]; then
        source "$HOME/export-esp.sh"
    elif [ -f "$HOME/.espressif/export-esp.sh" ]; then
        source "$HOME/.espressif/export-esp.sh"
    fi

    # Install espflash
    if ! command -v espflash &> /dev/null; then
        echo "Installing espflash..."
        cargo install espflash
    else
        echo -e "${GREEN}✓ espflash already installed${NC}"
    fi

    # Install ldproxy
    if ! command -v ldproxy &> /dev/null; then
        echo "Installing ldproxy..."
        cargo install ldproxy
    else
        echo -e "${GREEN}✓ ldproxy already installed${NC}"
    fi
}

# Build the project
build_project() {
    echo -e "\n${YELLOW}Building RuvLLM ESP32...${NC}"

    # Source ESP environment
    if [ -f "$HOME/export-esp.sh" ]; then
        source "$HOME/export-esp.sh"
    fi

    cargo build --release

    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✓ Build successful!${NC}"
    else
        echo -e "${RED}✗ Build failed${NC}"
        exit 1
    fi
}

# Flash to device
flash_device() {
    local PORT="${1:-/dev/ttyUSB0}"

    echo -e "\n${YELLOW}Flashing to $PORT...${NC}"

    # Detect port if not specified
    if [ ! -e "$PORT" ]; then
        echo "Detecting ESP32 port..."
        if [ "$OS" = "macos" ]; then
            PORT=$(ls /dev/cu.usbserial-* 2>/dev/null | head -1)
            [ -z "$PORT" ] && PORT=$(ls /dev/cu.SLAB_USBtoUART* 2>/dev/null | head -1)
        else
            PORT=$(ls /dev/ttyUSB* 2>/dev/null | head -1)
            [ -z "$PORT" ] && PORT=$(ls /dev/ttyACM* 2>/dev/null | head -1)
        fi
    fi

    if [ -z "$PORT" ] || [ ! -e "$PORT" ]; then
        echo -e "${RED}No ESP32 device found. Please specify port:${NC}"
        echo "  ./install.sh flash /dev/ttyUSB0"
        exit 1
    fi

    echo -e "${GREEN}Found device at: $PORT${NC}"

    espflash flash --port "$PORT" --monitor target/xtensa-esp32-espidf/release/ruvllm-esp32-flash
}

# Print usage
usage() {
    echo "Usage: ./install.sh [command] [options]"
    echo ""
    echo "Commands:"
    echo "  install     Install all dependencies and build (default)"
    echo "  build       Build the project only"
    echo "  flash       Flash to ESP32 (optionally specify port)"
    echo "  deps        Install dependencies only"
    echo "  cluster     Setup cluster configuration"
    echo "  help        Show this help"
    echo ""
    echo "Examples:"
    echo "  ./install.sh                    # Full install and build"
    echo "  ./install.sh flash /dev/ttyUSB0 # Flash to specific port"
    echo "  ./install.sh flash COM6         # Flash on Windows/WSL"
    echo "  ./install.sh cluster 5          # Setup 5-chip cluster"
}

# Cluster setup
setup_cluster() {
    local NUM_CHIPS="${1:-2}"

    echo -e "\n${YELLOW}Setting up $NUM_CHIPS-chip cluster...${NC}"

    # Create cluster config
    cat > cluster.toml << EOF
# RuvLLM ESP32 Cluster Configuration
# Generated by install.sh

[cluster]
name = "ruvllm-cluster"
chips = $NUM_CHIPS
topology = "pipeline"  # pipeline, tensor, hybrid

[chips]
EOF

    for i in $(seq 1 $NUM_CHIPS); do
        if [ "$OS" = "macos" ]; then
            DEFAULT_PORT="/dev/cu.usbserial-$i"
        else
            DEFAULT_PORT="/dev/ttyUSB$((i-1))"
        fi

        cat >> cluster.toml << EOF
[[chips.nodes]]
id = $i
role = "$([ $i -eq 1 ] && echo 'master' || echo 'worker')"
port = "$DEFAULT_PORT"
layers = [$(( (i-1) * 2 / NUM_CHIPS )), $(( i * 2 / NUM_CHIPS - 1 ))]

EOF
    done

    echo -e "${GREEN}✓ Created cluster.toml${NC}"
    echo ""
    echo "Edit cluster.toml to set correct ports, then run:"
    echo "  ./cluster-flash.sh"
}

# Main
main() {
    detect_os

    case "${1:-install}" in
        install)
            check_deps
            install_esp_toolchain
            build_project
            echo -e "\n${GREEN}Installation complete!${NC}"
            echo "To flash: ./install.sh flash [port]"
            ;;
        build)
            build_project
            ;;
        flash)
            flash_device "$2"
            ;;
        deps)
            check_deps
            install_esp_toolchain
            ;;
        cluster)
            setup_cluster "$2"
            ;;
        help|--help|-h)
            usage
            ;;
        *)
            echo -e "${RED}Unknown command: $1${NC}"
            usage
            exit 1
            ;;
    esac
}

main "$@"