#!/bin/bash

monitor_resolution="2960x1848@120"  # Default resolution for the virtual display
monitor_scale="1.6"  # Default scaling for the virtual display

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

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

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $*"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $*"
}

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

log_devices() {
    echo -e "${CYAN}[DEVICES]${NC} $*"
}

# Function to show ADB devices with live updates
show_adb_devices_live() {
    local prompt_message="$1"
    local exit_key="$2"
    local monitor_port_forwarding="$3"  # Optional: monitor and restore port forwarding
    local lines_to_clear=0
    local last_device_status=""
    local port_forwarding_active=true
    
    echo -e "\n${YELLOW}Press '$exit_key' when ready, or Ctrl+C to cancel${NC}"
    echo -e "${YELLOW}$prompt_message${NC}\n"
    
    while true; do
        # If not first iteration, clear the previous output
        if [[ $lines_to_clear -gt 0 ]]; then
            # Move cursor up and clear lines
            for ((i=0; i<lines_to_clear; i++)); do
                tput cuu1  # Move up one line
                tput el    # Clear line
            done
        fi
        
        # Count lines we're about to print (so we can clear them next time)
        local output_lines=0
        
        # Show current devices
        log_devices "Connected devices:"
        ((output_lines++))
        echo "----------------------------------------"
        ((output_lines++))
        
        local devices_output=$(adb devices 2>/dev/null)
        local device_found=false
        local current_device_status=""
        
        if [[ -n "$devices_output" ]]; then
            while IFS= read -r line; do
                if [[ -n "$line" && "$line" != *"List of devices attached"* ]]; then
                    device_found=true
                    current_device_status="$line"
                    if [[ "$line" == *"device"* ]]; then
                        echo -e "${GREEN}  ✓ $line${NC}"
                        
                        # Check if port forwarding needs to be restored
                        if [[ "$monitor_port_forwarding" == "true" ]]; then
                            # Check if port forwarding is actually active by testing the connection
                            if ! adb reverse --list 2>/dev/null | grep -q "tcp:5900"; then
                                port_forwarding_active=false
                            fi
                            
                            # If port forwarding is inactive but device is connected, restore it
                            if [[ $port_forwarding_active == false ]]; then
                                echo -e "${CYAN}  → Restoring port forwarding...${NC}"
                                ((output_lines++))
                                
                                # Restore port forwarding
                                if adb reverse tcp:5900 tcp:5900 2>/dev/null; then
                                    echo -e "${GREEN}  ✓ Port forwarding restored successfully${NC}"
                                    port_forwarding_active=true
                                else
                                    echo -e "${RED}  ✗ Failed to restore port forwarding${NC}"
                                    port_forwarding_active=false
                                fi
                                ((output_lines++))
                            fi
                        fi
                    elif [[ "$line" == *"unauthorized"* ]]; then
                        echo -e "${RED}  ✗ $line (Please authorize on device)${NC}"
                        port_forwarding_active=false
                    elif [[ "$line" == *"offline"* ]]; then
                        echo -e "${YELLOW}  ? $line${NC}"
                        port_forwarding_active=false
                    else
                        echo -e "${YELLOW}  ? $line${NC}"
                        port_forwarding_active=false
                    fi
                    ((output_lines++))
                fi
            done <<< "$devices_output"
        fi
        
        if [[ $device_found == false ]]; then
            echo -e "${RED}  No devices connected${NC}"
            ((output_lines++))
            port_forwarding_active=false
        fi
        
        echo "----------------------------------------"
        ((output_lines++))
        
        # Show port forwarding status if monitoring
        if [[ "$monitor_port_forwarding" == "true" ]]; then
            if [[ $port_forwarding_active == true && $device_found == true ]]; then
                echo -e "${GREEN}Port forwarding: Active${NC}"
            else
                echo -e "${RED}Port forwarding: Inactive${NC}"
            fi
            ((output_lines++))
        fi
        
        echo -e "${CYAN}Refreshing every 2 seconds...${NC}"
        ((output_lines++))
        echo -e "${YELLOW}Press '$exit_key' when ready${NC}"
        ((output_lines++))
        
        # Remember how many lines to clear next time and device status
        lines_to_clear=$output_lines
        last_device_status="$current_device_status"
        
        # Check if user wants to continue (non-blocking)
        read -t 2 -n 1 user_input
        if [[ "$user_input" == "$exit_key" ]] || [[ "$user_input" == "${exit_key^^}" ]]; then
            echo # New line after the key press
            break
        fi
    done
}

# Function to cleanup on exit
cleanup() {
    log_warning "Cleaning up..."
    
    # Kill wayvnc if it's running
    if [[ -n "${wayvnc_pid:-}" ]] && kill -0 "$wayvnc_pid" 2>/dev/null; then
        log_info "Stopping VNC server..."
        kill "$wayvnc_pid"
        wait "$wayvnc_pid" 2>/dev/null
    fi
    
    # Remove virtual display
    if hyprctl monitors | grep -q "Virtual-1"; then
        log_info "Removing virtual display..."
        hyprctl output remove Virtual-1
    fi
    
    # Remove port forward
    log_info "Removing port forwarding..."
    adb reverse --remove tcp:5900 2>/dev/null
    
    # Kill adb server
    log_info "Stopping ADB server..."
    adb kill-server
    
    log_success "Cleanup completed"
}

# Set up signal handlers
trap cleanup EXIT INT TERM

### STARTUP ###
clear
echo -e "${BLUE}╔══════════════════════════════════════╗${NC}"
echo -e "${BLUE}║        VNC Tablet Setup Script       ║${NC}"
echo -e "${BLUE}╚══════════════════════════════════════╝${NC}\n"

log_info "Creating headless monitor for Hyprland..."

# Create a new monitor configuration for Hyprland
if ! hyprctl output create headless Virtual-1; then
    log_error "Failed to create headless output"
    exit 1
fi

if ! hyprctl keyword monitor Virtual-1,${monitor_resolution},auto-left,${monitor_scale}; then
    log_error "Failed to configure monitor"
    exit 1
fi

log_success "Virtual display created successfully"

# Start adb server
log_info "Starting ADB server..."
if ! adb start-server; then
    log_error "Failed to start ADB server"
    exit 1
fi

log_success "ADB server started"

# Show devices with live updates
log_warning "Please connect your tablet and enable USB debugging"
show_adb_devices_live "Connect your tablet via USB and authorize debugging if prompted" "y"

# Final device check
log_info "Checking connected devices..."
device_count=$(adb devices | tail -n +2 | grep -c "device$")
if [[ $device_count -eq 0 ]]; then
    log_error "No devices connected or authorized. Please check your connection."
    exit 1
fi

log_success "$device_count device(s) connected and ready"

# Start VNC server
log_info "Starting VNC server on 127.0.0.1:5900..."
wayvnc -o Virtual-1 127.0.0.1 5900 &
wayvnc_pid=$!

# Give VNC server time to start
sleep 2

# Verify VNC server is running
if ! kill -0 "$wayvnc_pid" 2>/dev/null; then
    log_error "VNC server failed to start properly"
    exit 1
fi

log_success "VNC server started successfully (PID: $wayvnc_pid)"

# Set up port forwarding
log_info "Setting up port forwarding..."
if ! adb reverse tcp:5900 tcp:5900; then
    log_error "Failed to set up port forwarding"
    exit 1
fi

log_success "Port forwarding established (device:5900 -> host:5900)"

# Show final status
echo -e "\n${GREEN}╔══════════════════════════════════════╗${NC}"
echo -e "${GREEN}║          Setup Complete!             ║${NC}"
echo -e "${GREEN}╚══════════════════════════════════════╝${NC}"
echo -e "${GREEN}VNC server is running on 127.0.0.1:5900${NC}"
echo -e "${GREEN}Your tablet can now connect to localhost:5900${NC}\n"

# Wait for user to finish with live device monitoring
log_info "VNC server is running. Connect from your tablet to localhost:5900"
show_adb_devices_live "VNC server is active. When you're done, press 'q' to quit and cleanup" "q" "true"

log_info "Shutting down VNC setup..."