#!/usr/bin/env bash
# Copyright (c) 2025 Tylt LLC. All rights reserved.
# Derivative works may be released by researchers,
# but original files may not be redistributed or used beyond research purposes.
#
# DO NOT EDIT THIS FILE - Generated by cudag framework

# Usage:
#   ./scripts/generate.sh [options]           # Generate and auto-upload
#   ./scripts/generate.sh --dry [options]     # Generate only, no upload
#   ./scripts/generate.sh --verify            # Generate with interactive verification
#   ./scripts/generate.sh --verify --verbose  # Verify with streaming output

set -euo pipefail

DRY_RUN=false
VERIFY_MODE=false
VERBOSE=false
EXTRA_ARGS=()

# Parse args - extract --dry, --verify, --verbose, pass everything else through
for arg in "$@"; do
    if [[ "$arg" == "--dry" ]]; then
        DRY_RUN=true
    elif [[ "$arg" == "--verify" ]]; then
        VERIFY_MODE=true
        DRY_RUN=true  # verify implies dry run
    elif [[ "$arg" == "--verbose" ]] || [[ "$arg" == "-v" ]]; then
        VERBOSE=true
    else
        EXTRA_ARGS+=("$arg")
    fi
done

# If verify mode, delegate to verify.py which handles the loop
if [[ "$VERIFY_MODE" == "true" ]]; then
    echo "========================================"
    echo "Interactive Dataset Verification"
    echo "========================================"
    echo ""

    # Use prod config by default
    CONFIG_PATH="config/dataset.prod.yaml"
    for i in "${!EXTRA_ARGS[@]}"; do
        if [[ "${EXTRA_ARGS[$i]}" == "--config" ]] && [[ $((i+1)) -lt ${#EXTRA_ARGS[@]} ]]; then
            CONFIG_PATH="${EXTRA_ARGS[$((i+1))]}"
            break
        fi
    done

    VERIFY_ARGS=(--config "$CONFIG_PATH")
    if [[ "$VERBOSE" == "true" ]]; then
        VERIFY_ARGS+=(--verbose)
    fi

    exec uv run python scripts/verify.py "${VERIFY_ARGS[@]}"
fi

echo "========================================"
echo "STAGE 1: Generate Dataset"
echo "========================================"
echo ""

if [[ "$DRY_RUN" == "true" ]]; then
    echo "[DRY RUN] Will generate but NOT upload"
    echo ""
fi

# Run the dataset generation using uv run (uses pyproject.toml dependencies)
# Set env var so generator.py knows it was called from this script
export CUDAG_FROM_SCRIPT=1

if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then
    uv run python generator.py "${EXTRA_ARGS[@]}"
else
    uv run python generator.py
fi

if [[ $? -ne 0 ]]; then
    echo ""
    echo "Dataset generation failed"
    exit 1
fi

# Find the most recently created dataset directory
LATEST_DATASET=$(ls -td datasets/*/ 2>/dev/null | head -1)

if [[ -z "$LATEST_DATASET" ]]; then
    echo ""
    echo "No dataset directory found"
    exit 1
fi

DATASET_NAME=$(basename "$LATEST_DATASET")
echo ""
echo "Generated dataset: $DATASET_NAME"

if [[ "$DRY_RUN" == "true" ]]; then
    echo ""
    echo "[DRY RUN] Skipping upload"
    echo ""
    echo "To verify interactively:"
    echo "  ./scripts/generate.sh --verify"
    echo ""
    echo "To upload manually:"
    echo "  ./scripts/upload.sh datasets/$DATASET_NAME"
    exit 0
fi

echo ""
echo "========================================"
echo "Auto-starting upload..."
echo "========================================"
echo ""

exec ./scripts/upload.sh "datasets/$DATASET_NAME"
