#!/usr/bin/env bash

# Kubernetes Manifest Apply Script
# Smart apply with validation, dry-run, and rollback support

set -euo pipefail

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# Configuration
DRY_RUN="${DRY_RUN:-false}"
NAMESPACE=""
VALIDATE="${VALIDATE:-true}"
WAIT="${WAIT:-true}"
TIMEOUT="${TIMEOUT:-300s}"

# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VALIDATOR="${SCRIPT_DIR}/validate-k8s.sh"

# Helper functions
print_header() {
    echo -e "\n${BLUE}========================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}========================================${NC}\n"
}

print_success() {
    echo -e "${GREEN}✓ $1${NC}"
}

print_error() {
    echo -e "${RED}✗ $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}⚠ $1${NC}"
}

print_info() {
    echo -e "${BLUE}ℹ $1${NC}"
}

command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Prompt for confirmation
confirm() {
    local prompt="$1"
    local response

    read -p "$(echo -e ${YELLOW}${prompt}${NC}) [y/N]: " response
    [[ "$response" =~ ^[Yy]$ ]]
}

# Validate manifests before applying
validate_manifests() {
    local paths=("$@")

    if [ "$VALIDATE" != "true" ]; then
        print_info "Validation skipped"
        return 0
    fi

    print_header "Validating Manifests"

    if [ ! -x "$VALIDATOR" ]; then
        print_warning "Validator script not found or not executable: $VALIDATOR"
        print_info "Skipping validation..."
        return 0
    fi

    if "$VALIDATOR" "${paths[@]}"; then
        print_success "All manifests validated successfully"
        return 0
    else
        print_error "Validation failed"
        return 1
    fi
}

# Perform dry-run
dry_run_apply() {
    local file="$1"
    local ns_args=()

    if [ -n "$NAMESPACE" ]; then
        ns_args=("-n" "$NAMESPACE")
    fi

    print_info "Dry-run for: $file"

    if kubectl apply --dry-run=server "${ns_args[@]}" -f "$file" 2>&1; then
        print_success "Dry-run successful"
        return 0
    else
        print_error "Dry-run failed"
        return 1
    fi
}

# Show diff before applying
show_diff() {
    local file="$1"
    local ns_args=()

    if [ -n "$NAMESPACE" ]; then
        ns_args=("-n" "$NAMESPACE")
    fi

    if ! command_exists kubectl; then
        return 0
    fi

    print_info "Showing changes for: $file"
    echo ""

    local diff_output
    if diff_output=$(kubectl diff "${ns_args[@]}" -f "$file" 2>&1); then
        if [ -z "$diff_output" ]; then
            print_info "No changes detected"
        else
            echo "$diff_output"
        fi
    else
        # diff returns non-zero when there are differences
        # This is expected, so we just show the output
        if [ -n "$diff_output" ]; then
            echo "$diff_output"
        fi
    fi

    echo ""
}

# Apply a single manifest file
apply_file() {
    local file="$1"
    local ns_args=()
    local wait_args=()

    if [ -n "$NAMESPACE" ]; then
        ns_args=("-n" "$NAMESPACE")
    fi

    if [ "$WAIT" = "true" ]; then
        wait_args=("--wait" "--timeout=$TIMEOUT")
    fi

    print_info "Applying: $file"

    if kubectl apply "${ns_args[@]}" "${wait_args[@]}" -f "$file"; then
        print_success "Applied successfully: $file"
        return 0
    else
        print_error "Failed to apply: $file"
        return 1
    fi
}

# Apply manifests
apply_manifests() {
    local paths=("$@")
    local files=()

    # Collect all YAML files
    for path in "${paths[@]}"; do
        if [ -f "$path" ]; then
            if [[ "$path" =~ \.(yaml|yml)$ ]]; then
                files+=("$path")
            fi
        elif [ -d "$path" ]; then
            while IFS= read -r file; do
                files+=("$file")
            done < <(find "$path" -type f \( -name "*.yaml" -o -name "*.yml" \))
        fi
    done

    if [ ${#files[@]} -eq 0 ]; then
        print_error "No YAML files found"
        return 1
    fi

    print_header "Applying Manifests"

    print_info "Files to apply: ${#files[@]}"

    # Show all files
    for file in "${files[@]}"; do
        echo "  - $file"
    done
    echo ""

    # Show diffs if not in dry-run mode
    if [ "$DRY_RUN" != "true" ]; then
        if command_exists kubectl; then
            for file in "${files[@]}"; do
                show_diff "$file"
            done
        fi
    fi

    # Confirm before applying
    if [ "$DRY_RUN" != "true" ]; then
        if ! confirm "Apply these manifests?"; then
            print_warning "Aborted by user"
            return 1
        fi
    fi

    # Apply each file
    local failed=0
    for file in "${files[@]}"; do
        if [ "$DRY_RUN" = "true" ]; then
            if ! dry_run_apply "$file"; then
                ((failed++))
            fi
        else
            if ! apply_file "$file"; then
                ((failed++))
            fi
        fi
    done

    echo ""

    if [ $failed -eq 0 ]; then
        print_success "All manifests applied successfully (${#files[@]}/${#files[@]})"
        return 0
    else
        print_error "Some manifests failed to apply ($failed/${#files[@]})"
        return 1
    fi
}

# Check status of deployed resources
check_status() {
    local ns_args=()

    if [ -n "$NAMESPACE" ]; then
        ns_args=("-n" "$NAMESPACE")
    fi

    print_header "Resource Status"

    print_info "Deployments:"
    kubectl get deployments "${ns_args[@]}" 2>/dev/null || print_info "No deployments found"

    echo ""
    print_info "StatefulSets:"
    kubectl get statefulsets "${ns_args[@]}" 2>/dev/null || print_info "No statefulsets found"

    echo ""
    print_info "Services:"
    kubectl get services "${ns_args[@]}" 2>/dev/null || print_info "No services found"

    echo ""
    print_info "Pods:"
    kubectl get pods "${ns_args[@]}" 2>/dev/null || print_info "No pods found"

    echo ""
    print_info "Ingresses:"
    kubectl get ingress "${ns_args[@]}" 2>/dev/null || print_info "No ingresses found"

    echo ""
}

# Wait for rollout to complete
wait_for_rollout() {
    local ns_args=()

    if [ -n "$NAMESPACE" ]; then
        ns_args=("-n" "$NAMESPACE")
    fi

    print_header "Waiting for Rollout"

    # Get all deployments
    local deployments
    deployments=$(kubectl get deployments "${ns_args[@]}" -o name 2>/dev/null || echo "")

    if [ -z "$deployments" ]; then
        print_info "No deployments to wait for"
        return 0
    fi

    while IFS= read -r deployment; do
        if [ -n "$deployment" ]; then
            print_info "Waiting for: $deployment"
            if kubectl rollout status "${ns_args[@]}" "$deployment" --timeout="$TIMEOUT"; then
                print_success "Rollout complete: $deployment"
            else
                print_error "Rollout failed: $deployment"
                return 1
            fi
        fi
    done <<< "$deployments"

    return 0
}

# Rollback last deployment
rollback() {
    local resource="$1"
    local ns_args=()

    if [ -n "$NAMESPACE" ]; then
        ns_args=("-n" "$NAMESPACE")
    fi

    print_header "Rolling Back"

    print_warning "Rolling back: $resource"

    if kubectl rollout undo "${ns_args[@]}" "$resource"; then
        print_success "Rollback initiated"

        if [ "$WAIT" = "true" ]; then
            if kubectl rollout status "${ns_args[@]}" "$resource" --timeout="$TIMEOUT"; then
                print_success "Rollback complete"
                return 0
            else
                print_error "Rollback failed"
                return 1
            fi
        fi
        return 0
    else
        print_error "Rollback failed"
        return 1
    fi
}

# Usage information
usage() {
    cat <<EOF
Kubernetes Manifest Apply Script

Usage: $0 [OPTIONS] <file|directory>...

Safely apply Kubernetes manifests with validation, dry-run, and rollback support.

Options:
  -h, --help              Show this help message
  -n, --namespace NS      Target namespace
  -d, --dry-run           Perform dry-run only (no changes)
  --no-validate           Skip manifest validation
  --no-wait              Don't wait for rollout completion
  --timeout TIMEOUT      Timeout for operations (default: 300s)
  -s, --status           Show resource status after apply
  -r, --rollback RES     Rollback a deployment/statefulset

Examples:
  # Dry-run
  $0 --dry-run deployment.yaml

  # Apply with validation
  $0 manifests/

  # Apply to specific namespace
  $0 -n production deployment.yaml

  # Apply and check status
  $0 -s manifests/

  # Rollback a deployment
  $0 --rollback deployment/my-app

Environment Variables:
  DRY_RUN               Enable dry-run mode (true/false)
  VALIDATE              Enable validation (true/false)
  WAIT                  Wait for rollout (true/false)
  TIMEOUT               Timeout duration (default: 300s)

EOF
}

# Parse arguments
parse_args() {
    local paths=()
    local show_status=false
    local rollback_resource=""

    while [ $# -gt 0 ]; do
        case $1 in
            -h|--help)
                usage
                exit 0
                ;;
            -n|--namespace)
                NAMESPACE="$2"
                shift 2
                ;;
            -d|--dry-run)
                DRY_RUN=true
                shift
                ;;
            --no-validate)
                VALIDATE=false
                shift
                ;;
            --no-wait)
                WAIT=false
                shift
                ;;
            --timeout)
                TIMEOUT="$2"
                shift 2
                ;;
            -s|--status)
                show_status=true
                shift
                ;;
            -r|--rollback)
                rollback_resource="$2"
                shift 2
                ;;
            -*)
                print_error "Unknown option: $1"
                usage
                exit 1
                ;;
            *)
                paths+=("$1")
                shift
                ;;
        esac
    done

    # Handle rollback
    if [ -n "$rollback_resource" ]; then
        rollback "$rollback_resource"
        exit $?
    fi

    # Need at least one path
    if [ ${#paths[@]} -eq 0 ]; then
        usage
        exit 1
    fi

    echo "${paths[@]}" "$show_status"
}

# Check prerequisites
check_prerequisites() {
    if ! command_exists kubectl; then
        print_error "kubectl is not installed or not in PATH"
        exit 1
    fi

    # Check cluster connectivity
    if ! kubectl cluster-info >/dev/null 2>&1; then
        print_error "Cannot connect to Kubernetes cluster"
        print_info "Check your kubeconfig and cluster connectivity"
        exit 1
    fi

    print_info "Connected to cluster: $(kubectl config current-context)"

    if [ -n "$NAMESPACE" ]; then
        print_info "Target namespace: $NAMESPACE"

        # Check if namespace exists
        if ! kubectl get namespace "$NAMESPACE" >/dev/null 2>&1; then
            print_warning "Namespace '$NAMESPACE' does not exist"
            if confirm "Create namespace '$NAMESPACE'?"; then
                kubectl create namespace "$NAMESPACE"
                print_success "Namespace created"
            else
                exit 1
            fi
        fi
    fi

    echo ""
}

# Main function
main() {
    local result
    result=($(parse_args "$@"))

    local paths=("${result[@]:0:${#result[@]}-1}")
    local show_status="${result[${#result[@]}-1]}"

    print_header "Kubernetes Manifest Apply"

    check_prerequisites

    # Validate
    if ! validate_manifests "${paths[@]}"; then
        print_error "Validation failed. Aborting."
        exit 1
    fi

    # Apply
    if ! apply_manifests "${paths[@]}"; then
        print_error "Apply failed"
        exit 1
    fi

    # Wait for rollout
    if [ "$DRY_RUN" != "true" ] && [ "$WAIT" = "true" ]; then
        if ! wait_for_rollout; then
            print_error "Rollout failed"
            print_warning "You may want to check the status and consider rollback"
            exit 1
        fi
    fi

    # Show status
    if [ "$show_status" = "true" ] && [ "$DRY_RUN" != "true" ]; then
        check_status
    fi

    print_header "Complete"

    if [ "$DRY_RUN" = "true" ]; then
        print_success "Dry-run completed successfully"
    else
        print_success "Manifests applied successfully"
    fi
}

# Run main function
main "$@"
