#!/bin/bash
# Apply systematic fixes to ML test modules
# This script adds common test imports to test modules

set -e

echo "=== Applying ML Test Compilation Fixes ==="
echo "This script will add common imports to test modules"
echo ""

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

# Counter
fixes_applied=0

# Function to add imports to a test module
fix_test_module() {
    local file=$1
    local needs_device=false
    local needs_file=false
    local needs_tempdir=false

    # Check what imports are needed
    if grep -q "Device" "$file" && ! grep -q "use candle_core::Device" "$file"; then
        needs_device=true
    fi

    if grep -q "File\|Write\|Read" "$file" && ! grep -q "use std::fs::File" "$file"; then
        needs_file=true
    fi

    if grep -q "tempdir" "$file" && ! grep -q "use tempfile::tempdir" "$file"; then
        needs_tempdir=true
    fi

    # If any imports needed, apply fix
    if [ "$needs_device" = true ] || [ "$needs_file" = true ] || [ "$needs_tempdir" = true ]; then
        echo -e "${YELLOW}Fixing: $file${NC}"

        # Find the test module line
        test_mod_line=$(grep -n "^#\[cfg(test)\]" "$file" | head -1 | cut -d: -f1)

        if [ ! -z "$test_mod_line" ]; then
            # Calculate line after "mod tests {"
            mod_tests_line=$((test_mod_line + 1))

            # Build import string
            imports="    use super::*;"

            if [ "$needs_device" = true ]; then
                imports="$imports\n    use candle_core::{Device, DType};"
            fi

            if [ "$needs_file" = true ]; then
                imports="$imports\n    use std::fs::File;\n    use std::io::Write;"
            fi

            if [ "$needs_tempdir" = true ]; then
                imports="$imports\n    use tempfile::tempdir;"
            fi

            # Apply fix (insert after mod tests { line)
            awk -v modline="$mod_tests_line" -v imports="$imports" '
                NR == modline && /^mod tests \{/ {
                    print $0
                    print imports
                    next
                }
                {print}
            ' "$file" > "$file.tmp" && mv "$file.tmp" "$file"

            fixes_applied=$((fixes_applied + 1))
            echo -e "${GREEN}  ✓ Fixed${NC}"
        fi
    fi
}

# Find all Rust files with test modules
echo "Scanning for test modules..."
test_files=$(find ml/src -name "*.rs" -type f -exec grep -l "#\[cfg(test)\]" {} \;)

for file in $test_files; do
    fix_test_module "$file"
done

echo ""
echo -e "${GREEN}=== Summary ===${NC}"
echo "Files processed: $(echo "$test_files" | wc -l)"
echo "Fixes applied: $fixes_applied"
echo ""
echo "Next steps:"
echo "1. Run: cargo check -p ml --lib"
echo "2. Check for remaining errors"
echo "3. Apply additional fixes as needed"