#!/usr/bin/env bash

# earlGreyParTEA - Pangenome TE Analysis Pipeline (Full Mode)
# Version 0.1.3
# A Snakemake-based multi-genome transposable element annotation pipeline

set -e

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

usage() {
    cat << EOF
#############################
earlGreyParTEA version ${VERSION}
Pangenome TE Analysis - Full Pipeline Mode

This runs the complete pipeline: library construction + annotation

Required Parameters:
    -c, --config FILE       Path to config.yaml file
    -t, --threads INT       Number of threads to use

Optional Parameters:
    -m, --memory INT        Scheduling memory budget in MB (default: no limit). Controls how many
                            jobs run simultaneously — Snakemake will not start a new job if doing
                            so would exceed this budget. Does NOT cap per-process memory at OS
                            level. If the pipeline crashes with OOM, lower this value to reduce
                            concurrency rather than raising it.
    -n, --dry-run           Perform dry run (show what would be executed)
    --generate-config FILE  Generate an example config.yaml template
    --genome-dir DIR        Auto-populate genomes from a directory (use with --generate-config)
    --from-csv FILE         Populate genomes from a CSV file (use with --generate-config)
    --output-dir DIR        Set output_dir in the generated config (use with --generate-config)
    --unlock                Unlock working directory after crash
    --rerun-incomplete      Rerun incomplete jobs from previous run
    --snakemake-args "FLAGS"  Pass additional arguments directly to snakemake (e.g. "--until rule_name")
    --slurm                 Submit jobs to a SLURM cluster via snakemake-executor-plugin-slurm
    --slurm-jobs N          Max concurrent SLURM jobs (default: genomes × 3, capped at 200)
    --slurm-partition PART  SLURM partition/queue to submit to (required with --slurm)
    --slurm-account ACCT    SLURM account string (optional)
    --slurm-extra "FLAGS"   Additional raw sbatch flags (e.g. "--constraint=avx2")
    -h, --help              Show this help message

Example Usage:
    # Generate example config (blank template)
    earlGreyParTEA --generate-config my_config.yaml

    # Generate config from a directory of genomes
    earlGreyParTEA --generate-config my_config.yaml --genome-dir /path/to/genomes/ --output-dir /path/to/results/

    # Generate config from a CSV file (columns: species, genome_path)
    earlGreyParTEA --generate-config my_config.yaml --from-csv genomes.csv --output-dir /path/to/results/

    # Run full pipeline
    earlGreyParTEA -c my_config.yaml -t 16
    
    # Run with memory limit
    earlGreyParTEA -c my_config.yaml -t 32 -m 128000
    
    # Dry run to check what will execute
    earlGreyParTEA -c my_config.yaml -t 16 --dry-run

Documentation:
    https://github.com/TobyBaril/EarlGreyParTEA

Questions/Issues:
    tobias.baril[at]unine.ch
    https://github.com/TobyBaril/EarlGreyParTEA/issues
#############################
EOF
}

generate_example_config() {
    local output_file="$1"
    cat > "$output_file" << 'EOF'
# EarlGrey Pangenome Pipeline Configuration
# Full Mode: Complete library construction and annotation

# Input genomes
genome:
  genome1: /path/to/genome1.fasta
  genome2: /path/to/genome2.fasta
  genome3: /path/to/genome3.fasta

# Species identifiers (must match keys in 'genome' above)
species:
  - genome1
  - genome2
  - genome3

# Output directory for all results
output_dir: /path/to/output/directory

# Pipeline mode (do not change for earlGreyParTEA)
pipeline_mode: "full"

# Initial masking with known repeats (choose ONE or leave both empty)
repeatmasker_species: ""  # e.g., "fungi", "arthropoda", "viridiplantae"
custom_library: ""        # path to custom TE library in fasta format

# Library construction parameters
iterations: 10           # Number of BLAST-extend-align cycles
flank: 1000             # Flanking basepairs to extract
max_consensus_seqs: 20  # Max sequences for consensus building
min_consensus_seqs: 3   # Min sequences required for consensus

# Clustering options (for combining TE libraries from multiple genomes)
skip_clustering: false          # Set to true to skip clustering (just concatenate)
clustering_identity: 0.8        # cd-hit sequence identity threshold (0.0-1.0)
clustering_coverage: 0.8        # cd-hit alignment coverage (0.0-1.0)

# Output options
softmask: false         # Generate softmasked genome for each input
margin: false           # Remove short TE sequences (<100bp)
run_heliano: true       # Run HELIANO for Helitron detection

# Workflow visualization (requires graphviz/dot installed)
generate_dag: true      # Generate workflow DAG visualizations
dag_format: "svg"       # Options: "svg", "png", "pdf"

# Saturation analysis options
saturation_permutations: 100  # Number of random genome-addition permutations to average
                              # over for the TE family saturation plot. Higher values
                              # give smoother confidence intervals but increase runtime.

# ---------------------------------------------------------------------------
# Optional analysis modules (v0.1.6+)
# ---------------------------------------------------------------------------

# Shared/unique TE content analysis (full and annotate modes only)
run_shared_unique: false

# BUSCO-based phylogenomics
run_busco_phylo: false
busco_lineage: ""        # REQUIRED if run_busco_phylo: true  (e.g. "fungi_odb10")
busco_prefix: "busco"   # Prefix for BUSCO run directory names
busco_min_occupancy: 0.5  # Min fraction of species a gene must appear in (0.0-1.0)

# Advanced options (usually not needed)
# script_dir: "/path/to/earlgrey/scripts"  # Auto-detected if installed via conda/mamba
EOF
    echo "[INFO] Example config file generated: $output_file"
    echo "[INFO] Edit this file with your genome paths and parameters before running"
    exit 0
}

# Parse arguments
CONFIG=""
THREADS=""
MEMORY=""
DRY_RUN=""
UNLOCK=""
RERUN=""
EXTRA_SNAKEMAKE_ARGS=""
GENERATE_CONFIG=""
GENOME_DIR=""
FROM_CSV=""
OUTPUT_DIR=""
SLURM=false
SLURM_JOBS=""
SLURM_PARTITION=""
SLURM_ACCOUNT=""
SLURM_EXTRA=""

while [[ $# -gt 0 ]]; do
    case $1 in
        -c|--config)
            CONFIG="$2"
            shift 2
            ;;
        -t|--threads)
            THREADS="$2"
            shift 2
            ;;
        -m|--memory)
            MEMORY="$2"
            shift 2
            ;;
        -n|--dry-run)
            DRY_RUN="--dry-run"
            shift
            ;;
        --generate-config)
            GENERATE_CONFIG="$2"
            shift 2
            ;;
        --genome-dir)
            GENOME_DIR="$2"
            shift 2
            ;;
        --from-csv)
            FROM_CSV="$2"
            shift 2
            ;;
        --output-dir)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        --unlock)
            UNLOCK="--unlock"
            shift
            ;;
        --rerun-incomplete)
            RERUN="--rerun-incomplete"
            shift
            ;;
        --snakemake-args)
            EXTRA_SNAKEMAKE_ARGS="$2"
            shift 2
            ;;
        --slurm)
            SLURM=true
            shift
            ;;
        --slurm-jobs)
            SLURM_JOBS="$2"
            shift 2
            ;;
        --slurm-partition)
            SLURM_PARTITION="$2"
            shift 2
            ;;
        --slurm-account)
            SLURM_ACCOUNT="$2"
            shift 2
            ;;
        --slurm-extra)
            SLURM_EXTRA="$2"
            shift 2
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "[ERROR] Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Handle --generate-config (with optional --genome-dir / --from-csv / --output-dir)
if [ -n "$GENERATE_CONFIG" ]; then
    if [ -n "$GENOME_DIR" ] && [ -n "$FROM_CSV" ]; then
        echo "[ERROR] --genome-dir and --from-csv are mutually exclusive"
        exit 1
    fi
    if [ -n "$GENOME_DIR" ] || [ -n "$FROM_CSV" ]; then
        GEN_CONFIG_PY=""
        GEN_POSSIBLE=()
        if [ -n "$CONDA_PREFIX" ] && [ -d "$CONDA_PREFIX/share" ]; then
            for partea_dir in "$CONDA_PREFIX/share/earlgrey-partea"-*; do
                if [ -d "$partea_dir" ]; then
                    GEN_POSSIBLE+=("$partea_dir/scripts/generate_config.py")
                fi
            done
        fi
        GEN_POSSIBLE+=(
            "$SCRIPT_DIR/../share/earlgrey-partea/scripts/generate_config.py"
            "$SCRIPT_DIR/../scripts/generate_config.py"
            "$SCRIPT_DIR/scripts/generate_config.py"
            "$(dirname "$SCRIPT_DIR")/scripts/generate_config.py"
        )
        for loc in "${GEN_POSSIBLE[@]}"; do
            if [ -f "$loc" ]; then
                GEN_CONFIG_PY="$loc"
                break
            fi
        done
        if [ -z "$GEN_CONFIG_PY" ]; then
            echo "[ERROR] Could not locate generate_config.py. Please ensure EarlGrey is properly installed."
            exit 1
        fi
        CMD_ARGS=(python "$GEN_CONFIG_PY" --output "$GENERATE_CONFIG" --mode full)
        [ -n "$GENOME_DIR" ] && CMD_ARGS+=(--genome-dir "$GENOME_DIR")
        [ -n "$FROM_CSV" ]   && CMD_ARGS+=(--from-csv   "$FROM_CSV")
        [ -n "$OUTPUT_DIR" ] && CMD_ARGS+=(--output-dir "$OUTPUT_DIR")
        "${CMD_ARGS[@]}"
        exit 0
    else
        generate_example_config "$GENERATE_CONFIG"
    fi
fi

# Validate required arguments
if [ -z "$CONFIG" ]; then
    echo "[ERROR] Config file required (-c/--config)"
    echo ""
    usage
    exit 1
fi

if [ -z "$THREADS" ]; then
    echo "[ERROR] Number of threads required (-t/--threads)"
    echo ""
    usage
    exit 1
fi

if [ ! -f "$CONFIG" ]; then
    echo "[ERROR] Config file not found: $CONFIG"
    exit 1
fi

# Find Snakefile (check multiple possible locations, version-agnostic)
SNAKEFILE=""
POSSIBLE_LOCATIONS=()

# Check for conda package installation
if [ -n "$CONDA_PREFIX" ] && [ -d "$CONDA_PREFIX/share" ]; then
    # Find all earlgrey-partea-* directories (any version)
    for partea_dir in "$CONDA_PREFIX/share/earlgrey-partea"-*; do
        if [ -d "$partea_dir" ]; then
            POSSIBLE_LOCATIONS+=("$partea_dir/Snakefile")
        fi
    done
fi

# Check relative to script location (for development or custom installs)
POSSIBLE_LOCATIONS+=(
    "$SCRIPT_DIR/../share/earlgrey-partea/Snakefile"
    "$SCRIPT_DIR/../Snakefile"
    "$SCRIPT_DIR/Snakefile"
    "$(dirname "$SCRIPT_DIR")/Snakefile"
)

for location in "${POSSIBLE_LOCATIONS[@]}"; do
    if [ -f "$location" ]; then
        SNAKEFILE="$location"
        break
    fi
done

if [ -z "$SNAKEFILE" ]; then
    echo "[ERROR] Could not locate Snakefile. Please ensure EarlGrey is properly installed."
    echo "[ERROR] Searched locations:"
    for loc in "${POSSIBLE_LOCATIONS[@]}"; do
        echo "        $loc"
    done
    exit 1
fi

echo "[INFO] Using Snakefile: $SNAKEFILE"
echo "[INFO] Using config: $CONFIG"
echo "[INFO] Threads: $THREADS"

# Build snakemake command
if $SLURM; then
    # --- SLURM mode ---
    # Fall back to SLURM settings from config file if not set via CLI
    if [ -z "$SLURM_PARTITION" ]; then
        SLURM_PARTITION=$(python3 -c "import yaml; cfg=yaml.safe_load(open('$CONFIG')); print(cfg.get('slurm_partition', ''))" 2>/dev/null || echo "")
    fi
    if [ -z "$SLURM_ACCOUNT" ]; then
        SLURM_ACCOUNT=$(python3 -c "import yaml; cfg=yaml.safe_load(open('$CONFIG')); print(cfg.get('slurm_account', ''))" 2>/dev/null || echo "")
    fi
    if [ -z "$SLURM_EXTRA" ]; then
        SLURM_EXTRA=$(python3 -c "import yaml; cfg=yaml.safe_load(open('$CONFIG')); print(cfg.get('slurm_extra', ''))" 2>/dev/null || echo "")
    fi
    if [ -z "$SLURM_PARTITION" ]; then
        echo "[ERROR] --slurm-partition is required when using --slurm (or set 'slurm_partition:' in the config file)"
        exit 1
    fi
    # Default: genomes × 3, capped at 200
    if [ -z "$SLURM_JOBS" ]; then
        N_GENOMES=$(python3 -c "import yaml; cfg=yaml.safe_load(open('$CONFIG')); print(len(cfg.get('species', cfg.get('genome', {}))))" 2>/dev/null || echo 1)
        SLURM_JOBS=$(( N_GENOMES * 3 ))
        [ "$SLURM_JOBS" -gt 200 ] && SLURM_JOBS=200
    fi
    # Build default-resources as an array — Snakemake 9 requires each key=value as a
    # separate argument; passing them as a single quoted string passes the literal " char.
    DEFAULT_RESOURCES=("slurm_partition=$SLURM_PARTITION")
    [ -n "$SLURM_ACCOUNT" ] && DEFAULT_RESOURCES+=("slurm_account=$SLURM_ACCOUNT")
    [ -n "$SLURM_EXTRA" ]   && DEFAULT_RESOURCES+=("slurm_extra=$SLURM_EXTRA")
    echo "[INFO] SLURM mode: partition=$SLURM_PARTITION jobs=$SLURM_JOBS"
    SNAKEMAKE_ARGS=(
        snakemake
        --snakefile "$SNAKEFILE"
        --configfile "$CONFIG"
        --config slurm_mode=true
        --executor slurm
        --default-resources "${DEFAULT_RESOURCES[@]}"
        --jobs "$SLURM_JOBS"
        --cores "$THREADS"
        --latency-wait 60
        --retries 1
    )
else
    # --- Local mode ---
    SNAKEMAKE_ARGS=(snakemake --snakefile "$SNAKEFILE" --configfile "$CONFIG" --cores "$THREADS" --default-resources "tmpdir='/tmp'")
    if [ -n "$MEMORY" ]; then
        if [ "$MEMORY" -lt 32000 ]; then
            echo "[WARNING] -m ${MEMORY}MB is below the recommended minimum of 32000MB. TEstrainer"
            echo "          requires approximately 32 GB per job and may crash with less. Consider"
            echo "          increasing -m or running on a machine with more RAM."
        fi
        SNAKEMAKE_ARGS+=(--resources "mem_mb=$MEMORY")
        SNAKEMAKE_ARGS+=(--config "total_memory_mb=$MEMORY")
        echo "[INFO] Memory limit: ${MEMORY}MB"
    fi
fi

if [ -n "$DRY_RUN" ]; then
    SNAKEMAKE_ARGS+=("$DRY_RUN")
    echo "[INFO] Performing dry run"
fi

if [ -n "$UNLOCK" ]; then
    SNAKEMAKE_ARGS+=("$UNLOCK")
    echo "[INFO] Unlocking working directory"
fi

if [ -n "$RERUN" ]; then
    SNAKEMAKE_ARGS+=("$RERUN")
    echo "[INFO] Rerunning incomplete jobs"
fi

if [ -n "$EXTRA_SNAKEMAKE_ARGS" ]; then
    # Split on whitespace into array elements
    read -ra _EXTRA <<< "$EXTRA_SNAKEMAKE_ARGS"
    SNAKEMAKE_ARGS+=("${_EXTRA[@]}")
    echo "[INFO] Extra snakemake args: $EXTRA_SNAKEMAKE_ARGS"
fi

# Execute
echo ""
echo "Executing: ${SNAKEMAKE_ARGS[*]}"
echo ""

exec "${SNAKEMAKE_ARGS[@]}"
