#!/bin/bash
# Script to create an SSH tunnel to the remote server

# Load environment variables from local_test.env
if [ -f "local_test.env" ]; then
    source local_test.env
else
    echo "Error: local_test.env file not found. Please run setup_local_test_env.sh first."
    exit 1
fi

# Check if required variables are set
if [ -z "$SERVER_IP" ] || [ -z "$SERVER_PORT" ] || [ -z "$SERVER_USER" ]; then
    echo "Error: SERVER_IP, SERVER_PORT, or SERVER_USER not set in local_test.env."
    exit 1
fi

# Check if password is set
if [ -z "$SERVER_PASSWORD" ]; then
    echo "Warning: SERVER_PASSWORD not set in local_test.env. You will be prompted for the password."
fi

# Function to create SSH tunnel
create_tunnel() {
    local local_port=$1
    local remote_port=$2
    
    echo "Creating SSH tunnel from localhost:$local_port to $SERVER_IP:$remote_port..."
    
    # Check if the tunnel is already established
    if netstat -tuln | grep -q ":$local_port "; then
        echo "Port $local_port is already in use. Tunnel may already be established."
        return 0
    fi
    
    # Create the SSH tunnel
    if [ -n "$SERVER_PASSWORD" ]; then
        # Use sshpass if available
        if command -v sshpass &> /dev/null; then
            sshpass -p "$SERVER_PASSWORD" ssh -N -L $local_port:localhost:$remote_port -p $SERVER_PORT $SERVER_USER@$SERVER_IP &
        else
            echo "sshpass not installed. You will be prompted for the password."
            ssh -N -L $local_port:localhost:$remote_port -p $SERVER_PORT $SERVER_USER@$SERVER_IP &
        fi
    else
        ssh -N -L $local_port:localhost:$remote_port -p $SERVER_PORT $SERVER_USER@$SERVER_IP &
    fi
    
    # Check if the tunnel was established
    if [ $? -eq 0 ]; then
        echo "SSH tunnel established. Local port $local_port is now forwarded to $SERVER_IP:$remote_port."
        echo "Tunnel process ID: $!"
        echo $! >> tunnel_pids.txt
    else
        echo "Failed to establish SSH tunnel."
        return 1
    fi
}

# Create tunnels for Ollama and OpenWebUI
echo "Setting up SSH tunnels to remote server at $SERVER_IP..."

# Create a file to store tunnel process IDs
touch tunnel_pids.txt

# Create tunnel for Ollama (port 11434)
create_tunnel 11434 11434

# Create tunnel for OpenWebUI (port 8080)
create_tunnel 8080 8080

echo "SSH tunnels established. You can now access:"
echo "- Ollama API at http://localhost:11434"
echo "- OpenWebUI at http://localhost:8080"
echo ""
echo "To stop the tunnels, run: ./stop_ssh_tunnels.sh"