#!/usr/bin/env bash
#
# RustyTorch++ Medical Demos - Unified Launch Script
#
# Usage:
#   ./run-demo.sh           # Default: dev mode (landing page)
#   ./run-demo.sh --dev     # Development mode (hot reload)
#   ./run-demo.sh --prod    # Production build and run
#   ./run-demo.sh --test    # Run tests only
#   ./run-demo.sh --check   # Prerequisites check only
#   ./run-demo.sh --cleanup # Kill stuck processes
#   ./run-demo.sh --validate-mre  # Validate MRE backend (compile + tests)
#   ./run-demo.sh --help    # Show usage
#
# Exit codes:
#   0 - Success
#   1 - Prerequisites check failed
#   2 - Tests failed
#   3 - Build failed
#   4 - Launch failed
#   5 - Validation failed

set -euo pipefail

# Script directory (demos/)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
UI_DIR="$SCRIPT_DIR/ui"

# Process tracking
TAURI_PID=""
VITE_PID=""

# Ports used by the demo
VITE_PORT=9090
TAURI_PORT=1420

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

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[OK]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_header() {
    echo ""
    echo -e "${BOLD}${CYAN}=== $1 ===${NC}"
    echo ""
}

# Show usage
show_usage() {
    cat << EOF
${BOLD}RustyTorch++ Medical Demos${NC}

${BOLD}USAGE:${NC}
    ./run-demo.sh [OPTIONS]

${BOLD}OPTIONS:${NC}
    --dev           Development mode with hot reload (default)
    --prod          Build and run production binary
    --test          Run tests only (no UI launch)
    --check         Check prerequisites only
    --cleanup       Kill any stuck demo processes
    --validate-mre  Validate MRE backend (compile + run tests)
    --help          Show this help message

${BOLD}DEMOS AVAILABLE:${NC}
    Virtual Catheter (Hemodynamics) - Navier-Stokes PINN for blood flow
    MRE Elastography - Inverse Helmholtz for tissue stiffness

${BOLD}EXAMPLES:${NC}
    ./run-demo.sh               # Start in dev mode (landing page)
    ./run-demo.sh --test --dev  # Run tests, then start dev mode
    ./run-demo.sh --prod        # Build and run production version
    ./run-demo.sh --validate-mre # Validate MRE backend compiles and tests pass
    ./run-demo.sh --cleanup     # Clean up stuck processes

${BOLD}EXIT CODES:${NC}
    0 - Success
    1 - Prerequisites check failed
    2 - Tests failed
    3 - Build failed
    4 - Launch failed
    5 - Validation failed

EOF
}

# Check if a command exists
command_exists() {
    command -v "$1" &> /dev/null
}

# Check if a port is in use
port_in_use() {
    local port=$1
    if command_exists lsof; then
        lsof -i :"$port" &> /dev/null
    elif command_exists ss; then
        ss -tuln | grep -q ":$port "
    elif command_exists netstat; then
        netstat -tuln | grep -q ":$port "
    else
        # Fallback: try to connect
        (echo > /dev/tcp/localhost/"$port") 2>/dev/null
    fi
}

# Get PID using a port
get_pid_on_port() {
    local port=$1
    if command_exists lsof; then
        lsof -ti :"$port" 2>/dev/null || true
    elif command_exists fuser; then
        fuser "$port"/tcp 2>/dev/null | awk '{print $1}' || true
    else
        echo ""
    fi
}

# Check prerequisites
check_prerequisites() {
    log_header "Checking Prerequisites"
    local failed=0

    # Check Node.js
    if command_exists node; then
        local node_version
        node_version=$(node --version | sed 's/v//')
        local major_version
        major_version=$(echo "$node_version" | cut -d. -f1)
        if [ "$major_version" -ge 18 ]; then
            log_success "Node.js v$node_version (>= 18 required)"
        else
            log_error "Node.js v$node_version is too old (>= 18 required)"
            failed=1
        fi
    else
        log_error "Node.js not found"
        failed=1
    fi

    # Check npm
    if command_exists npm; then
        local npm_version
        npm_version=$(npm --version)
        log_success "npm v$npm_version"
    else
        log_error "npm not found"
        failed=1
    fi

    # Check Rust/cargo
    if command_exists cargo; then
        local cargo_version
        cargo_version=$(cargo --version | awk '{print $2}')
        log_success "cargo v$cargo_version"
    else
        log_error "cargo not found (Rust toolchain required)"
        failed=1
    fi

    # Check rustc
    if command_exists rustc; then
        local rustc_version
        rustc_version=$(rustc --version | awk '{print $2}')
        log_success "rustc v$rustc_version"
    else
        log_error "rustc not found"
        failed=1
    fi

    # Check UI directory exists
    if [ -d "$UI_DIR" ]; then
        log_success "UI directory found: $UI_DIR"
    else
        log_error "UI directory not found: $UI_DIR"
        failed=1
    fi

    # Check/install node_modules
    if [ -d "$UI_DIR/node_modules" ]; then
        log_success "node_modules directory exists"
    else
        log_warn "node_modules not found, running npm install..."
        if ! (cd "$UI_DIR" && npm install); then
            log_error "npm install failed"
            failed=1
        else
            log_success "npm install completed"
        fi
    fi

    # Check Tauri CLI
    if (cd "$UI_DIR" && npm run tauri -- --version &> /dev/null); then
        local tauri_version
        tauri_version=$(cd "$UI_DIR" && npm run tauri -- --version 2>/dev/null | tail -1)
        log_success "Tauri CLI $tauri_version"
    else
        log_warn "Tauri CLI not responding, may need to build first"
    fi

    # Check for port conflicts
    if port_in_use $VITE_PORT; then
        log_warn "Port $VITE_PORT is already in use (vite dev server port)"
    fi

    if port_in_use $TAURI_PORT; then
        log_warn "Port $TAURI_PORT is already in use (Tauri port)"
    fi

    if [ $failed -eq 1 ]; then
        log_error "Prerequisites check failed"
        return 1
    fi

    log_success "All prerequisites satisfied"
    return 0
}

# Run tests
run_tests() {
    log_header "Running Tests"

    cd "$UI_DIR"

    if npm test; then
        log_success "All tests passed"
        return 0
    else
        log_error "Tests failed"
        return 2
    fi
}

# Cleanup function
cleanup() {
    log_header "Cleaning Up"

    local killed=0

    # Kill tracked processes
    if [ -n "$TAURI_PID" ] && kill -0 "$TAURI_PID" 2>/dev/null; then
        log_info "Stopping Tauri process (PID: $TAURI_PID)..."
        kill -TERM "$TAURI_PID" 2>/dev/null || true
        sleep 1
        kill -KILL "$TAURI_PID" 2>/dev/null || true
        killed=$((killed + 1))
    fi

    if [ -n "$VITE_PID" ] && kill -0 "$VITE_PID" 2>/dev/null; then
        log_info "Stopping Vite process (PID: $VITE_PID)..."
        kill -TERM "$VITE_PID" 2>/dev/null || true
        sleep 1
        kill -KILL "$VITE_PID" 2>/dev/null || true
        killed=$((killed + 1))
    fi

    # Kill processes on ports
    for port in $VITE_PORT $TAURI_PORT; do
        local pid
        pid=$(get_pid_on_port "$port")
        if [ -n "$pid" ]; then
            log_info "Killing process on port $port (PID: $pid)..."
            kill -TERM "$pid" 2>/dev/null || true
            sleep 0.5
            kill -KILL "$pid" 2>/dev/null || true
            killed=$((killed + 1))
        fi
    done

    # Kill any remaining tauri-related processes
    local tauri_pids
    tauri_pids=$(pgrep -f "rtx-hemodynamics-ui" 2>/dev/null || true)
    if [ -n "$tauri_pids" ]; then
        for pid in $tauri_pids; do
            log_info "Killing remaining Tauri process (PID: $pid)..."
            kill -TERM "$pid" 2>/dev/null || true
            killed=$((killed + 1))
        done
    fi

    # Kill any vite processes in the UI directory
    local vite_pids
    vite_pids=$(pgrep -f "vite.*$UI_DIR" 2>/dev/null || true)
    if [ -n "$vite_pids" ]; then
        for pid in $vite_pids; do
            log_info "Killing remaining Vite process (PID: $pid)..."
            kill -TERM "$pid" 2>/dev/null || true
            killed=$((killed + 1))
        done
    fi

    if [ $killed -gt 0 ]; then
        log_success "Cleaned up $killed process(es)"
    else
        log_info "No processes to clean up"
    fi

    return 0
}

# Validate that services are running
validate_running() {
    log_header "Validating Services"

    local dev_mode=${1:-false}
    local failed=0

    # Check vite server (dev mode only)
    if [ "$dev_mode" = true ]; then
        if port_in_use $VITE_PORT; then
            log_success "Vite dev server running on port $VITE_PORT"
        else
            log_error "Vite dev server not responding on port $VITE_PORT"
            failed=1
        fi
    fi

    # Check for Tauri/app process
    if pgrep -f "rtx-hemodynamics-ui" > /dev/null 2>&1; then
        log_success "Tauri application process running"
    else
        log_warn "Tauri application process not detected (may still be starting)"
    fi

    if [ $failed -eq 1 ]; then
        return 5
    fi

    return 0
}

# Wait for port to be available
wait_for_port() {
    local port=$1
    local timeout=${2:-30}
    local elapsed=0

    log_info "Waiting for port $port to be ready..."

    while ! port_in_use "$port"; do
        sleep 1
        elapsed=$((elapsed + 1))
        if [ $elapsed -ge $timeout ]; then
            log_error "Timeout waiting for port $port after ${timeout}s"
            return 1
        fi
    done

    log_success "Port $port is ready (waited ${elapsed}s)"
    return 0
}

# Launch development mode
launch_dev() {
    log_header "Launching Development Mode"

    cd "$UI_DIR"

    # Check for port conflicts first
    if port_in_use $VITE_PORT; then
        log_error "Port $VITE_PORT is already in use"
        log_info "Run './run-demo.sh --cleanup' to kill stuck processes"
        return 4
    fi

    log_info "Starting Tauri development server..."
    log_info "This may take a moment on first run (compiling Rust backend)..."
    echo ""

    # Start tauri dev (this starts both vite and the Tauri app)
    npm run tauri:dev &
    TAURI_PID=$!

    log_info "Tauri dev started (PID: $TAURI_PID)"

    # Wait for vite server to be ready
    if ! wait_for_port $VITE_PORT 60; then
        log_error "Failed to start Vite dev server"
        cleanup
        return 4
    fi

    # Validate services
    sleep 2
    if ! validate_running true; then
        log_warn "Validation incomplete, but services may still be starting"
    fi

    echo ""
    log_success "Demo is running!"
    echo ""
    echo -e "${BOLD}Frontend:${NC} http://localhost:$VITE_PORT"
    echo -e "${BOLD}Process:${NC}  PID $TAURI_PID"
    echo ""
    echo -e "${YELLOW}Press Ctrl+C to stop the demo${NC}"
    echo ""

    # Wait for the Tauri process
    wait $TAURI_PID 2>/dev/null || true

    return 0
}

# Launch production mode
launch_prod() {
    log_header "Building Production Version"

    cd "$UI_DIR"

    log_info "Building production bundle..."
    log_info "This may take several minutes on first build..."
    echo ""

    if ! npm run tauri:build; then
        log_error "Production build failed"
        return 3
    fi

    log_success "Build completed"

    log_header "Launching Production Binary"

    # Find the built binary
    local binary_path=""
    local possible_paths=(
        "$UI_DIR/src-tauri/target/release/rtx-hemodynamics-ui"
        "$UI_DIR/src-tauri/target/release/bundle/appimage/"*.AppImage
        "$UI_DIR/src-tauri/target/release/bundle/deb/"*.deb
    )

    for path in "${possible_paths[@]}"; do
        if [ -f "$path" ] || compgen -G "$path" > /dev/null; then
            binary_path=$(compgen -G "$path" | head -1)
            break
        fi
    done

    if [ -z "$binary_path" ]; then
        log_error "Could not find built binary"
        log_info "Check: $UI_DIR/src-tauri/target/release/"
        return 4
    fi

    log_info "Found binary: $binary_path"
    log_info "Launching..."
    echo ""

    # Make executable if needed
    chmod +x "$binary_path" 2>/dev/null || true

    # Launch the binary
    "$binary_path" &
    TAURI_PID=$!

    log_success "Production app launched (PID: $TAURI_PID)"
    echo ""
    echo -e "${YELLOW}Press Ctrl+C to stop the demo${NC}"
    echo ""

    # Wait for the process
    wait $TAURI_PID 2>/dev/null || true

    return 0
}

# Signal handler
handle_signal() {
    echo ""
    log_info "Received shutdown signal..."
    cleanup
    exit 0
}

# Validate MRE backend
validate_mre() {
    log_header "Validating MRE Backend"

    local failed=0

    # Check mre-shared compiles
    log_info "Checking mre-shared crate..."
    if cargo check -p mre-shared --quiet 2>/dev/null; then
        log_success "mre-shared compiles"
    else
        log_error "mre-shared compilation failed"
        failed=1
    fi

    # Check rtx-mre compiles
    log_info "Checking rtx-mre crate..."
    if cargo check -p rtx-mre --quiet 2>/dev/null; then
        log_success "rtx-mre compiles"
    else
        log_error "rtx-mre compilation failed"
        failed=1
    fi

    # Check rtx-hemodynamics-server compiles (includes MreService)
    log_info "Checking server crate with MreService..."
    if cargo check -p rtx-hemodynamics-server --quiet 2>/dev/null; then
        log_success "rtx-hemodynamics-server compiles"
    else
        log_error "rtx-hemodynamics-server compilation failed"
        failed=1
    fi

    # Run unit tests
    log_info "Running MRE unit tests..."
    if cargo test -p mre-shared -p rtx-mre --quiet 2>/dev/null; then
        log_success "All MRE tests passed"
    else
        log_error "MRE tests failed"
        failed=1
    fi

    if [ $failed -eq 1 ]; then
        log_error "MRE validation failed"
        return 5
    fi

    log_success "MRE backend validation complete"
    return 0
}

# Main function
main() {
    local mode="dev"
    local run_test=false
    local check_only=false
    local cleanup_only=false
    local validate_mre_only=false

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --dev)
                mode="dev"
                shift
                ;;
            --prod)
                mode="prod"
                shift
                ;;
            --test)
                run_test=true
                shift
                ;;
            --check)
                check_only=true
                shift
                ;;
            --cleanup)
                cleanup_only=true
                shift
                ;;
            --validate-mre)
                validate_mre_only=true
                shift
                ;;
            --help|-h)
                show_usage
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_usage
                exit 1
                ;;
        esac
    done

    # Setup signal handlers
    trap handle_signal SIGINT SIGTERM

    echo ""
    echo -e "${BOLD}${CYAN}RustyTorch++ Medical Demos${NC}"
    echo ""

    # Cleanup only mode
    if [ "$cleanup_only" = true ]; then
        cleanup
        exit 0
    fi

    # Validate MRE only mode
    if [ "$validate_mre_only" = true ]; then
        validate_mre
        exit $?
    fi

    # Prerequisites check
    if ! check_prerequisites; then
        exit 1
    fi

    # Check only mode
    if [ "$check_only" = true ]; then
        exit 0
    fi

    # Run tests if requested
    if [ "$run_test" = true ]; then
        if ! run_tests; then
            exit 2
        fi

        # If only --test was specified, exit here
        if [ "$mode" = "dev" ] && [ $# -eq 0 ]; then
            # Check if --dev was explicitly passed
            if [[ ! " ${BASH_ARGV[*]} " =~ " --dev " ]]; then
                exit 0
            fi
        fi
    fi

    # Launch based on mode
    case $mode in
        dev)
            launch_dev
            ;;
        prod)
            launch_prod
            ;;
    esac

    # Cleanup on exit
    cleanup
}

# Run main
main "$@"