#!/bin/bash
#
# E2E Test for ControlPlan v0.1 Bottle Line Scenario
#
# This script runs the full deterministic flow:
# 1. (Optional) Generate configuration.json via LLM
# 2. Build scenario with control plan
# 3. Validate generated logic
# 4. Print command to start ICS-SimLab (does not execute)
#
# Usage:
#   ./scripts/e2e_bottle_control_plan.sh [OPTIONS]
#
# Options:
#   --skip-llm    Skip LLM generation, use existing outputs/configuration.json
#   --use-config PATH  Use a specific configuration.json (implies --skip-llm)
#   --help        Show this help message
#
# Prerequisites:
#   - Python virtual environment at .venv/
#   - OPENAI_API_KEY set (unless --skip-llm)
#

set -e

# Configuration
REPO_DIR="$(cd "$(dirname "$0")/.." && pwd)"
VENV_DIR="$REPO_DIR/.venv"
CONTROL_PLAN="$REPO_DIR/examples/control_plans/bottle_line_v0.1.json"
PROMPT_FILE="$REPO_DIR/prompts/e2e_bottle.txt"
OUT_DIR="$REPO_DIR/outputs/scenario_bottle_cp"
CONFIG_OUT="$REPO_DIR/outputs/configuration.json"

# Parse arguments
SKIP_LLM=false
USE_CONFIG=""

while [[ $# -gt 0 ]]; do
    case $1 in
        --skip-llm)
            SKIP_LLM=true
            shift
            ;;
        --use-config)
            USE_CONFIG="$2"
            SKIP_LLM=true
            shift 2
            ;;
        --help)
            head -30 "$0" | grep "^#" | sed 's/^# //'
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
log_step() { echo -e "\n${BLUE}===== $1 =====${NC}"; }

# ==============================================================================
# Step 0: Verify prerequisites
# ==============================================================================
log_step "Step 0: Verifying prerequisites"

if [ ! -d "$VENV_DIR" ]; then
    log_error "Virtual environment not found: $VENV_DIR"
    log_info "Create it with: python3 -m venv .venv && .venv/bin/pip install -r requirements.txt"
    exit 1
fi

if [ ! -f "$CONTROL_PLAN" ]; then
    log_error "Control plan not found: $CONTROL_PLAN"
    exit 1
fi

if [ ! -f "$PROMPT_FILE" ]; then
    log_error "Prompt file not found: $PROMPT_FILE"
    exit 1
fi

# Activate virtual environment
source "$VENV_DIR/bin/activate"
log_info "Activated virtual environment: $VENV_DIR"

# ==============================================================================
# Step 1: Generate configuration.json (optional)
# ==============================================================================
log_step "Step 1: Generate configuration.json"

if [ -n "$USE_CONFIG" ]; then
    if [ ! -f "$USE_CONFIG" ]; then
        log_error "Config file not found: $USE_CONFIG"
        exit 1
    fi
    log_info "Using provided config: $USE_CONFIG"
    mkdir -p "$(dirname "$CONFIG_OUT")"
    cp "$USE_CONFIG" "$CONFIG_OUT"
elif [ "$SKIP_LLM" = true ]; then
    if [ ! -f "$CONFIG_OUT" ]; then
        log_error "No configuration.json found and --skip-llm specified"
        log_info "Either run without --skip-llm, or provide --use-config PATH"
        exit 1
    fi
    log_info "Skipping LLM generation, using existing: $CONFIG_OUT"
else
    if [ -z "$OPENAI_API_KEY" ]; then
        log_error "OPENAI_API_KEY not set"
        log_info "Set it with: export OPENAI_API_KEY='...'"
        log_info "Or use --skip-llm to skip LLM generation"
        exit 1
    fi
    log_info "Generating configuration via LLM..."
    python3 "$REPO_DIR/main.py" \
        --input-file "$PROMPT_FILE" \
        --out "$CONFIG_OUT"

    if [ ! -f "$CONFIG_OUT" ]; then
        log_error "LLM generation failed: $CONFIG_OUT not created"
        exit 1
    fi
    log_info "Generated: $CONFIG_OUT"
fi

# ==============================================================================
# Step 2: Build scenario with control plan
# ==============================================================================
log_step "Step 2: Build scenario with control plan"

log_info "Control plan: $CONTROL_PLAN"
log_info "Output dir: $OUT_DIR"

python3 "$REPO_DIR/build_scenario.py" \
    --config "$CONFIG_OUT" \
    --out "$OUT_DIR" \
    --control-plan "$CONTROL_PLAN" \
    --overwrite

if [ ! -f "$OUT_DIR/configuration.json" ]; then
    log_error "Build failed: $OUT_DIR/configuration.json not created"
    exit 1
fi

log_info "Scenario built: $OUT_DIR"

# ==============================================================================
# Step 3: Validate generated logic
# ==============================================================================
log_step "Step 3: Validate generated logic"

python3 -m tools.validate_logic \
    --config "$OUT_DIR/configuration.json" \
    --logic-dir "$OUT_DIR/logic" \
    --check-callbacks \
    --check-hil-init

log_info "Logic validation passed"

# ==============================================================================
# Step 4: Verify control plan compiled files exist
# ==============================================================================
log_step "Step 4: Verify control plan compiled files"

# Check that HIL files from control plan exist
for hil_name in water_hil filler_hil; do
    hil_file="$OUT_DIR/logic/${hil_name}.py"
    if [ -f "$hil_file" ]; then
        log_info "Found: $hil_file"
        # Verify it contains logic(physical_values)
        if grep -q "def logic(physical_values)" "$hil_file"; then
            log_info "  Contains logic(physical_values): OK"
        else
            log_warn "  Missing logic(physical_values) signature"
        fi
    else
        log_warn "HIL file not found: $hil_file (may have different name in config)"
    fi
done

# ==============================================================================
# Summary
# ==============================================================================
log_step "SUCCESS: Scenario ready"

echo ""
echo "Scenario contents:"
echo "  Configuration: $OUT_DIR/configuration.json"
echo "  Logic files:"
for f in "$OUT_DIR/logic"/*.py; do
    echo "    $(basename "$f")"
done

echo ""
echo "To run with ICS-SimLab (requires Docker and ICS-SimLab repo):"
echo "  cd ~/projects/ICS-SimLab-main/curtin-ics-simlab"
echo "  sudo ./start.sh $(realpath "$OUT_DIR")"

exit 0