#!/bin/bash

# WebSocket SSH Tunnel - Bridge Mode Test Script
# This script demonstrates how to use wsssht in bridge mode
# Bridge mode allows programmatic control via JSON stdin/stdout

set -e

# Configuration
WSSSHT="./wssshtools/wsssht"  # Path to wsssht binary
CLIENT_ID="${CLIENT_ID:-testclient}"
WSSSHD_HOST="${WSSSHD_HOST:-mbetter.nexlab.net}"
WSSSHD_PORT="${WSSSHD_PORT:-9898}"
DEBUG="${DEBUG:-0}"

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

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

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

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

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

# Function to send JSON command to wsssht
send_command() {
    local command="$1"
    echo "$command"
    log_info "Sending command: $command"
}

# Function to test bridge mode
test_bridge_mode() {
    log_info "Starting bridge mode test..."
    log_info "Client ID: $CLIENT_ID"
    log_info "WSSSHD Host: $WSSSHD_HOST"
    log_info "WSSSHD Port: $WSSSHD_PORT"

    # Build command arguments
    local cmd_args="--bridge --clientid $CLIENT_ID --wssshd-host $WSSSHD_HOST --wssshd-port $WSSSHD_PORT"
    if [ "$DEBUG" = "1" ]; then
        cmd_args="$cmd_args --debug"
    fi

    log_info "Starting wsssht with command: $WSSSHT $cmd_args"

    # Start wsssht in bridge mode with coprocess
    coproc WSSSHT_PROC { $WSSSHT $cmd_args; }

    # Get the process IDs
    local wsssht_pid=$WSSSHT_PROC_PID
    local wsssht_in=${WSSSHT_PROC[1]}
    local wsssht_out=${WSSSHT_PROC[0]}

    log_success "wsssht started with PID: $wsssht_pid"

    # Function to read JSON response
    read_response() {
        local timeout=${1:-5}
        local response=""

        # Set up timeout
        if command -v timeout >/dev/null 2>&1; then
            response=$(timeout $timeout cat <&${wsssht_out} 2>/dev/null || true)
        else
            # Fallback for systems without timeout
            response=$(head -1 <&${wsssht_out} 2>/dev/null || true)
        fi

        if [ -n "$response" ]; then
            log_info "Received: $response"
            echo "$response"
        else
            log_warning "No response received within $timeout seconds"
        fi
    }

    # Wait for initial status message
    log_info "Waiting for initial status message..."
    sleep 2

    # Read initial messages
    while read -t 1 -u ${wsssht_out} line 2>/dev/null; do
        if [ -n "$line" ]; then
            log_info "Initial response: $line"
            # Check if bridge is ready
            if echo "$line" | grep -q '"type":"ready"'; then
                log_success "Bridge mode is ready!"
                break
            fi
        fi
    done

    # Generate a unique request ID for the tunnel
    REQUEST_ID="test_$(date +%s)_$$"

    # Test 1: Send tunnel request
    log_info "Test 1: Sending tunnel request"
    send_command '{"type":"tunnel_request","client_id":"'$CLIENT_ID'","request_id":"'$REQUEST_ID'","tunnel":"any","tunnel_control":"any","service":"ssh","timestamp":'$(date +%s)'}' >&${wsssht_in}

    # Wait for tunnel acknowledgment
    log_info "Waiting for tunnel acknowledgment..."
    local tunnel_ack_received=false
    local timeout_counter=0
    while [ $timeout_counter -lt 10 ]; do  # 10 second timeout
        if read -t 1 -u ${wsssht_out} line 2>/dev/null; then
            if [ -n "$line" ]; then
                log_info "Response: $line"
                # Check if we received tunnel_ack
                if echo "$line" | grep -q '"message".*"tunnel_ack"'; then
                    log_success "Tunnel acknowledgment received!"
                    tunnel_ack_received=true
                    break
                fi
            fi
        fi
        timeout_counter=$((timeout_counter + 1))
    done

    if [ "$tunnel_ack_received" = false ]; then
        log_error "Timeout waiting for tunnel acknowledgment"
        # Send quit to clean up
        send_command '{"command":"quit","timestamp":'$(date +%s)'}' >&${wsssht_in}
        return 1
    fi

    # Test 2: Send data message with text string
    log_info "Test 2: Sending data message with text string"
    TEST_MESSAGE="Hello from bridge mode test!"
    # Convert to hex for tunnel_data format
    HEX_MESSAGE=$(echo -n "$TEST_MESSAGE" | xxd -p | tr -d '\n')
    send_command '{"type":"tunnel_data","request_id":"'$REQUEST_ID'","data":"'$HEX_MESSAGE'","timestamp":'$(date +%s)'}' >&${wsssht_in}

    # Wait for tunnel response
    log_info "Waiting for tunnel response..."
    local tunnel_response_received=false
    timeout_counter=0
    while [ $timeout_counter -lt 10 ]; do  # 10 second timeout
        if read -t 1 -u ${wsssht_out} line 2>/dev/null; then
            if [ -n "$line" ]; then
                log_info "Response: $line"
                # Check if we received tunnel_response
                if echo "$line" | grep -q '"message".*"tunnel_response"'; then
                    log_success "Tunnel response received!"
                    tunnel_response_received=true
                    break
                fi
            fi
        fi
        timeout_counter=$((timeout_counter + 1))
    done

    if [ "$tunnel_response_received" = false ]; then
        log_warning "No tunnel response received within timeout"
    fi

    # Test 3: Send tunnel close
    log_info "Test 3: Sending tunnel close"
    send_command '{"type":"tunnel_close","request_id":"'$REQUEST_ID'","timestamp":'$(date +%s)'}' >&${wsssht_in}

    sleep 2

    # Test 4: Send quit command to end bridge mode
    log_info "Test 4: Sending quit command to end bridge mode"
    send_command '{"command":"quit","timestamp":'$(date +%s)'}' >&${wsssht_in}

    sleep 2

    # Read final messages
    while read -t 2 -u ${wsssht_out} line 2>/dev/null; do
        if [ -n "$line" ]; then
            log_info "Final response: $line"
        fi
    done

    # Wait for process to finish
    wait $wsssht_pid 2>/dev/null || true

    log_success "Bridge mode test completed!"
}

# Function to show usage
show_usage() {
    cat << EOF
WebSocket SSH Tunnel - Bridge Mode Test Script

Usage: $0 [OPTIONS]

Options:
    -c, --client-id ID     Client ID for the tunnel (default: testclient)
    -h, --host HOST        wssshd server hostname (default: mbetter.nexlab.net)
    -p, --port PORT        wssshd server port (default: 9898)
    -d, --debug            Enable debug output
    --help                 Show this help message

Environment Variables:
    CLIENT_ID              Client ID (overrides -c)
    WSSSHD_HOST            wssshd hostname (overrides -h)
    WSSSHD_PORT            wssshd port (overrides -p)
    DEBUG                  Enable debug (set to 1)

Examples:
    $0
    $0 --client-id myclient --host example.com --debug
    CLIENT_ID=myclient WSSSHD_HOST=example.com $0

EOF
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -c|--client-id)
            CLIENT_ID="$2"
            shift 2
            ;;
        -h|--host)
            WSSSHD_HOST="$2"
            shift 2
            ;;
        -p|--port)
            WSSSHD_PORT="$2"
            shift 2
            ;;
        -d|--debug)
            DEBUG=1
            shift
            ;;
        --help)
            show_usage
            exit 0
            ;;
        *)
            log_error "Unknown option: $1"
            show_usage
            exit 1
            ;;
    esac
done

# Check if wsssht binary exists
if [ ! -x "$WSSSHT" ]; then
    log_error "wsssht binary not found at: $WSSSHT"
    log_error "Please build wsssht first or update WSSSHT path in this script"
    exit 1
fi

log_info "WebSocket SSH Tunnel - Bridge Mode Test"
log_info "========================================"

# Run the test
test_bridge_mode

log_success "Test script completed successfully!"