#!/bin/bash

# Enhanced RDP Script with Microsoft RDP Client-like features
# Author: Enhanced version
# Date: $(date +%Y-%m-%d)

# Configuration directories
CONFIG_DIR="$HOME/.config/rdp-client"
CREDENTIALS_DIR="$CONFIG_DIR/credentials"
CONNECTIONS_DIR="$CONFIG_DIR/connections"

# Create necessary directories
mkdir -p "$CONFIG_DIR" "$CREDENTIALS_DIR" "$CONNECTIONS_DIR"

# Function to encrypt credentials
encrypt_password() {
    local password="$1"
    echo "$password" | openssl enc -aes-256-cbc -a -salt -pass pass:"$(whoami)@$(hostname)"
}

# Function to decrypt credentials
decrypt_password() {
    local encrypted="$1"
    echo "$encrypted" | openssl enc -aes-256-cbc -d -a -pass pass:"$(whoami)@$(hostname)" 2>/dev/null
}

# Function to save credentials
save_credentials() {
    local server="$1"
    local username="$2"
    local password="$3"
    local domain="$4"
    
    local cred_file="$CREDENTIALS_DIR/${server}.cred"
    
    cat > "$cred_file" << EOL
username=$username
domain=$domain
password=$(encrypt_password "$password")
EOL
    chmod 600 "$cred_file"
}

# Function to load credentials
load_credentials() {
    local server="$1"
    local cred_file="$CREDENTIALS_DIR/${server}.cred"
    
    if [[ -f "$cred_file" ]]; then
        source "$cred_file"
        password=$(decrypt_password "$password")
        echo "$username|$domain|$password"
    else
        echo "||"
    fi
}

# Function to get saved connections
get_saved_connections() {
    local connections=()
    if [[ -d "$CONNECTIONS_DIR" ]]; then
        for file in "$CONNECTIONS_DIR"/*.conn; do
            if [[ -f "$file" ]]; then
                local basename=$(basename "$file" .conn)
                connections+=("$basename")
            fi
        done
    fi
    printf '%s\n' "${connections[@]}"
}

# Function to save connection profile
save_connection_profile() {
    local name="$1"
    local server="$2"
    local username="$3"
    local domain="$4"
    local resolution="$5"
    local multimon="$6"
    local sound="$7"
    local clipboard="$8"
    local drives="$9"
    
    local conn_file="$CONNECTIONS_DIR/${name}.conn"
    
    cat > "$conn_file" << EOL
server=$server
username=$username
domain=$domain
resolution=$resolution
multimon=$multimon
sound=$sound
clipboard=$clipboard
drives=$drives
created=$(date)
EOL
}

# Function to load connection profile
load_connection_profile() {
    local name="$1"
    local conn_file="$CONNECTIONS_DIR/${name}.conn"
    
    if [[ -f "$conn_file" ]]; then
        source "$conn_file"
        echo "$server|$username|$domain|$resolution|$multimon|$sound|$clipboard|$drives"
    fi
}

# Function to show main menu
show_main_menu() {
    local saved_connections=$(get_saved_connections)
    
    if [[ -n "$saved_connections" ]]; then
        zenity --list \
            --title="RDP Client - Microsoft Style" \
            --text="Choose an option:" \
            --column="Action" \
            --width=400 \
            --height=300 \
            "New Connection" \
            "Saved Connections" \
            "Manage Connections"
    else
        zenity --list \
            --title="RDP Client - Microsoft Style" \
            --text="Choose an option:" \
            --column="Action" \
            --width=400 \
            --height=300 \
            "New Connection"
    fi
}

# Function to show saved connections
show_saved_connections() {
    local connections=($(get_saved_connections))
    
    if [[ ${#connections[@]} -eq 0 ]]; then
        zenity --info --text="No saved connections found."
        return 1
    fi
    
    zenity --list \
        --title="Saved Connections" \
        --text="Select a saved connection:" \
        --column="Connection Name" \
        --width=400 \
        --height=300 \
        "${connections[@]}"
}

# Function to create new connection dialog
create_new_connection() {
    local form_data
    form_data=$(zenity --forms \
        --title="New RDP Connection" \
        --text="Enter connection details:" \
        --add-entry="Server/IP:" \
        --add-entry="Username:" \
        --add-entry="Domain (optional):" \
        --add-password="Password:" \
        --add-combo="Resolution:" --combo-values="1920x1080|1366x768|1280x1024|1024x768|Full Screen" \
        --add-combo="Multiple Monitors:" --combo-values="No|Yes" \
        --add-combo="Sound:" --combo-values="Yes|No" \
        --add-combo="Clipboard:" --combo-values="Yes|No" \
        --add-combo="Share Home Drive:" --combo-values="Yes|No" \
        --width=500)
    
    if [[ $? -ne 0 ]]; then
        return 1
    fi
    
    IFS='|' read -r server username domain password resolution multimon sound clipboard drives <<< "$form_data"
    
    # Validate required fields
    if [[ -z "$server" || -z "$username" ]]; then
        zenity --error --text="Server and Username are required!"
        return 1
    fi
    
    # Ask if user wants to save credentials
    if zenity --question --text="Save credentials for future use?" --width=300; then
        save_credentials "$server" "$username" "$password" "$domain"
        
        # Ask if user wants to save connection profile
        if zenity --question --text="Save this connection profile?" --width=300; then
            local conn_name
            conn_name=$(zenity --entry --text="Enter a name for this connection:" --entry-text="$server")
            if [[ -n "$conn_name" ]]; then
                save_connection_profile "$conn_name" "$server" "$username" "$domain" "$resolution" "$multimon" "$sound" "$clipboard" "$drives"
            fi
        fi
    fi
    
    # Connect
    connect_rdp "$server" "$username" "$password" "$domain" "$resolution" "$multimon" "$sound" "$clipboard" "$drives"
}

# Function to connect using saved connection
connect_saved() {
    local conn_name="$1"
    local conn_data=$(load_connection_profile "$conn_name")
    
    if [[ -z "$conn_data" ]]; then
        zenity --error --text="Connection profile not found!"
        return 1
    fi
    
    IFS='|' read -r server username domain resolution multimon sound clipboard drives <<< "$conn_data"
    
    # Load saved credentials
    local cred_data=$(load_credentials "$server")
    IFS='|' read -r saved_username saved_domain saved_password <<< "$cred_data"
    
    local password=""
    if [[ -n "$saved_password" ]]; then
        password="$saved_password"
        if [[ -n "$saved_username" ]]; then
            username="$saved_username"
        fi
        if [[ -n "$saved_domain" ]]; then
            domain="$saved_domain"
        fi
    else
        # Ask for password if not saved
        password=$(zenity --password --text="Enter password for $username@$server:")
        if [[ $? -ne 0 ]]; then
            return 1
        fi
    fi
    
    connect_rdp "$server" "$username" "$password" "$domain" "$resolution" "$multimon" "$sound" "$clipboard" "$drives"
}

# Function to execute RDP connection
connect_rdp() {
    local server="$1"
    local username="$2"
    local password="$3"
    local domain="$4"
    local resolution="$5"
    local multimon="$6"
    local sound="$7"
    local clipboard="$8"
    local drives="$9"
    
    # Build xfreerdp command
    local cmd="/usr/bin/xfreerdp"
    
    # Basic options
    cmd="$cmd +window-drag +smart-sizing /cert-ignore"
    
    # Server and authentication
    cmd="$cmd /v:$server /u:$username"
    if [[ -n "$password" ]]; then
        cmd="$cmd /p:$password"
    fi
    if [[ -n "$domain" ]]; then
        cmd="$cmd /d:$domain"
    fi
    
    # Resolution settings
    case "$resolution" in
        "Full Screen")
            cmd="$cmd /f"
            ;;
        *)
            cmd="$cmd /size:$resolution"
            ;;
    esac
    
    # Multiple monitors
    if [[ "$multimon" == "Yes" ]]; then
        cmd="$cmd /multimon /monitors:0,1"
    else
        cmd="$cmd /monitors:0"
    fi
    
    # Sound
    if [[ "$sound" == "Yes" ]]; then
        cmd="$cmd /sound /microphone"
    fi
    
    # Clipboard
    if [[ "$clipboard" == "Yes" ]]; then
        cmd="$cmd /clipboard"
    fi
    
    # Drive sharing
    if [[ "$drives" == "Yes" ]]; then
        cmd="$cmd /drive:home,/home/rwiegand"
    fi
    
    # Execute connection
    echo "Connecting to $server..."
    eval "$cmd"
}

# Function to manage connections
manage_connections() {
    local action
    action=$(zenity --list \
        --title="Manage Connections" \
        --text="Choose an action:" \
        --column="Action" \
        --width=300 \
        --height=200 \
        "Delete Connection" \
        "Clear All Credentials" \
        "Back to Main Menu")
    
    case "$action" in
        "Delete Connection")
            local conn_to_delete=$(show_saved_connections)
            if [[ -n "$conn_to_delete" ]]; then
                if zenity --question --text="Delete connection '$conn_to_delete'?" --width=300; then
                    rm -f "$CONNECTIONS_DIR/${conn_to_delete}.conn"
                    rm -f "$CREDENTIALS_DIR/${conn_to_delete}.cred"
                    zenity --info --text="Connection deleted successfully."
                fi
            fi
            ;;
        "Clear All Credentials")
            if zenity --question --text="This will delete ALL saved credentials. Continue?" --width=300; then
                rm -f "$CREDENTIALS_DIR"/*.cred
                zenity --info --text="All credentials cleared."
            fi
            ;;
    esac
}

# Main program flow
main() {
    while true; do
        local choice=$(show_main_menu)
        
        # Debug: Show what was selected
        echo "Selected: '$choice'"
        
        case "$choice" in
            "New Connection")
                create_new_connection
                ;;
            "Saved Connections")
                local selected_conn=$(show_saved_connections)
                if [[ -n "$selected_conn" ]]; then
                    connect_saved "$selected_conn"
                fi
                ;;
            "Manage Connections")
                manage_connections
                ;;
            *)
                echo "Exiting..."
                break
                ;;
        esac
        
        # Ask if user wants to make another connection
        if ! zenity --question --text="Make another connection?" --width=300; then
            break
        fi
    done
}

# Check dependencies
if ! command -v xfreerdp &> /dev/null; then
    zenity --error --text="xfreerdp is not installed. Please install freerdp package."
    exit 1
fi

if ! command -v zenity &> /dev/null; then
    zenity --error --text="zenity is not installed. Please install zenity package."
    exit 1
fi

if ! command -v openssl &> /dev/null; then
    zenity --error --text="openssl is not installed. Please install openssl package."
    exit 1
fi

# Start the application
main