# Simple Workflow Example
# Both coordinator and workflow images are built from this Dockerfile and pushed to ECR.

tf_dir := "../terraform"

# =============================================================================
# Build and Push
# =============================================================================

# Build the workflow image
build-workflow:
    #!/usr/bin/env bash
    set -euo pipefail
    workflow_ecr=$(terraform -chdir={{tf_dir}} output -raw workflow_ecr_repository_url)
    docker build --target runtime -t "$workflow_ecr:latest" .

# Build the coordinator image (includes workflow files)
build-coordinator:
    #!/usr/bin/env bash
    set -euo pipefail
    coordinator_ecr=$(terraform -chdir={{tf_dir}} output -raw coordinator_ecr_repository_url)
    docker build --target coordinator -t "$coordinator_ecr:latest" .

# Build both images
build: build-workflow build-coordinator

# Login to ECR
ecr-login:
    #!/usr/bin/env bash
    set -euo pipefail
    region=$(terraform -chdir={{tf_dir}} output -raw region)
    workflow_ecr=$(terraform -chdir={{tf_dir}} output -raw workflow_ecr_repository_url)
    ecr_host="${workflow_ecr%%/*}"
    aws ecr get-login-password --region "$region" | docker login --username AWS --password-stdin "$ecr_host"

# Push the workflow image to ECR
push-workflow: ecr-login
    #!/usr/bin/env bash
    set -euo pipefail
    workflow_ecr=$(terraform -chdir={{tf_dir}} output -raw workflow_ecr_repository_url)
    docker push "$workflow_ecr:latest"

# Push the coordinator image to ECR
push-coordinator: ecr-login
    #!/usr/bin/env bash
    set -euo pipefail
    coordinator_ecr=$(terraform -chdir={{tf_dir}} output -raw coordinator_ecr_repository_url)
    docker push "$coordinator_ecr:latest"

# Push both images to ECR
push: push-workflow push-coordinator

# Build and push both images
build-push: build push

# =============================================================================
# Run and Monitor
# =============================================================================

# Directory containing this justfile
workflow_dir := justfile_directory()

# Submit coordinator job via snakemake plugin
run *args:
    #!/usr/bin/env bash
    set -euo pipefail
    region=$(terraform -chdir={{tf_dir}} output -raw region)
    coordinator_queue=$(terraform -chdir={{tf_dir}} output -raw coordinator_job_queue_name)
    workflow_queue=$(terraform -chdir={{tf_dir}} output -raw workflow_job_queue_name)
    bucket=$(terraform -chdir={{tf_dir}} output -raw bucket_name)
    workflow_job_def=$(terraform -chdir={{tf_dir}} output -raw workflow_job_definition_name)
    coordinator_job_def=$(terraform -chdir={{tf_dir}} output -raw coordinator_job_definition_name)
    snakemake \
        --snakefile Snakefile \
        --executor aws-basic-batch \
        --aws-basic-batch-coordinator true \
        --aws-basic-batch-region "$region" \
        --aws-basic-batch-job-queue "$workflow_queue" \
        --aws-basic-batch-job-definition "$workflow_job_def" \
        --aws-basic-batch-coordinator-job-definition "$coordinator_job_def" \
        --aws-basic-batch-coordinator-queue "$coordinator_queue" \
        --default-storage-provider s3 \
        --default-storage-prefix "s3://$bucket" \
        --jobs 10 \
        --forceall \
        {{args}} 2>&1 | tee /dev/stderr | grep -oP 'Coordinator job submitted: \K[a-f0-9-]+' > "{{workflow_dir}}/.last-job-id" || true

# Check status of last submitted job
status:
    #!/usr/bin/env bash
    set -euo pipefail
    job_id=$(cat "{{workflow_dir}}/.last-job-id" 2>/dev/null || echo "")
    if [ -z "$job_id" ]; then
        echo "No job ID found. Run 'just run' first."
        exit 1
    fi
    aws batch describe-jobs --jobs "$job_id" \
        --query 'jobs[0].{status:status,statusReason:statusReason,logStreamName:container.logStreamName}' \
        --output table

# Get logs for last submitted job
logs:
    #!/usr/bin/env bash
    set -euo pipefail
    job_id=$(cat "{{workflow_dir}}/.last-job-id" 2>/dev/null || echo "")
    if [ -z "$job_id" ]; then
        echo "No job ID found. Run 'just run' first."
        exit 1
    fi
    log_stream=$(aws batch describe-jobs --jobs "$job_id" \
        --query 'jobs[0].container.logStreamName' --output text)
    if [ "$log_stream" = "None" ] || [ -z "$log_stream" ]; then
        echo "No log stream available yet. Job may still be starting."
        exit 1
    fi
    log_group=$(terraform -chdir={{tf_dir}} output -raw log_group_name)
    aws logs get-log-events \
        --log-group-name "$log_group" \
        --log-stream-name "$log_stream" \
        --query 'events[*].message' --output text

# Watch job status until completion
watch:
    #!/usr/bin/env bash
    set -euo pipefail
    job_id=$(cat "{{workflow_dir}}/.last-job-id" 2>/dev/null || echo "")
    if [ -z "$job_id" ]; then
        echo "No job ID found. Run 'just run' first."
        exit 1
    fi
    while true; do
        status=$(aws batch describe-jobs --jobs "$job_id" \
            --query 'jobs[0].status' --output text)
        echo "$(date '+%H:%M:%S') - Status: $status"
        if [ "$status" = "SUCCEEDED" ] || [ "$status" = "FAILED" ]; then
            break
        fi
        sleep 10
    done
    just logs

# =============================================================================
# Infrastructure
# =============================================================================

# Initialize terraform
tf-init:
    terraform -chdir={{tf_dir}} init

# Plan infrastructure
tf-plan *args:
    terraform -chdir={{tf_dir}} plan {{args}}

# Apply infrastructure
tf-apply *args:
    terraform -chdir={{tf_dir}} apply {{args}}

# Destroy infrastructure
tf-destroy *args:
    terraform -chdir={{tf_dir}} destroy {{args}}

# Show terraform outputs
tf-output *args:
    terraform -chdir={{tf_dir}} output {{args}}

# Generate snakemake command using deployed infrastructure
tf-snakemake-cmd:
    #!/usr/bin/env bash
    set -euo pipefail
    region=$(terraform -chdir={{tf_dir}} output -raw region)
    coordinator_queue=$(terraform -chdir={{tf_dir}} output -raw coordinator_job_queue_name)
    workflow_queue=$(terraform -chdir={{tf_dir}} output -raw workflow_job_queue_name)
    bucket=$(terraform -chdir={{tf_dir}} output -raw bucket_name)
    workflow_job_def=$(terraform -chdir={{tf_dir}} output -raw workflow_job_definition_name)
    coordinator_job_def=$(terraform -chdir={{tf_dir}} output -raw coordinator_job_definition_name)

    echo "snakemake --executor aws-basic-batch \\"
    echo "    --aws-basic-batch-region $region \\"
    echo "    --aws-basic-batch-job-queue $workflow_queue \\"
    echo "    --aws-basic-batch-job-definition $workflow_job_def \\"
    echo "    --aws-basic-batch-coordinator true \\"
    echo "    --aws-basic-batch-coordinator-job-definition $coordinator_job_def \\"
    echo "    --aws-basic-batch-coordinator-queue $coordinator_queue \\"
    echo "    --default-storage-provider s3 \\"
    echo "    --default-storage-prefix s3://$bucket"
