#!/usr/bin/env bash

set -euo pipefail

DB_NAME="${DB_NAME:-dosis2}"
DB_USER="${DB_USER:-dsuper}"
DB_HOST="${DB_HOST:-localhost}"
DB_PORT="${DB_PORT:-5432}"
STATE_DIR="${STATE_DIR:-/var/lib/tower-tunnels}"
RUNTIME_DIR="${RUNTIME_DIR:-/run/tower-tunnels}"
LOG_DIR="${LOG_DIR:-/var/log/tower-tunnels}"
DEFAULT_SSH_USER="${DEFAULT_SSH_USER:-dPG}"
DEFAULT_PRIVATE_KEY="${DEFAULT_PRIVATE_KEY:-/home/dPG/.ssh/fuse/id_rsa}"
DEFAULT_KNOWN_HOSTS_FILE="${DEFAULT_KNOWN_HOSTS_FILE:-/root/.ssh/known_hosts}"
DEFAULT_PRINTER_PORT="${DEFAULT_PRINTER_PORT:-8080}"
TELEPORT_CONFIG="${TELEPORT_CONFIG:-/etc/teleport.yaml}"
SETUP_IMAGE_SCRIPT="/opt/dosis2/bin/setup-image-share"

mkdir -p "$STATE_DIR" "$RUNTIME_DIR" "$LOG_DIR"

for cmd in psql ssh systemctl systemd-escape setsid; do
    command -v "$cmd" >/dev/null 2>&1 || {
        echo "ERROR: Required command '$cmd' is not installed or not in PATH."
        exit 1
    }
done

is_pid_running() {
    local pid="$1"
    [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null
}

stop_tunnel() {
    local instance_id="$1"
    local pid_file="$RUNTIME_DIR/${instance_id}.pid"
    local meta_file="$STATE_DIR/${instance_id}.meta"

    if [[ -f "$pid_file" ]]; then
        local pid
        pid="$(<"$pid_file")"
        if is_pid_running "$pid"; then
            kill "$pid" 2>/dev/null || true
            for _ in $(seq 1 10); do
                if ! is_pid_running "$pid"; then
                    break
                fi
                sleep 1
            done
            if is_pid_running "$pid"; then
                kill -9 "$pid" 2>/dev/null || true
            fi
        fi
    fi

    rm -f "$pid_file" "$meta_file"
}

start_tunnel() {
    local instance_name="$1"
    local instance_id="$2"
    local local_port="$3"
    local ssh_host="$4"
    local printer_host="$5"

    local pid_file="$RUNTIME_DIR/${instance_id}.pid"
    local meta_file="$STATE_DIR/${instance_id}.meta"
    local log_file="$LOG_DIR/${instance_id}.log"

    : >"$log_file"
    chmod 0600 "$log_file"

    setsid /usr/bin/ssh \
        -i "$DEFAULT_PRIVATE_KEY" \
        -N \
        -T \
        -L "127.0.0.1:${local_port}:${printer_host}:${DEFAULT_PRINTER_PORT}" \
        -o UserKnownHostsFile="${DEFAULT_KNOWN_HOSTS_FILE}" \
        -o StrictHostKeyChecking=accept-new \
        -o ExitOnForwardFailure=yes \
        -o ServerAliveInterval=30 \
        -o ServerAliveCountMax=3 \
        "${DEFAULT_SSH_USER}@${ssh_host}" >>"$log_file" 2>&1 &

    local pid=$!
    echo "$pid" >"$pid_file"

    cat >"$meta_file" <<EOF
INSTANCE_NAME=${instance_name}
LOCAL_PORT=${local_port}
SSH_HOST=${ssh_host}
PRINTER_HOST=${printer_host}
EOF
    chmod 0600 "$meta_file" "$pid_file"
}

read -r -d '' TOWER_QUERY <<SQL || true
SELECT
    dafe.tag AS instance_name,
    dafe.site_ip AS ssh_host,
    '172.20.4.101' AS printer_host
FROM equipment.dafe dafe
WHERE dafe.active = TRUE AND dafe.id > 1 AND dafe.tag LIKE 'U60%'
ORDER BY dafe.id;
SQL

read -r -d '' HMI_QUERY <<SQL || true
SELECT
    dafe.tag AS instance_name,
    dafe.site_ip AS ssh_host
FROM equipment.dafe dafe
WHERE dafe.active = TRUE AND dafe.id > 1
ORDER BY dafe.id;
SQL

mapfile -t tower_rows < <(
    PGPASSWORD="${PGPASSWORD:-}" \
    psql \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USER" \
        --dbname="$DB_NAME" \
        --no-align \
        --tuples-only \
        --field-separator=$'\t' \
        --set=ON_ERROR_STOP=1 \
        -c "$TOWER_QUERY"
)

mapfile -t hmi_rows < <(
    PGPASSWORD="${PGPASSWORD:-}" \
    psql \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USER" \
        --dbname="$DB_NAME" \
        --no-align \
        --tuples-only \
        --field-separator=$'\t' \
        --set=ON_ERROR_STOP=1 \
        -c "$HMI_QUERY"
)

declare -A desired_meta=()
declare -A desired_names=()
ordered_instance_ids=()

port_offset=0
for idx in "${!tower_rows[@]}"; do
    row="${tower_rows[$idx]}"
    [[ -n "$row" ]] || continue

    local_port=$((18080 + port_offset))
    port_offset=$((port_offset + 1))
    IFS=$'\t' read -r instance_name ssh_host printer_host <<<"$row"

    if [[ -z "${instance_name:-}" || -z "${ssh_host:-}" || -z "${printer_host:-}" ]]; then
        echo "ERROR: Query returned an incomplete row: $row"
        exit 1
    fi

    instance_id="$(systemd-escape -- "$instance_name")"
    meta_value="${instance_name}"$'\t'"${local_port}"$'\t'"${ssh_host}"$'\t'"${printer_host}"

    if [[ -n "${desired_meta[$instance_id]:-}" ]]; then
        echo "ERROR: Query returned duplicate instance name after escaping: ${instance_name}"
        exit 1
    fi

    desired_meta["$instance_id"]="$meta_value"
    desired_names["$instance_id"]="$instance_name"
    ordered_instance_ids+=("$instance_id")
done

declare -A hmi_meta=()
ordered_hmi_instance_ids=()

for row in "${hmi_rows[@]}"; do
    [[ -n "$row" ]] || continue

    IFS=$'\t' read -r instance_name ssh_host <<<"$row"

    if [[ -z "${instance_name:-}" || -z "${ssh_host:-}" ]]; then
        echo "ERROR: HMI query returned an incomplete row: $row"
        exit 1
    fi

    instance_id="$(systemd-escape -- "$instance_name")"

    if [[ -n "${hmi_meta[$instance_id]:-}" ]]; then
        echo "ERROR: HMI query returned duplicate instance name after escaping: ${instance_name}"
        exit 1
    fi

    hmi_meta["$instance_id"]="${instance_name}"$'\t'"${ssh_host}"
    ordered_hmi_instance_ids+=("$instance_id")
done

tower_serial_name() {
    local instance_name="$1"
    local serial="$instance_name"

    if [[ "$instance_name" =~ ^[UL]60_(.+)$ ]]; then
        serial="${BASH_REMATCH[1]}"
    fi

    printf '%s' "${serial,,}"
}

for pid_path in "$RUNTIME_DIR"/*.pid; do
    [[ -e "$pid_path" ]] || break

    instance_id="$(basename "$pid_path" .pid)"
    if [[ -z "${desired_meta[$instance_id]:-}" ]]; then
        echo "Stopping stale tunnel: ${instance_id}"
        stop_tunnel "$instance_id"
        continue
    fi

    meta_file="$STATE_DIR/${instance_id}.meta"
    expected_meta="${desired_meta[$instance_id]}"
    current_meta=''

    if [[ -f "$meta_file" ]]; then
        current_meta="$(
            source "$meta_file"
            printf '%s\t%s\t%s\t%s' \
                "${INSTANCE_NAME:-}" \
                "${LOCAL_PORT:-}" \
                "${SSH_HOST:-}" \
                "${PRINTER_HOST:-}"
        )"
    fi

    current_pid="$(<"$pid_path")"
    if [[ "$current_meta" != "$expected_meta" ]] || ! is_pid_running "$current_pid"; then
        echo "Refreshing tunnel: ${desired_names[$instance_id]}"
        stop_tunnel "$instance_id"
    fi
done

for instance_id in "${!desired_meta[@]}"; do
    pid_file="$RUNTIME_DIR/${instance_id}.pid"
    if [[ -f "$pid_file" ]]; then
        current_pid="$(<"$pid_file")"
        if is_pid_running "$current_pid"; then
            continue
        fi
        rm -f "$pid_file"
    fi

    IFS=$'\t' read -r instance_name local_port ssh_host printer_host <<<"${desired_meta[$instance_id]}"
    echo "Starting tunnel: ${instance_name} on localhost:${local_port}"
    start_tunnel "$instance_name" "$instance_id" "$local_port" "$ssh_host" "$printer_host"
done

refresh_image_shares() {
    if [[ ! -f "$SETUP_IMAGE_SCRIPT" ]]; then
        echo "Cannot refresh tower image mount dirs. Exiting..."
        exit 1
    fi

    eval $SETUP_IMAGE_SCRIPT
}

refresh_image_shares
