#!/usr/bin/env bash

set -Eeuo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

CONFIG_FILE="${SCRIPT_DIR}/config"
DATA_DIR="${SCRIPT_DIR}/data"
CLIENT_DIR="${SCRIPT_DIR}/clients"
REGISTRY="${DATA_DIR}/clients.json"

die() {
    printf "ERROR: %s\n" "$*" >&2
    exit 1
}

info() {
    printf "INFO: %s\n" "$*"
}

require_binary() {
    command -v "$1" >/dev/null 2>&1 \
        || die "$1 not installed"
}

load_config() {
    [[ -f "$CONFIG_FILE" ]] || die "Config not found"
    # shellcheck source=/dev/null
    source "$CONFIG_FILE"
}

init_storage() {

    mkdir -p "$DATA_DIR"
    mkdir -p "$CLIENT_DIR"

    if [[ ! -f "$REGISTRY" ]]; then
        echo "[]" > "$REGISTRY"
    fi

    # Проверка валидности JSON
    jq empty "$REGISTRY" \
        || die "Registry JSON is invalid"
}

sanitize() {
    local v="$1"
    v="${v//[^a-zA-Z0-9._-]/_}"
    printf "%s" "$v"
}

get_next_id() {

    local id

    id="$(jq -r '
        if length == 0 then 1
        else (map(.id) | max) + 1
        end
    ' "$REGISTRY")"

    [[ "$id" =~ ^[0-9]+$ ]] \
        || die "Failed to calculate next ID"

    echo "$id"
}

get_last_ip() {

    jq -r '
        if length == 0 then empty
        else (sort_by(.id) | last.ip)
        end
    ' "$REGISTRY"
}

get_first_client_ip() {

    local network="${SERVER_NETWORK}"

    # Удаляем CIDR suffix
    network="${network%%/*}"

    IFS='.' read -r o1 o2 o3 o4 <<< "$network"

    # Первый клиент всегда .2
    echo "${o1}.${o2}.${o3}.2"
}

increment_ip() {

    local current_ip="$1"

    # Если клиентов еще нет —
    # выдаем первый IP из SERVER_NETWORK
    if [[ -z "$current_ip" ]]; then
        get_first_client_ip
        return
    fi

    IFS='.' read -r o1 o2 o3 o4 <<< "$current_ip"

    ((o4++))

    # Защита от выхода за пределы /24
    if ((o4 >= 255)); then
        die "IP pool exhausted"
    fi

    echo "${o1}.${o2}.${o3}.${o4}"
}

generate_client_keys() {

    CLIENT_PRIV="$(wg genkey)"

    CLIENT_PUB="$(
        printf "%s" "$CLIENT_PRIV" | wg pubkey
    )"

    CLIENT_PSK="$(wg genpsk)"
}

create_client_config() {

    local id="$1"
    local name="$2"
    local ip="$3"

    local filename="${CLIENT_DIR}/${name}.conf"

    cat > "$filename" <<EOF
# Client: $name
# ID: $id

[Interface]
PrivateKey = $CLIENT_PRIV
Address = ${ip}/32
DNS = $DNS_SERVER
MTU = $SERVER_MTU

[Peer]
PublicKey = $SERVER_PUBLIC_KEY
PresharedKey = $CLIENT_PSK
AllowedIPs = 0.0.0.0/0
Endpoint = ${SERVER_PUBLIC_IP}:${SERVER_PUBLIC_PORT}
PersistentKeepalive = 25
EOF

    qrencode \
        -s 8 \
        -o "${CLIENT_DIR}/${name}.png" \
        < "$filename"

    info "Client config created: ${name}.conf"
}

registry_add() {

    local id="$1"
    local name="$2"
    local ip="$3"
    local c_status
    local tmp_file
    tmp_file="$(mktemp)"
    c_status="ACTIVE"
    echo "cstatus is: "$c_status""
    jq \
        --argjson id "$id" \
        --arg name "$name" \
        --arg ip "$ip" \
        --arg pub "$CLIENT_PUB" \
        --arg priv "$CLIENT_PRIV" \
        --arg psk "$CLIENT_PSK" \
        --arg status "$c_status" \
        '
        . + [{
            id: $id,
            name: $name,
            ip: $ip,
            public_key: $pub,
            private_key: $priv,
            psk_key: $psk,
            is_enabled: $status,
            created_at: now|floor
        }]
        ' "$REGISTRY" > "$tmp_file" \
        || die "Failed to update registry"

    mv "$tmp_file" "$REGISTRY"

    # Проверка после записи
    jq empty "$REGISTRY" \
        || die "Registry corrupted after write"
}

create_client() {

    local raw_name=""
    local apply_ros=false

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --apply_ros)
                apply_ros=true
                shift
                ;;
            *)
                if [[ -z "$raw_name" ]]; then
                    raw_name="$1"
                else
                    die "Unexpected argument: $1"
                fi
                shift
                ;;
        esac
    done

    [[ -n "$raw_name" ]] \
        || die "Client name required"

    local name
    name="$(sanitize "$raw_name")"

    [[ -n "$name" ]] \
        || die "Sanitized name is empty"

    local id
    id="$(get_next_id)"

    local last_ip
    last_ip="$(get_last_ip)"

    local next_ip
    next_ip="$(increment_ip "$last_ip")"

    generate_client_keys

    create_client_config "$id" "$name" "$next_ip"

    registry_add "$id" "$name" "$next_ip"

    info "Client registered: id=$id name=$name ip=$next_ip"

    print_mikrotik_commands

    info "RouterOS Configuration Script has been Created: id=$id name=$name"

    if [[ "$apply_ros" == true ]]; then
        apply_mikrotik_commands
        info "RouterOS WG Peer has been Configured: id=$id name=$name"
    else
        info "RouterOS WG Peer was not Configured (use --apply_ros to apply configuration)"
    fi

}

list_clients() {

    jq -r '
        if length == 0 then
            "No clients"
        else
            .[] | "ID=\(.id) NAME=\(.name) IP=\(.ip) STATUS=\(.is_enabled)"
        end
    ' "$REGISTRY"
}

disable_client() {

    local id="${1:-}"

    [[ "$id" =~ ^[0-9]+$ ]] \
        || die "Invalid ID"

    local current_status name

    read -r current_status name < <(jq -r --argjson id "$id" '.[] | select(.id == $id) | "\(.is_enabled) \(.name)"' "$REGISTRY" 2>/dev/null)

    if [[ -z "$name" ]]; then
        die "Client with ID $id not found"
    fi

    local new_status
    if [[ "$current_status" == "ACTIVE" ]]; then
        new_status="DISABLED"
    else
        new_status="ACTIVE"
    fi

    local tmp_file
    tmp_file="$(mktemp)"

    jq \
        --argjson id "$id" \
        --arg new_status "$new_status" \
        '
        map(
            if .id == $id then
                .is_enabled = $new_status
            else
                .
            end
        )
        ' "$REGISTRY" > "$tmp_file" \
        || { rm -f "$tmp_file"; die "Failed to update registry status"; }

    mv "$tmp_file" "$REGISTRY"

    info "Client $id ($name) status changed from $current_status to $new_status"

}

enable_client() {

    local id="${1:-}"

    [[ "$id" =~ ^[0-9]+$ ]] \
        || die "Invalid ID"

    local current_status name

    read -r current_status name < <(jq -r --argjson id "$id" '.[] | select(.id == $id) | "\(.is_enabled) \(.name)"' "$REGISTRY" 2>/dev/null)

    if [[ -z "$name" ]]; then
        die "Client with ID $id not found"
    fi

    local new_status
    if [[ "$current_status" == "DISABLED" ]]; then
        new_status="ACTIVE"
    else
        new_status="DISABLED"
    fi

    local tmp_file
    tmp_file="$(mktemp)"

    jq \
        --argjson id "$id" \
        --arg new_status "$new_status" \
        '
        map(
            if .id == $id then
                .is_enabled = $new_status
            else
                .
            end
        )
        ' "$REGISTRY" > "$tmp_file" \
        || { rm -f "$tmp_file"; die "Failed to update registry status"; }

    mv "$tmp_file" "$REGISTRY"

    info "Client $id ($name) status changed from $current_status to $new_status"

}

delete_client() {
    local id="${1:-}"

    [[ "$id" =~ ^[0-9]+$ ]] \
        || die "Invalid ID"

    local name
    name="$(jq -r --argjson id "$id" '.[] | select(.id == $id) | .name' "$REGISTRY" 2>/dev/null)"

    if [[ -z "$name" ]]; then
        die "Client with ID $id not found"
    fi

    local tmp_file
    tmp_file="$(mktemp)"

    jq \
        --argjson id "$id" \
        'map(select(.id != $id))' \
        "$REGISTRY" > "$tmp_file" \
        || { rm -f "$tmp_file"; die "Failed to update registry"; }

    mv "$tmp_file" "$REGISTRY"

    if [[ -n "$name" ]]; then
        rm -f "${CLIENT_DIR}/${name}.conf"
        rm -f "${CLIENT_DIR}/${name}.png"
        rm -f "${CLIENT_DIR}/${name}.rsc"
    fi

    info "Client $id ($name) removed"
}

print_mikrotik_commands() {
    client_ip=$(jq -r --arg val "$name" '.[] | select(.name == $val) | .ip' "$REGISTRY")

    local mt_filename="${CLIENT_DIR}/${name}.rsc"

    cat > "$mt_filename" <<EOF
# RouterOS Config for $name
/interface wireguard peers
add interface="$SERVER_INTERFACE" \\
    public-key="$CLIENT_PUB" \\
    allowed-address="$client_ip/32" \\
    preshared-key="$CLIENT_PSK" \\
    name="$name" \\
    comment="$name"
EOF
}

test_mikrotik_commands() {
    local HOST="$MT_HOST"
    local USER="$MT_USER"
    local PASS="$MT_PASSW"

    sshpass -p "$PASS" ssh \
    -o StrictHostKeyChecking=no \
    -o UserKnownHostsFile=/dev/null \
    "$USER@$HOST" <<'EOF'
echo "start"
/sys identity print
/sys routerboard print
echo "done"
EOF
}

apply_mikrotik_commands() {
    local HOST="$MT_HOST"
    local USER="$MT_USER"
    local PASS="$MT_PASSW"
    LOCAL_FILE="${CLIENT_DIR}/${name}.rsc"
    REMOTE_FILE="${name}.rsc"
    # Upload file via SSH
    sshpass -p "$PASS" scp \
        -O \
        -o StrictHostKeyChecking=no \
        -o UserKnownHostsFile=/dev/null \
        -o PubkeyAuthentication=no \
         "$LOCAL_FILE" "$USER@$HOST:$REMOTE_FILE"
    # Execute file remotely
    sshpass -p "$PASS" ssh \
        -T \
        -o StrictHostKeyChecking=no \
        -o UserKnownHostsFile=/dev/null \
        -o PubkeyAuthentication=no \
        "$USER@$HOST" \
        "/import file-name=$REMOTE_FILE"
    # Cleanup file remotely
    sshpass -p "$PASS" ssh \
        -T \
        -o StrictHostKeyChecking=no \
        -o UserKnownHostsFile=/dev/null \
        -o PubkeyAuthentication=no \
        "$USER@$HOST" \
        "/file/remove $REMOTE_FILE"
}

main() {

    require_binary wg
    require_binary jq
    require_binary qrencode

    load_config
    init_storage

    case "${1:-}" in

        create)
            shift
            create_client "$@"
            ;;

        delete)
            shift
            delete_client "$@"
            ;;

        disable)
            shift
            disable_client "$@"
            ;;

        enable)
            shift
            enable_client "$@"
            ;;

        list)
            list_clients
            ;;

        test)
            test_mikrotik_commands
            ;;

        *)
            echo "Usage:

- create  <name> <option>   # create new client

  NOTE: use --apply_ros command to configure RouterOS WG Peer

- delete  <id>              # delete existing client by ID
- disable <id>              # disable existing client by ID
- enable  <id>              # enable existing client by ID
- list                      # list all clients
- test                      # test remote exec"

            exit 1
            ;;
    esac
}

main "$@"