#!/bin/bash
#
# RunPod Control Center - A Beautiful CLI Wrapper for runpodctl
# Manage RunPod GPU instances with style
#
# Usage: ./service_runpod_control.sh [COMMAND] [OPTIONS]
#
# Commands:
#   create      Create a new pod from configuration
#   remove      Remove a pod (by name or ID)
#   get         Get pod status (all pods)
#   start       Start a stopped pod
#   stop        Stop a running pod
#   status      Alias for 'get'
#   help        Show help message
#
# Options:
#   -c, --config FILE     Configuration file (default: $PWD/runpod.yml)
#   -n, --dry-run         Show what would be done without executing
#   -h, --help            Show help message
#

set -euo pipefail

# ============================================================================
# COLOR PALETTE - GPU Computing Vibe (Orange/Yellow Theme)
# ============================================================================

RESET='\033[0m'

# Foreground Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
CYAN='\033[0;36m'
ORANGE='\033[38;5;208m'

# Bold
BOLD_RED='\033[1;31m'
BOLD_GREEN='\033[1;32m'
BOLD_YELLOW='\033[1;33m'
BOLD_BLUE='\033[1;34m'
BOLD_MAGENTA='\033[1;35m'
BOLD_CYAN='\033[1;36m'
BOLD_WHITE='\033[1;37m'
BOLD_ORANGE='\033[1;38;5;208m'

# Styles
DIM='\033[2m'

# ============================================================================
# UNICODE CHARACTERS
# ============================================================================

CHECK_MARK="✓"
CROSS_MARK="✗"
ROCKET="🚀"
GPU="🎮"
SERVER="🖥️"
POWER="⚡"
WARNING="⚠️"
INFO="ℹ️"
SPARKLES="✨"
MONEY="💰"
GLOBE="🌐"
PLUG="🔌"
ID="🆔"
STATUS_DOT="●"
ARROW_RIGHT="→"
BOX_LIGHT="─"
BOX_DOUBLE="═"
BOX_VERT="│"
BOX_TL="╔"
BOX_TR="╗"
BOX_BL="╚"
BOX_BR="╝"

# ============================================================================
# CONFIGURATION
# ============================================================================

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG_FILE="${PWD}/runpod.yml"
COMMAND=""
DRY_RUN=false
POD_ID=""

# ============================================================================
# LOGGING FUNCTIONS
# ============================================================================

print_banner() {
    local width=70
    local title="$SERVER  RunPod Control Center  $POWER"
    local title_plain="  RunPod Control Center  "
    local padding=$(( (width - ${#title_plain} - 4) / 2 ))

    echo ""
    echo -e "${BOLD_ORANGE}${BOX_TL}$(printf '%*s' "$width" '' | tr ' ' "$BOX_DOUBLE")${BOX_TR}${RESET}"
    echo -e "${BOLD_ORANGE}${BOX_VERT}$(printf '%*s' "$padding" '')${BOLD_WHITE}${title}${BOLD_ORANGE}$(printf '%*s' "$padding" '')${BOX_VERT}${RESET}"
    echo -e "${BOLD_ORANGE}${BOX_BL}$(printf '%*s' "$width" '' | tr ' ' "$BOX_DOUBLE")${BOX_BR}${RESET}"
    echo ""
}

print_section() {
    local title="$1"
    echo -e "${BOLD_ORANGE}» ${BOLD_WHITE}${title}${RESET}"
    echo -e "${DIM}$(printf '%*s' 70 '' | tr ' ' "$BOX_LIGHT")${RESET}"
}

print_success() {
    echo -e "${BOLD_GREEN}${CHECK_MARK}${RESET} $1"
}

print_error() {
    echo -e "${BOLD_RED}${CROSS_MARK}${RESET} $1" >&2
}

print_warning() {
    echo -e "${BOLD_YELLOW}${WARNING}${RESET} $1"
}

print_info() {
    echo -e "${BOLD_CYAN}${INFO}${RESET} $1"
}

print_detail() {
    echo -e "  ${DIM}${ARROW_RIGHT}${RESET} $1"
}

print_field() {
    local icon="$1"
    local label="$2"
    local value="$3"
    printf "  ${icon} ${BOLD_WHITE}%-12s${RESET} %s\n" "${label}:" "$value"
}

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

check_dependencies() {
    local missing=()

    if ! command -v yq &>/dev/null; then
        missing+=("yq")
    fi

    if ! command -v runpodctl &>/dev/null; then
        missing+=("runpodctl")
    fi

    if [[ ${#missing[@]} -gt 0 ]]; then
        print_error "Missing required dependencies: ${missing[*]}"
        echo ""
        echo "Install with:"
        for dep in "${missing[@]}"; do
            case "$dep" in
                yq)
                    echo "  brew install yq  # or snap install yq"
                    ;;
                runpodctl)
                    echo "  See: https://github.com/runpod/runpodctl"
                    ;;
            esac
        done
        exit 1
    fi
}

load_config() {
    if [[ ! -f "$CONFIG_FILE" ]]; then
        print_error "Configuration file not found: $CONFIG_FILE"
        exit 1
    fi

    # Validate YAML
    if ! yq eval '.' "$CONFIG_FILE" &>/dev/null; then
        print_error "Invalid YAML in configuration file: $CONFIG_FILE"
        exit 1
    fi

    print_info "Using config: ${DIM}${CONFIG_FILE}${RESET}"
}

get_pod_field() {
    local field="$1"
    local default="${2:-}"
    local value

    value=$(yq eval ".$field // \"\"" "$CONFIG_FILE")

    if [[ -z "$value" || "$value" == "null" ]]; then
        echo "$default"
    else
        echo "$value"
    fi
}

get_pod_array() {
    local field="$1"
    yq eval ".$field[]? // \"\"" "$CONFIG_FILE" 2>/dev/null | grep -v '^$' || true
}

# Find pod ID by name from runpodctl get pod output
find_pod_id_by_name() {
    local name="$1"
    local output

    output=$(runpodctl get pod 2>/dev/null) || {
        print_error "Failed to get pod list"
        return 1
    }

    # Parse the output - runpodctl outputs a table
    # Skip header line and find matching name
    echo "$output" | tail -n +2 | while read -r line; do
        local pod_id pod_name
        pod_id=$(echo "$line" | awk '{print $1}')
        pod_name=$(echo "$line" | awk '{print $2}')

        if [[ "$pod_name" == "$name" ]]; then
            echo "$pod_id"
            return 0
        fi
    done
}

get_status_color() {
    local status="$1"
    case "${status^^}" in
        RUNNING)
            echo -e "${BOLD_GREEN}${STATUS_DOT} RUNNING${RESET}"
            ;;
        EXITED|STOPPED)
            echo -e "${BOLD_YELLOW}${STATUS_DOT} STOPPED${RESET}"
            ;;
        ERROR|FAILED)
            echo -e "${BOLD_RED}${STATUS_DOT} ERROR${RESET}"
            ;;
        CREATING|STARTING)
            echo -e "${BOLD_CYAN}${STATUS_DOT} STARTING${RESET}"
            ;;
        *)
            echo -e "${DIM}${STATUS_DOT} ${status}${RESET}"
            ;;
    esac
}

# ============================================================================
# COMMAND IMPLEMENTATIONS
# ============================================================================

cmd_create() {
    print_section "Creating Pod"

    load_config

    local name gpu_type gpu_count template_id network_volume_id
    local container_disk_size volume_size volume_path mem vcpu
    local secure_cloud community_cloud image_name

    # Required fields
    name=$(get_pod_field "pod.name")
    gpu_type=$(get_pod_field "pod.gpuType")

    if [[ -z "$name" ]]; then
        print_error "Pod name is required in configuration"
        exit 1
    fi

    if [[ -z "$gpu_type" ]]; then
        print_error "GPU type is required in configuration"
        exit 1
    fi

    # Optional fields
    gpu_count=$(get_pod_field "pod.gpuCount" "1")
    template_id=$(get_pod_field "pod.templateId")
    network_volume_id=$(get_pod_field "pod.networkVolumeId")
    container_disk_size=$(get_pod_field "pod.containerDiskSize" "20")
    volume_size=$(get_pod_field "pod.volumeSize" "0")
    volume_path=$(get_pod_field "pod.volumePath" "/runpod")
    mem=$(get_pod_field "pod.mem")
    vcpu=$(get_pod_field "pod.vcpu")
    secure_cloud=$(get_pod_field "pod.secureCloud" "false")
    community_cloud=$(get_pod_field "pod.communityCloud" "false")
    image_name=$(get_pod_field "pod.imageName")

    # Display configuration
    echo ""
    print_field "$GPU" "Name" "$name"
    print_field "$POWER" "GPU" "$gpu_type (x$gpu_count)"
    [[ -n "$template_id" ]] && print_field "$SERVER" "Template" "$template_id"
    [[ -n "$network_volume_id" ]] && print_field "$PLUG" "Network Vol" "$network_volume_id"
    [[ -n "$image_name" ]] && print_field "$ROCKET" "Image" "$image_name"

    # Ports
    local ports=()
    while IFS= read -r port; do
        [[ -n "$port" ]] && ports+=("$port")
    done < <(get_pod_array "pod.ports")

    if [[ ${#ports[@]} -gt 0 ]]; then
        print_field "$PLUG" "Ports" "${ports[*]}"
    fi

    # Environment variables
    local env_vars=()
    while IFS= read -r env; do
        [[ -n "$env" ]] && env_vars+=("$env")
    done < <(get_pod_array "pod.env")

    echo ""

    # Build command arguments
    local args=()
    args+=(--name "$name")
    args+=(--gpuType "$gpu_type")
    args+=(--gpuCount "$gpu_count")
    args+=(--containerDiskSize "$container_disk_size")

    [[ -n "$template_id" ]] && args+=(--templateId "$template_id")
    [[ -n "$network_volume_id" ]] && args+=(--networkVolumeId "$network_volume_id")
    [[ -n "$image_name" ]] && args+=(--imageName "$image_name")
    [[ "$volume_size" != "0" ]] && args+=(--volumeSize "$volume_size")
    [[ -n "$volume_path" ]] && args+=(--volumePath "$volume_path")
    [[ -n "$mem" ]] && args+=(--mem "$mem")
    [[ -n "$vcpu" ]] && args+=(--vcpu "$vcpu")
    [[ "$secure_cloud" == "true" ]] && args+=(--secureCloud)
    [[ "$community_cloud" == "true" ]] && args+=(--communityCloud)

    for port in "${ports[@]}"; do
        args+=(--ports "$port")
    done

    for env in "${env_vars[@]}"; do
        args+=(--env "$env")
    done

    if [[ "$DRY_RUN" == "true" ]]; then
        print_warning "Dry run - would execute:"
        echo ""
        echo -e "  ${DIM}runpodctl create pod ${args[*]}${RESET}"
        echo ""
    else
        print_info "Executing: runpodctl create pod..."
        echo ""

        local output
        if output=$(runpodctl create pod "${args[@]}" 2>&1); then
            print_success "Pod created successfully!"
            echo ""
            echo -e "  ${DIM}${output}${RESET}"
            echo ""
        else
            print_error "Failed to create pod"
            echo -e "  ${DIM}${output}${RESET}"
            exit 1
        fi
    fi
}

cmd_get() {
    print_section "Pod Status"
    echo ""

    local output
    if ! output=$(runpodctl get pod 2>&1); then
        print_error "Failed to get pod status"
        echo -e "  ${DIM}${output}${RESET}"
        exit 1
    fi

    # Check if empty
    if [[ -z "$output" || "$output" == *"No pods found"* ]]; then
        print_warning "No pods found"
        return 0
    fi

    # Parse and display pods beautifully
    local header_shown=false
    local count=0

    while IFS= read -r line; do
        # Skip empty lines
        [[ -z "$line" ]] && continue

        # Skip header line
        if [[ "$header_shown" == "false" ]]; then
            header_shown=true
            continue
        fi

        # Parse fields - runpodctl output format varies
        # Typically: ID NAME STATUS GPU COST
        local pod_id pod_name pod_status gpu_info cost_info

        pod_id=$(echo "$line" | awk '{print $1}')
        pod_name=$(echo "$line" | awk '{print $2}')
        pod_status=$(echo "$line" | awk '{print $3}')
        gpu_info=$(echo "$line" | awk '{print $4}')
        cost_info=$(echo "$line" | awk '{print $5}')

        [[ -z "$pod_id" ]] && continue

        ((count++))

        if [[ $count -gt 1 ]]; then
            echo ""
        fi

        print_field "$GPU" "Name" "$pod_name"
        print_field "$ID" "ID" "$pod_id"
        echo -e "  ${STATUS_DOT} ${BOLD_WHITE}Status:${RESET}      $(get_status_color "$pod_status")"
        [[ -n "$gpu_info" && "$gpu_info" != "-" ]] && print_field "$POWER" "GPU" "$gpu_info"
        [[ -n "$cost_info" && "$cost_info" != "-" ]] && print_field "$MONEY" "Cost" "$cost_info"

    done <<< "$output"

    if [[ $count -eq 0 ]]; then
        print_warning "No pods found"
    else
        echo ""
        print_info "Total pods: $count"
    fi
    echo ""
}

cmd_start() {
    print_section "Starting Pod"

    load_config

    local name pod_id
    name=$(get_pod_field "pod.name")

    if [[ -n "$POD_ID" ]]; then
        pod_id="$POD_ID"
    elif [[ -n "$name" ]]; then
        print_info "Looking up pod by name: $name"
        pod_id=$(find_pod_id_by_name "$name")
    fi

    if [[ -z "$pod_id" ]]; then
        print_error "Pod not found. Specify --id or configure pod.name in config."
        exit 1
    fi

    echo ""
    print_field "$ID" "Pod ID" "$pod_id"
    echo ""

    if [[ "$DRY_RUN" == "true" ]]; then
        print_warning "Dry run - would execute:"
        echo ""
        echo -e "  ${DIM}runpodctl start pod $pod_id${RESET}"
        echo ""
    else
        local output
        if output=$(runpodctl start pod "$pod_id" 2>&1); then
            print_success "Pod starting!"
            echo -e "  ${DIM}${output}${RESET}"
        else
            print_error "Failed to start pod"
            echo -e "  ${DIM}${output}${RESET}"
            exit 1
        fi
    fi
    echo ""
}

cmd_stop() {
    print_section "Stopping Pod"

    load_config

    local name pod_id
    name=$(get_pod_field "pod.name")

    if [[ -n "$POD_ID" ]]; then
        pod_id="$POD_ID"
    elif [[ -n "$name" ]]; then
        print_info "Looking up pod by name: $name"
        pod_id=$(find_pod_id_by_name "$name")
    fi

    if [[ -z "$pod_id" ]]; then
        print_error "Pod not found. Specify --id or configure pod.name in config."
        exit 1
    fi

    echo ""
    print_field "$ID" "Pod ID" "$pod_id"
    echo ""

    if [[ "$DRY_RUN" == "true" ]]; then
        print_warning "Dry run - would execute:"
        echo ""
        echo -e "  ${DIM}runpodctl stop pod $pod_id${RESET}"
        echo ""
    else
        local output
        if output=$(runpodctl stop pod "$pod_id" 2>&1); then
            print_success "Pod stopping!"
            echo -e "  ${DIM}${output}${RESET}"
        else
            print_error "Failed to stop pod"
            echo -e "  ${DIM}${output}${RESET}"
            exit 1
        fi
    fi
    echo ""
}

cmd_remove() {
    print_section "Removing Pod"

    load_config

    local name pod_id
    name=$(get_pod_field "pod.name")

    if [[ -n "$POD_ID" ]]; then
        pod_id="$POD_ID"
    elif [[ -n "$name" ]]; then
        print_info "Looking up pod by name: $name"
        pod_id=$(find_pod_id_by_name "$name")
    fi

    if [[ -z "$pod_id" ]]; then
        print_error "Pod not found. Specify --id or configure pod.name in config."
        exit 1
    fi

    echo ""
    print_field "$ID" "Pod ID" "$pod_id"
    print_field "$GPU" "Name" "${name:-unknown}"
    echo ""

    if [[ "$DRY_RUN" == "true" ]]; then
        print_warning "Dry run - would execute:"
        echo ""
        echo -e "  ${DIM}runpodctl remove pod $pod_id${RESET}"
        echo ""
    else
        # Confirm removal
        print_warning "This will permanently delete the pod!"
        echo -n "  Continue? [y/N] "
        read -r confirm

        if [[ "${confirm,,}" != "y" ]]; then
            print_info "Cancelled"
            exit 0
        fi

        echo ""
        local output
        if output=$(runpodctl remove pod "$pod_id" 2>&1); then
            print_success "Pod removed!"
            echo -e "  ${DIM}${output}${RESET}"
        else
            print_error "Failed to remove pod"
            echo -e "  ${DIM}${output}${RESET}"
            exit 1
        fi
    fi
    echo ""
}

cmd_help() {
    echo -e "${BOLD_WHITE}USAGE${RESET}"
    echo "  service_runpod_control.sh [COMMAND] [OPTIONS]"
    echo ""

    echo -e "${BOLD_WHITE}COMMANDS${RESET}"
    echo -e "  ${BOLD_ORANGE}create${RESET}      Create a new pod from configuration"
    echo -e "  ${BOLD_ORANGE}remove${RESET}      Remove a pod (by name or ID)"
    echo -e "  ${BOLD_ORANGE}get${RESET}         Get pod status (all pods)"
    echo -e "  ${BOLD_ORANGE}start${RESET}       Start a stopped pod"
    echo -e "  ${BOLD_ORANGE}stop${RESET}        Stop a running pod"
    echo -e "  ${BOLD_ORANGE}status${RESET}      Alias for 'get'"
    echo -e "  ${BOLD_ORANGE}help${RESET}        Show this help message"
    echo ""

    echo -e "${BOLD_WHITE}OPTIONS${RESET}"
    echo "  -c, --config FILE     Configuration file (default: \$PWD/runpod.yml)"
    echo "  --id POD_ID           Pod ID for start/stop/remove commands"
    echo "  -n, --dry-run         Show what would be done without executing"
    echo "  -h, --help            Show this help message"
    echo ""

    echo -e "${BOLD_WHITE}CONFIGURATION${RESET}"
    echo "  Create a runpod.yml file with your pod configuration:"
    echo ""
    echo -e "  ${DIM}pod:"
    echo "    name: \"my-gpu-pod\""
    echo "    gpuType: \"NVIDIA GeForce RTX 4090\""
    echo "    gpuCount: 1"
    echo "    templateId: \"your-template-id\""
    echo "    networkVolumeId: \"your-volume-id\""
    echo "    ports:"
    echo -e "      - \"22/tcp\"${RESET}"
    echo ""

    echo -e "${BOLD_WHITE}EXAMPLES${RESET}"
    echo "  # Create a pod from config"
    echo "  ./service_runpod_control.sh create"
    echo ""
    echo "  # Get status of all pods"
    echo "  ./service_runpod_control.sh get"
    echo ""
    echo "  # Stop a pod by name (from config)"
    echo "  ./service_runpod_control.sh stop"
    echo ""
    echo "  # Stop a specific pod by ID"
    echo "  ./service_runpod_control.sh stop --id abc123xyz"
    echo ""
}

# ============================================================================
# ARGUMENT PARSING
# ============================================================================

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            create|remove|get|start|stop|status|help)
                COMMAND="$1"
                shift
                ;;
            -c|--config)
                CONFIG_FILE="$2"
                shift 2
                ;;
            --id)
                POD_ID="$2"
                shift 2
                ;;
            -n|--dry-run)
                DRY_RUN=true
                shift
                ;;
            -h|--help)
                COMMAND="help"
                shift
                ;;
            *)
                print_error "Unknown argument: $1"
                echo "Use 'help' command for usage information"
                exit 1
                ;;
        esac
    done
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    parse_args "$@"

    # Default to help if no command
    if [[ -z "$COMMAND" ]]; then
        cmd_help
        exit 0
    fi

    # Check dependencies (except for help)
    if [[ "$COMMAND" != "help" ]]; then
        check_dependencies
    fi

    print_banner

    case "$COMMAND" in
        create)
            cmd_create
            ;;
        remove)
            cmd_remove
            ;;
        get|status)
            cmd_get
            ;;
        start)
            cmd_start
            ;;
        stop)
            cmd_stop
            ;;
        help)
            cmd_help
            ;;
        *)
            print_error "Unknown command: $COMMAND"
            exit 1
            ;;
    esac
}

main "$@"