#!/usr/bin/env bash

#########################################################################################
#### xfreerdp-connect.sh — RDP connection launcher with Zenity GUI prompt            ####
#### Connects to Windows hosts via xfreerdp with optional RemoteApp support.         ####
#### Requires: bash, xfreerdp, zenity                                                ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.2                                                                     ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./xfreerdp-connect.sh                                                         ####
####   ./xfreerdp-connect.sh --help                                                  ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Colors ────────────────────────────────────────────────────────────
if [[ -t 1 ]]; then
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    RESET='\033[0m'
else
    RED="" GREEN="" RESET=""
fi

log() { echo -e "${GREEN}[OK]${RESET}    $*"; }
err() { echo -e "${RED}[ERROR]${RESET} $*" >&2; }

# ── Profile support ──────────────────────────────────────────────────
PROFILE_DIR="$HOME/.config/xfreerdp-connect"
PROFILE_FILE="$PROFILE_DIR/profiles.conf"

profile_list() {
    if [[ ! -f "$PROFILE_FILE" ]]; then
        echo "No saved profiles."
        return
    fi
    local found=0
    while IFS= read -r line; do
        if [[ "$line" =~ ^\[(.+)\]$ ]]; then
            local name="${BASH_REMATCH[1]}"
            local srv=""
            while IFS= read -r kv; do
                [[ -z "$kv" || "$kv" =~ ^\[ ]] && break
                if [[ "$kv" =~ ^server= ]]; then
                    srv="${kv#server=}"
                fi
            done
            printf "  %-20s %s\n" "$name" "$srv"
            found=1
        fi
    done < "$PROFILE_FILE"
    if [[ "$found" -eq 0 ]]; then
        echo "No saved profiles."
    fi
}

profile_load() {
    local name="$1"
    if [[ ! -f "$PROFILE_FILE" ]]; then
        err "Profile '$name' not found (no profiles file)"
        exit 1
    fi
    local in_section=0 found=0
    while IFS= read -r line; do
        if [[ "$line" =~ ^\[(.+)\]$ ]]; then
            if [[ "${BASH_REMATCH[1]}" == "$name" ]]; then
                in_section=1
                found=1
                continue
            else
                [[ "$in_section" -eq 1 ]] && break
            fi
        fi
        if [[ "$in_section" -eq 1 && -n "$line" ]]; then
            case "$line" in
                server=*)     server="${line#server=}" ;;
                port=*)       port="${line#port=}" ;;
                domain=*)     domain="${line#domain=}" ;;
                username=*)   username="${line#username=}" ;;
                app=*)        app="${line#app=}" ;;
                fullscreen=*) fullscreen="${line#fullscreen=}" ;;
                multimon=*)   multimon="${line#multimon=}" ;;
                drive=*)      drive="${line#drive=}" ;;
                sound=*)      sound="${line#sound=}" ;;
                gateway=*)    gateway="${line#gateway=}" ;;
                admin=*)      admin="${line#admin=}" ;;
                no_nla=*)     no_nla="${line#no_nla=}" ;;
                resolution=*) resolution="${line#resolution=}" ;;
            esac
        fi
    done < "$PROFILE_FILE"
    if [[ "$found" -eq 0 ]]; then
        err "Profile '$name' not found"
        exit 1
    fi
    log "Loaded profile '$name'"
}

profile_save() {
    local name="$1"
    if [[ -f "$PROFILE_FILE" ]]; then
        while IFS= read -r line; do
            if [[ "$line" =~ ^\[(.+)\]$ && "${BASH_REMATCH[1]}" == "$name" ]]; then
                err "Profile '$name' already exists (use --delete first)"
                exit 1
            fi
        done < "$PROFILE_FILE"
    fi
    mkdir -p "$PROFILE_DIR"
    {
        echo "[$name]"
        [[ -n "$server" ]]     && echo "server=$server"
        [[ -n "$port" ]]       && echo "port=$port"
        [[ -n "$domain" ]]     && echo "domain=$domain"
        [[ -n "$username" ]]   && echo "username=$username"
        [[ -n "$app" ]]        && echo "app=$app"
        [[ "$fullscreen" -eq 1 ]] && echo "fullscreen=1"
        [[ "$multimon" -eq 1 ]]   && echo "multimon=1"
        [[ -n "$drive" ]]      && echo "drive=$drive"
        [[ "$sound" -eq 1 ]]   && echo "sound=1"
        [[ -n "$gateway" ]]    && echo "gateway=$gateway"
        [[ "$admin" -eq 1 ]]   && echo "admin=1"
        [[ "$no_nla" -eq 1 ]]  && echo "no_nla=1"
        [[ -n "$resolution" ]] && echo "resolution=$resolution"
        echo ""
    } >> "$PROFILE_FILE"
    log "Saved profile '$name'"
}

profile_delete() {
    local name="$1"
    if [[ ! -f "$PROFILE_FILE" ]]; then
        err "Profile '$name' not found (no profiles file)"
        exit 1
    fi
    local found=0 in_section=0
    local tmpfile
    tmpfile=$(mktemp)
    while IFS= read -r line; do
        if [[ "$line" =~ ^\[(.+)\]$ ]]; then
            if [[ "${BASH_REMATCH[1]}" == "$name" ]]; then
                in_section=1
                found=1
                continue
            else
                in_section=0
            fi
        fi
        if [[ "$in_section" -eq 0 ]]; then
            echo "$line" >> "$tmpfile"
        fi
    done < "$PROFILE_FILE"
    if [[ "$found" -eq 0 ]]; then
        rm -f "$tmpfile"
        err "Profile '$name' not found"
        exit 1
    fi
    mv "$tmpfile" "$PROFILE_FILE"
    log "Deleted profile '$name'"
}

# ── Defaults ──────────────────────────────────────────────────────────
server=""
port="3389"
domain=""
username=""
password=""
app=""
fullscreen=0
multimon=0
drive=""
sound=0
gateway=""
admin=0
no_nla=0
resolution=""
opt_profile=""
opt_save=""
opt_list=0
opt_delete=""

# ── Usage ─────────────────────────────────────────────────────────────
usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Launch an RDP session via xfreerdp. Opens a Zenity GUI prompt if connection
details are not fully provided on the command line.

Options:
  --server HOST     Hostname or IP (required)
  --port PORT       RDP port (default: 3389)
  --domain DOM      Active Directory domain (optional)
  --user USER       Login username (required)
  --pass PASS       Password (use - to prompt securely)
  --app NAME        RemoteApp name (optional — launches app instead of desktop)
  --fullscreen      Full screen mode (replaces workarea)
  --multimon        Span across multiple monitors
  --drive NAME,PATH Redirect a local folder to the remote session
  --sound           Enable audio redirection
  --gateway HOST    Connect through an RD Gateway
  --admin           Connect to the console/admin session
  --no-nla          Disable Network Level Authentication
  --resolution WxH  Fixed resolution (e.g., 1920x1080) instead of dynamic
  --profile NAME    Load a saved connection profile
  --save NAME       Save current CLI args as a named profile, then exit
  --list            List saved profiles and exit
  --delete NAME     Remove a saved profile and exit
  -h, --help        Show this help

If --server, --user, and --pass are all provided, connects directly.
Otherwise, opens a Zenity prompt with any provided values pre-filled.
If saved profiles exist, a profile picker is shown before the connection form.

Examples:
  # GUI prompt
  ./$(basename "$0")

  # Direct connect
  ./$(basename "$0") --server 10.0.1.50 --user admin --pass -

  # RemoteApp
  ./$(basename "$0") --server 10.0.1.50 --user admin --pass - --app notepad

  # Pre-fill server, prompt for the rest
  ./$(basename "$0") --server 10.0.1.50

  # Save a profile
  ./$(basename "$0") --server 10.0.1.50 --user pconnor --domain CORP --save work-dc01

  # Connect using a profile
  ./$(basename "$0") --profile work-dc01 --pass -

  # List / delete profiles
  ./$(basename "$0") --list
  ./$(basename "$0") --delete work-dc01
EOF
}

# ── Parse args ────────────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
    case "$1" in
        --server)  server="$2"; shift 2 ;;
        --port)    port="$2"; shift 2 ;;
        --domain)  domain="$2"; shift 2 ;;
        --user)    username="$2"; shift 2 ;;
        --pass)    password="$2"; shift 2 ;;
        --app)     app="$2"; shift 2 ;;
        --fullscreen)  fullscreen=1; shift ;;
        --multimon)    multimon=1; shift ;;
        --drive)       drive="$2"; shift 2 ;;
        --sound)       sound=1; shift ;;
        --gateway)     gateway="$2"; shift 2 ;;
        --admin)       admin=1; shift ;;
        --no-nla)      no_nla=1; shift ;;
        --resolution)  resolution="$2"; shift 2 ;;
        --profile) opt_profile="$2"; shift 2 ;;
        --save)    opt_save="$2"; shift 2 ;;
        --list)    opt_list=1; shift ;;
        --delete)  opt_delete="$2"; shift 2 ;;
        -h|--help) usage; exit 0 ;;
        *)         err "Unknown option: $1"; usage; exit 1 ;;
    esac
done

# ── Profile actions (run and exit) ────────────────────────────────────
if [[ "$opt_list" -eq 1 ]]; then
    profile_list
    exit 0
fi

if [[ -n "$opt_delete" ]]; then
    profile_delete "$opt_delete"
    exit 0
fi

if [[ -n "$opt_save" ]]; then
    profile_save "$opt_save"
    exit 0
fi

if [[ -n "$opt_profile" ]]; then
    profile_load "$opt_profile"
fi

# Secure password prompt if --pass -
if [[ "$password" == "-" ]]; then
    read -rsp "Password: " password
    echo ""
fi

# ── Preflight ─────────────────────────────────────────────────────────
if ! command -v xfreerdp &>/dev/null; then
    err "xfreerdp not found. Install with:"
    err "  Debian/Ubuntu: sudo apt install freerdp3-x11  (or freerdp2-x11)"
    err "  Fedora/RHEL:   sudo dnf install freerdp"
    exit 1
fi

# Detect FreeRDP major version (2 or 3)
FREERDP_VER=2
freerdp_ver_output=$(xfreerdp --version 2>&1 || true)
if [[ "$freerdp_ver_output" =~ ([0-9]+)\.[0-9]+ ]]; then
    FREERDP_VER="${BASH_REMATCH[1]}"
fi

# ── Prompt if needed ───────────────────────────────────────────────────
if [[ -z "$server" || -z "$username" || -z "$password" ]]; then
    if ! command -v zenity &>/dev/null; then
        err "Missing required fields and zenity not found for GUI prompt"
        err "Provide --server, --user, and --pass on the command line"
        exit 1
    fi

    # Profile picker — show if profiles exist and --profile was not given
    if [[ -z "$opt_profile" && -f "$PROFILE_FILE" ]]; then
        profiles=()
        while IFS= read -r line; do
            if [[ "$line" =~ ^\[(.+)\]$ ]]; then
                profiles+=("${BASH_REMATCH[1]}")
            fi
        done < "$PROFILE_FILE"
        if [[ ${#profiles[@]} -gt 0 ]]; then
            picker_args=()
            for p in "${profiles[@]}"; do
                picker_args+=("$p")
            done
            picker_args+=("New connection")
            picked=$(zenity --list --title="Select Profile" \
                --text="Choose a saved profile or start a new connection" \
                --column="Profile" "${picker_args[@]}") || {
                log "Cancelled"
                exit 0
            }
            if [[ "$picked" != "New connection" && -n "$picked" ]]; then
                profile_load "$picked"
            fi
        fi
    fi

    output=$(zenity --forms --title="RDP Connection" \
        --text="Enter connection details" \
        --separator="," \
        --add-entry="Server${server:+ [$server]}" \
        --add-entry="Port${port:+ [$port]}" \
        --add-entry="Domain${domain:+ [$domain]}" \
        --add-entry="Username${username:+ [$username]}" \
        --add-password="Password" \
        --add-entry="Remote App${app:+ [$app]}") || {
        log "Cancelled"
        exit 0
    }

    IFS=, read -r z_server z_port z_domain z_username z_password z_app <<< "$output"

    # Use Zenity values only if CLI values were not already set
    server="${z_server:-$server}"
    port="${z_port:-$port}"
    domain="${z_domain:-$domain}"
    username="${z_username:-$username}"
    password="${z_password:-$password}"
    app="${z_app:-$app}"
fi

if [[ -z "$server" ]]; then
    err "Server is required"
    exit 1
fi

if [[ -z "$username" ]]; then
    err "Username is required"
    exit 1
fi

# ── Build xfreerdp command ────────────────────────────────────────────
cmd=(xfreerdp
    /v:"${server}:${port}"
    /u:"${username}"
    +clipboard
    /from-stdin
)

# Version-specific flags
if [[ "$FREERDP_VER" -ge 3 ]]; then
    cmd+=(/cert:ignore)
else
    cmd+=(/cert-ignore /rfx)
fi

# Display mode
if [[ "$fullscreen" -eq 1 ]]; then
    cmd+=(/f)
elif [[ -n "$resolution" ]]; then
    cmd+=("/size:${resolution}")
else
    cmd+=(/workarea /dynamic-resolution)
fi

# Multi-monitor
[[ "$multimon" -eq 1 ]] && cmd+=(/multimon)

# Domain
[[ -n "$domain" ]] && cmd+=(/d:"${domain}")

# Drive redirection
[[ -n "$drive" ]] && cmd+=("/drive:${drive}")

# Sound
[[ "$sound" -eq 1 ]] && cmd+=(/sound)

# RD Gateway
[[ -n "$gateway" ]] && cmd+=("/g:${gateway}")

# Admin/console session
[[ "$admin" -eq 1 ]] && cmd+=(/admin)

# Disable NLA
if [[ "$no_nla" -eq 1 ]]; then
    if [[ "$FREERDP_VER" -ge 3 ]]; then
        cmd+=(/sec:nla:off)
    else
        cmd+=(-sec-nla)
    fi
fi

# RemoteApp
if [[ -n "$app" ]]; then
    cmd+=(/t:"${app} on ${server}" "/app:||${app}")
else
    cmd+=(/t:"${server}")
fi

# ── Connect ───────────────────────────────────────────────────────────
log "Connecting to ${server}:${port} as ${username}..."
# Password passed via stdin to keep it out of the process list
printf '/p:%s\n' "$password" | "${cmd[@]}"
password=""