#!/usr/bin/env bash
# build_all.sh — Generate synthetic data and train all 4 TFLite-Micro attack detection models
# Run from the models/ directory or provide MODELS_DIR env var
# Target: DGX (spark-b0c1) with TensorFlow 2.x, models deploy to ESP32-S3 via SPIFFS

set -euo pipefail

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

PYTHON="${PYTHON:-python3}"

echo "============================================================"
echo "  TFLite-Micro Attack Detection — Full Build Pipeline"
echo "  $(date)"
echo "============================================================"
echo ""
echo "Python: $($PYTHON --version 2>&1)"
echo "TensorFlow: $($PYTHON -c 'import tensorflow as tf; print(tf.__version__)' 2>&1)"
echo ""

MODELS=(
    "rf_attack_classifier"
    "mavlink_anomaly"
    "lora_mesh_attack"
    "protocol_anomaly"
)

FAILED=()

for model in "${MODELS[@]}"; do
    echo ""
    echo "============================================================"
    echo "  Building: $model"
    echo "============================================================"

    model_dir="$SCRIPT_DIR/$model"

    if [ ! -d "$model_dir" ]; then
        echo "ERROR: Directory not found: $model_dir"
        FAILED+=("$model")
        continue
    fi

    echo ""
    echo "--- Step 1: Generating synthetic data ---"
    if $PYTHON "$model_dir/generate_synthetic_data.py" --output-dir "$model_dir"; then
        echo "  Data generation: OK"
    else
        echo "  Data generation: FAILED"
        FAILED+=("$model")
        continue
    fi

    echo ""
    echo "--- Step 2: Training model ---"
    if $PYTHON "$model_dir/train.py"; then
        echo "  Training: OK"
    else
        echo "  Training: FAILED"
        FAILED+=("$model")
        continue
    fi

    echo ""
    echo "--- $model complete ---"
done

echo ""
echo "============================================================"
echo "  Build Summary"
echo "============================================================"
echo ""

# List generated artifacts
for model in "${MODELS[@]}"; do
    model_dir="$SCRIPT_DIR/$model"
    tflite_file="$model_dir/$model.tflite"
    if [ -f "$tflite_file" ]; then
        size=$(du -h "$tflite_file" | cut -f1)
        echo "  [OK]   $model.tflite — $size"
    else
        echo "  [MISS] $model.tflite — not generated"
    fi
done

echo ""

if [ ${#FAILED[@]} -eq 0 ]; then
    echo "All models built successfully."
    echo ""
    echo "Deployment: Copy .tflite files to ESP32-S3 SPIFFS partition."
    echo "  Each model also has scaler_params.npz for input normalization."
    exit 0
else
    echo "FAILED models: ${FAILED[*]}"
    exit 1
fi
