#!/usr/bin/env bash

set -euo pipefail

require_cmd() {
    if ! command -v "$1" >/dev/null 2>&1; then
        echo "Missing required command: $1" >&2
        exit 1
    fi
}

ensure_azure_login() {
    if az account show >/dev/null 2>&1; then
        return 0
    fi

    echo "No active Azure CLI session found. Starting az login..." >&2
    az login >/dev/null
}

pick_subscription() {
    local active_id
    local selected_line
    local marker
    local name
    local id
    local tenant

    active_id="$(az account show --query id -o tsv 2>/dev/null || true)"

    mapfile -t subscription_lines < <(
        az account list --all \
            --query "[?state=='Enabled'].[name,id,tenantId]" \
            -o tsv |
        while IFS=$'\t' read -r name id tenant; do
            marker=" "
            if [[ "$id" == "$active_id" ]]; then
                marker=">"
            fi

            printf '%s\t%s\t%s\t%s\n' "$marker" "$name" "$id" "$tenant"
        done
    )

    if [[ "${#subscription_lines[@]}" -eq 0 ]]; then
        echo "No enabled Azure subscriptions found for the current login." >&2
        exit 1
    fi

    local fzf_output
    local exit_status=0
    fzf_output="$(
        printf '%s\n' "${subscription_lines[@]}" |
        fzf \
            --height=40% \
            --layout=reverse \
            --border \
            --prompt="az subscription > " \
            --header=$'> current\nPick an Azure subscription' \
            --delimiter=$'\t' \
            --with-nth=1,2,4 \
            --preview='printf "Subscription ID: %s\nTenant ID: %s\n" {3} {4}' \
            --print-query
    )" || exit_status=$?

    if [[ "$exit_status" -ne 0 ]]; then
        return "$exit_status"
    fi

    local query
    query="$(head -n 1 <<<"$fzf_output")"
    selected_line="$(tail -n +2 <<<"$fzf_output")"

    local subscription_id
    subscription_id="$(awk -F '\t' '{print $3}' <<<"$selected_line")"

    printf '%s\n%s\n' "$query" "$subscription_id"
}

pick_kube_context() {
    local current_context
    local selected_line
    local context
    local marker
    local initial_query="${1:-}"

    current_context="$(kubectl config current-context 2>/dev/null || true)"

    mapfile -t context_lines < <(
        kubectl config get-contexts -o name 2>/dev/null |
        while IFS= read -r context; do
            marker=" "
            if [[ "$context" == "$current_context" ]]; then
                marker=">"
            fi

            printf '%s\t%s\n' "$marker" "$context"
        done
    )

    if [[ "${#context_lines[@]}" -eq 0 ]]; then
        return 2
    fi

    selected_line="$(
        printf '%s\n' "${context_lines[@]}" |
        fzf \
            --height=40% \
            --layout=reverse \
            --border \
            --prompt="kubectl context > " \
            --header=$'> current\nPick a kubectl context' \
            --delimiter=$'\t' \
            --with-nth=1,2 \
            --query="$initial_query"
    )" || return 1

    awk -F '\t' '{print $2}' <<<"$selected_line"
}

print_final_state() {
    local sub_name sub_id tenant_id kube_context

    sub_name="$(az account show --query name -o tsv)"
    sub_id="$(az account show --query id -o tsv)"
    tenant_id="$(az account show --query tenantId -o tsv)"
    kube_context="$(kubectl config current-context 2>/dev/null || true)"

    printf '\nActive Azure subscription: %s (%s)\n' "$sub_name" "$sub_id"
    printf 'Active Azure tenant: %s\n' "$tenant_id"

    if [[ -n "$kube_context" ]]; then
        printf 'Active kubectl context: %s\n' "$kube_context"
    else
        printf 'Active kubectl context: <none>\n'
    fi
}

main() {
    local subscription_id
    local kube_context

    require_cmd az
    require_cmd kubectl
    require_cmd fzf

    ensure_azure_login

    local subscription_output
    if ! subscription_output="$(pick_subscription)"; then
        echo "Subscription selection cancelled." >&2
        exit 130
    fi

    local query
    query="$(head -n 1 <<<"$subscription_output")"
    subscription_id="$(tail -n +2 <<<"$subscription_output")"

    az account set --subscription "$subscription_id"

    if kube_context="$(pick_kube_context "$query")"; then
        kubectl config use-context "$kube_context" >/dev/null
    else
        case $? in
            1)
                echo "Kubectl context selection cancelled. Azure subscription was still updated." >&2
                ;;
            2)
                echo "No kubectl contexts found. Azure subscription was updated." >&2
                ;;
            *)
                echo "Kubectl context selection failed. Azure subscription was still updated." >&2
                ;;
        esac
    fi

    print_final_state
}

main "$@"
