#!/bin/bash
set -euo pipefail

# Tree Runner - Execute cookbook trees with checkpoint support
#
# Usage:
#   ./tree-runner.sh run <tree> [--var-name value]...
#   ./tree-runner.sh resume <tree>
#   ./tree-runner.sh only <tree> <step>
#   ./tree-runner.sh status <tree>
#   ./tree-runner.sh teardown <tree>
#   ./tree-runner.sh list
#   ./tree-runner.sh clean <tree>
#
# Examples:
#   ./tree-runner.sh run landing-page --project-name my-test
#   ./tree-runner.sh resume landing-page
#   ./tree-runner.sh only landing-page wait-pipeline
#   ./tree-runner.sh status landing-page
#   ./tree-runner.sh teardown landing-page
#   ./tree-runner.sh list
#   ./tree-runner.sh clean landing-page

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

# Source dependencies
source "$SCRIPT_DIR/common.sh"
source "$SCRIPT_DIR/lib/checkpoint.sh"
source "$SCRIPT_DIR/lib/tree-parser.sh"

# Parse command
COMMAND="${1:-}"

if [[ -z "$COMMAND" ]]; then
    echo "Tree Runner - Execute cookbook trees with checkpoint support"
    echo ""
    echo "Usage: $0 <command> [args]"
    echo ""
    echo "Commands:"
    echo "  run <tree> [--var-name value]...  Run a tree from the beginning"
    echo "  resume <tree>                      Resume from last checkpoint"
    echo "  only <tree> <step>                 Run only a specific step"
    echo "  status <tree>                      Show checkpoint status"
    echo "  teardown <tree>                    Run tree's teardown steps"
    echo "  list                               List available trees"
    echo "  clean <tree>                       Delete checkpoint for a tree"
    echo ""
    echo "Examples:"
    echo "  $0 run landing-page --project-name my-test-\$(date +%s)"
    echo "  $0 resume landing-page"
    echo "  $0 only landing-page wait-pipeline"
    echo "  $0 status landing-page"
    echo "  $0 teardown landing-page"
    echo "  $0 list"
    echo "  $0 clean landing-page"
    echo ""
    exit 1
fi

shift

# ============================================================================
# Step Executors
# ============================================================================

# Execute an API step
# Arguments: step_json
# Returns: response JSON
execute_api_step() {
    local step_json="$1"

    local method endpoint body
    method=$(echo "$step_json" | jq -r '.method // "GET"')
    endpoint=$(echo "$step_json" | jq -r '.endpoint')
    body=$(echo "$step_json" | jq -c '.body // null')

    local response
    if [[ "$body" != "null" ]]; then
        response=$(api_call "$method" "$endpoint" "$body")
    else
        response=$(api_call "$method" "$endpoint")
    fi

    echo "$response"
}

# Execute a wait_pipeline step
# Arguments: step_json
# Returns: 0 on success, 1 on failure
execute_wait_pipeline_step() {
    local step_json="$1"

    local project_id max_attempts poll_interval
    project_id=$(echo "$step_json" | jq -r '.project_id')
    max_attempts=$(echo "$step_json" | jq -r '.max_attempts // 60')
    poll_interval=$(echo "$step_json" | jq -r '.poll_interval // 5')

    wait_for_pipeline "$project_id" "$max_attempts" "$poll_interval"
}

# Execute a wait_site step
# Arguments: step_json
# Returns: 0 on success, 1 on failure
execute_wait_site_step() {
    local step_json="$1"

    local domain project_id max_attempts poll_interval
    domain=$(echo "$step_json" | jq -r '.domain')
    project_id=$(echo "$step_json" | jq -r '.project_id // ""')
    max_attempts=$(echo "$step_json" | jq -r '.max_attempts // 30')
    poll_interval=$(echo "$step_json" | jq -r '.poll_interval // 5')

    wait_for_site "$domain" "$max_attempts" "$poll_interval" "$project_id"
}

# Execute a diagnose step
# Arguments: step_json
execute_diagnose_step() {
    local step_json="$1"

    local diag_type project_id domain
    diag_type=$(echo "$step_json" | jq -r '.type // "pipeline"')
    project_id=$(echo "$step_json" | jq -r '.project_id // ""')
    domain=$(echo "$step_json" | jq -r '.domain // ""')

    case "$diag_type" in
        pipeline)
            diagnose_pipeline_failure "$project_id"
            ;;
        site)
            diagnose_site_failure "$domain" "$project_id"
            ;;
        *)
            print_error "Unknown diagnose type: $diag_type"
            return 1
            ;;
    esac
}

# Execute a shell step
# Arguments: step_json
# Returns: command output
execute_shell_step() {
    local step_json="$1"

    local command
    command=$(echo "$step_json" | jq -r '.command')

    eval "$command"
}

# Extract outputs from response
# Arguments: response_json output_rules_json
# Returns: JSON object of extracted outputs
extract_outputs() {
    local response="$1"
    local output_rules="$2"

    if [[ "$output_rules" == "[]" || -z "$output_rules" ]]; then
        echo "{}"
        return
    fi

    # Extract each output using jq
    local outputs="{}"
    while IFS= read -r rule; do
        local name jq_path value
        name=$(echo "$rule" | jq -r '.name')
        jq_path=$(echo "$rule" | jq -r '.jq_path')

        value=$(echo "$response" | jq -r "$jq_path // \"\"")
        outputs=$(echo "$outputs" | jq --arg k "$name" --arg v "$value" '. + {($k): $v}')
    done < <(echo "$output_rules" | jq -c '.[]')

    echo "$outputs"
}

# Execute a single step
# Arguments: tree_name step_name vars_json outputs_json
# Returns: 0 on success, 1 on failure
# Updates outputs_json with new outputs
execute_step() {
    local tree_name="$1"
    local step_name="$2"
    local vars_json="$3"
    local outputs_json="$4"

    # Get and expand step definition
    local step_raw
    step_raw=$(tree_get_step "$tree_name" "$step_name") || return 1

    local step
    step=$(tree_expand_step "$step_raw" "$vars_json" "$outputs_json")

    local action description on_error
    action=$(tree_step_action "$step")
    description=$(echo "$step" | jq -r '.description // ""')
    on_error=$(tree_step_on_error "$step")

    # Print step header (to stderr so it doesn't mix with JSON output)
    if [[ -n "$description" ]]; then
        echo -e "${CYAN}Step: $step_name${NC} - $description" >&2
    else
        echo -e "${CYAN}Step: $step_name${NC}" >&2
    fi

    # Mark step as started
    checkpoint_step_start "$tree_name" "$step_name"

    local response=""
    local step_failed=0

    case "$action" in
        api)
            response=$(execute_api_step "$step") || step_failed=1
            if [[ $step_failed -eq 0 ]]; then
                # Check for error in response
                local error
                error=$(echo "$response" | jq -r '.error // ""')
                if [[ -n "$error" && "$error" != "null" ]]; then
                    print_error "API error: $error" >&2
                    step_failed=1
                fi
            fi
            ;;
        wait_pipeline)
            # Redirect status output to stderr so it doesn't pollute JSON return
            execute_wait_pipeline_step "$step" >&2 || step_failed=1
            response="{}"
            ;;
        wait_site)
            # Redirect status output to stderr so it doesn't pollute JSON return
            execute_wait_site_step "$step" >&2 || step_failed=1
            response="{}"
            ;;
        diagnose)
            execute_diagnose_step "$step" >&2
            response="{}"
            ;;
        shell)
            response=$(execute_shell_step "$step") || step_failed=1
            # Try to parse as JSON, otherwise wrap in object
            if ! echo "$response" | jq -e '.' > /dev/null 2>&1; then
                response=$(jq -n --arg out "$response" '{output: $out}')
            fi
            ;;
        *)
            print_error "Unknown action type: $action" >&2
            checkpoint_step_fail "$tree_name" "$step_name" "Unknown action: $action"
            return 1
            ;;
    esac

    if [[ $step_failed -eq 1 ]]; then
        checkpoint_step_fail "$tree_name" "$step_name" "Step failed"
        if [[ "$on_error" == "continue" ]]; then
            print_warning "Step failed but continuing (on_error: continue)" >&2
            checkpoint_step_complete "$tree_name" "$step_name" "{}"
            return 0
        fi
        return 1
    fi

    # Extract outputs if defined
    local output_rules
    output_rules=$(tree_step_outputs "$step")

    local step_outputs
    step_outputs=$(extract_outputs "$response" "$output_rules")

    # Save outputs to checkpoint
    checkpoint_step_complete "$tree_name" "$step_name" "$step_outputs"

    # Return outputs for use by subsequent steps (this is the only stdout)
    echo "$step_outputs"

    print_success "Step completed: $step_name" >&2
    return 0
}

# Build outputs JSON from checkpoint
# Arguments: tree_name
# Returns: outputs JSON object
build_outputs_from_checkpoint() {
    local tree_name="$1"

    local checkpoint
    checkpoint=$(checkpoint_load "$tree_name") || {
        echo "{}"
        return
    }

    echo "$checkpoint" | jq '
        .steps | to_entries |
        map(select(.value.status == "completed")) |
        map({key: .key, value: .value.output}) |
        from_entries
    '
}

# ============================================================================
# Commands
# ============================================================================

# Run a tree from the beginning
cmd_run() {
    local tree_name="${1:-}"
    shift || true

    if [[ -z "$tree_name" ]]; then
        echo "Usage: $0 run <tree> [--var-name value]..."
        exit 1
    fi

    # Validate tree exists
    if ! tree_parse "$tree_name" > /dev/null 2>&1; then
        print_error "Tree '$tree_name' not found"
        echo ""
        echo "Available trees:"
        tree_list_detail
        exit 1
    fi

    # Parse variables from args
    local vars_json
    vars_json=$(tree_get_default_vars "$tree_name")

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --*)
                local var_name="${1#--}"
                var_name="${var_name//-/_}"  # Convert dashes to underscores
                local var_value="${2:-}"
                if [[ -z "$var_value" ]]; then
                    print_error "Missing value for --$var_name"
                    exit 1
                fi
                vars_json=$(echo "$vars_json" | jq --arg k "$var_name" --arg v "$var_value" '. + {($k): $v}')
                shift 2
                ;;
            *)
                print_error "Unknown argument: $1"
                exit 1
                ;;
        esac
    done

    # Check required vars (empty string values)
    local missing_vars
    missing_vars=$(echo "$vars_json" | jq -r 'to_entries | .[] | select(.value == "") | .key')
    if [[ -n "$missing_vars" ]]; then
        print_error "Missing required variables:"
        echo "$missing_vars" | sed 's/^/  --/'
        exit 1
    fi

    # Initialize checkpoint
    local run_id
    run_id=$(checkpoint_init "$tree_name" "$vars_json")

    print_header "Running tree: $tree_name"
    echo "Run ID: $run_id"
    echo "Variables:"
    echo "$vars_json" | jq -r 'to_entries | .[] | "  \(.key): \(.value)"'
    echo ""

    # Get execution order
    local execution_order
    execution_order=$(tree_execution_order "$tree_name")

    # Build outputs as we go
    local outputs_json="{}"

    # Execute steps in order
    while IFS= read -r step_name; do
        local step_outputs
        if step_outputs=$(execute_step "$tree_name" "$step_name" "$vars_json" "$outputs_json"); then
            # Merge step outputs into cumulative outputs
            outputs_json=$(echo "$outputs_json" | jq --arg step "$step_name" --argjson out "$step_outputs" '. + {($step): $out}')
        else
            print_error "Tree execution failed at step: $step_name"
            echo ""
            echo "To resume from this point:"
            echo "  $0 resume $tree_name"
            echo ""
            echo "To run just this step:"
            echo "  $0 only $tree_name $step_name"
            echo ""
            echo "To teardown resources:"
            echo "  $0 teardown $tree_name"
            exit 1
        fi
        echo ""
    done <<< "$execution_order"

    # Mark as completed
    checkpoint_mark_completed "$tree_name"

    print_header "Tree completed successfully!"
    echo "Run ID: $run_id"
    echo ""
    echo "To view final state:"
    echo "  $0 status $tree_name"
    echo ""
    echo "To teardown resources:"
    echo "  $0 teardown $tree_name"
}

# Resume from last checkpoint
cmd_resume() {
    local tree_name="${1:-}"

    if [[ -z "$tree_name" ]]; then
        echo "Usage: $0 resume <tree>"
        exit 1
    fi

    # Load checkpoint
    local checkpoint
    if ! checkpoint=$(checkpoint_load "$tree_name"); then
        print_error "No checkpoint found for tree: $tree_name"
        echo ""
        echo "Run the tree first:"
        echo "  $0 run $tree_name [--var-name value]..."
        exit 1
    fi

    local status
    status=$(echo "$checkpoint" | jq -r '.status')

    if [[ "$status" == "completed" ]]; then
        print_success "Tree already completed. Nothing to resume."
        echo ""
        echo "To run again:"
        echo "  $0 clean $tree_name && $0 run $tree_name ..."
        exit 0
    fi

    print_header "Resuming tree: $tree_name"
    echo "Status: $status"

    # Get vars from checkpoint
    local vars_json
    vars_json=$(echo "$checkpoint" | jq '.vars')

    # Build outputs from completed steps
    local outputs_json
    outputs_json=$(build_outputs_from_checkpoint "$tree_name")

    # Get completed steps
    local completed_steps
    completed_steps=$(checkpoint_completed_steps "$tree_name")

    echo "Completed steps:"
    if [[ -n "$completed_steps" ]]; then
        echo "$completed_steps" | sed 's/^/  ✓ /'
    else
        echo "  (none)"
    fi
    echo ""

    # Get execution order
    local execution_order
    execution_order=$(tree_execution_order "$tree_name")

    # Execute remaining steps
    while IFS= read -r step_name; do
        # Skip completed steps
        if echo "$completed_steps" | grep -q "^${step_name}$"; then
            continue
        fi

        local step_outputs
        if step_outputs=$(execute_step "$tree_name" "$step_name" "$vars_json" "$outputs_json"); then
            outputs_json=$(echo "$outputs_json" | jq --arg step "$step_name" --argjson out "$step_outputs" '. + {($step): $out}')
        else
            print_error "Tree execution failed at step: $step_name"
            echo ""
            echo "To resume again:"
            echo "  $0 resume $tree_name"
            exit 1
        fi
        echo ""
    done <<< "$execution_order"

    checkpoint_mark_completed "$tree_name"

    print_header "Tree completed successfully!"
}

# Run only a specific step
cmd_only() {
    local tree_name="${1:-}"
    local step_name="${2:-}"

    if [[ -z "$tree_name" || -z "$step_name" ]]; then
        echo "Usage: $0 only <tree> <step>"
        exit 1
    fi

    # Validate step exists
    if ! tree_get_step "$tree_name" "$step_name" > /dev/null 2>&1; then
        print_error "Step '$step_name' not found in tree '$tree_name'"
        echo ""
        echo "Available steps:"
        tree_get_steps "$tree_name" | sed 's/^/  /'
        exit 1
    fi

    # Load checkpoint (or use empty state)
    local vars_json="{}"
    local outputs_json="{}"

    if checkpoint=$(checkpoint_load "$tree_name" 2>/dev/null); then
        vars_json=$(echo "$checkpoint" | jq '.vars')
        outputs_json=$(build_outputs_from_checkpoint "$tree_name")
    fi

    print_header "Running single step: $step_name"
    echo "Tree: $tree_name"
    echo ""

    local step_outputs
    if step_outputs=$(execute_step "$tree_name" "$step_name" "$vars_json" "$outputs_json"); then
        print_success "Step completed"
        echo ""
        echo "Outputs:"
        echo "$step_outputs" | jq '.'
    else
        print_error "Step failed"
        exit 1
    fi
}

# Show checkpoint status
cmd_status() {
    local tree_name="${1:-}"

    if [[ -z "$tree_name" ]]; then
        echo "Usage: $0 status <tree>"
        exit 1
    fi

    checkpoint_status_detail "$tree_name"
}

# Run teardown steps
cmd_teardown() {
    local tree_name="${1:-}"

    if [[ -z "$tree_name" ]]; then
        echo "Usage: $0 teardown <tree>"
        exit 1
    fi

    print_header "Teardown: $tree_name"

    # Load checkpoint for outputs
    local vars_json="{}"
    local outputs_json="{}"

    if checkpoint=$(checkpoint_load "$tree_name" 2>/dev/null); then
        vars_json=$(echo "$checkpoint" | jq '.vars')
        outputs_json=$(build_outputs_from_checkpoint "$tree_name")
    else
        print_warning "No checkpoint found - teardown may not have all variables"
    fi

    # Get teardown steps
    local teardown_steps
    teardown_steps=$(tree_get_teardown "$tree_name")

    local teardown_count
    teardown_count=$(echo "$teardown_steps" | jq 'length')

    if [[ "$teardown_count" -eq 0 ]]; then
        echo "No teardown steps defined for this tree."
        return 0
    fi

    echo "Running $teardown_count teardown steps..."
    echo ""

    # Execute teardown steps
    local i=0
    while IFS= read -r step_json; do
        ((i++))

        # Expand templates
        local step
        step=$(tree_expand_step "$step_json" "$vars_json" "$outputs_json")

        local action description
        action=$(echo "$step" | jq -r '.action // "unknown"')
        description=$(echo "$step" | jq -r '.description // "Teardown step $i"')

        echo -e "${CYAN}Teardown $i:${NC} $description"

        case "$action" in
            api)
                execute_api_step "$step" > /dev/null && print_success "Done" || print_warning "Failed (continuing)"
                ;;
            shell)
                execute_shell_step "$step" > /dev/null && print_success "Done" || print_warning "Failed (continuing)"
                ;;
            *)
                print_warning "Skipping unknown action: $action"
                ;;
        esac
    done < <(echo "$teardown_steps" | jq -c '.[]')

    echo ""
    print_success "Teardown complete"

    # Optionally clean checkpoint
    echo ""
    echo "Checkpoint preserved. To remove:"
    echo "  $0 clean $tree_name"
}

# List available trees
cmd_list() {
    print_header "Available Trees"
    tree_list_detail

    echo ""
    echo "Checkpoints:"
    local checkpoints
    checkpoints=$(checkpoint_list)
    if [[ -n "$checkpoints" ]]; then
        while IFS= read -r tree; do
            local status
            status=$(checkpoint_status "$tree")
            printf "  %-20s %s\n" "$tree" "($status)"
        done <<< "$checkpoints"
    else
        echo "  (none)"
    fi
}

# Clean checkpoint
cmd_clean() {
    local tree_name="${1:-}"

    if [[ -z "$tree_name" ]]; then
        echo "Usage: $0 clean <tree>"
        exit 1
    fi

    if checkpoint_delete "$tree_name"; then
        print_success "Checkpoint deleted for tree: $tree_name"
    else
        echo "No checkpoint found for tree: $tree_name"
    fi
}

# ============================================================================
# Main dispatch
# ============================================================================

case "$COMMAND" in
    run)
        cmd_run "$@"
        ;;
    resume)
        cmd_resume "$@"
        ;;
    only)
        cmd_only "$@"
        ;;
    status)
        cmd_status "$@"
        ;;
    teardown)
        cmd_teardown "$@"
        ;;
    list)
        cmd_list
        ;;
    clean)
        cmd_clean "$@"
        ;;
    *)
        echo "Unknown command: $COMMAND"
        echo "Valid commands: run, resume, only, status, teardown, list, clean"
        exit 1
        ;;
esac