#!/bin/bash
# Fix ML test compilation errors

echo "=== Fixing ML Test Compilation Errors ==="

# Find all test modules and add common missing imports
find ml/src -name "*.rs" -type f | while read file; do
    # Check if file has tests
    if grep -q "#\[cfg(test)\]" "$file" 2>/dev/null; then
        echo "Processing: $file"

        # Check if imports section exists
        if ! grep -q "^#\[cfg(test)\]" "$file"; then
            continue
        fi

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

        if [ ! -z "$test_line" ]; then
            # Check what's missing and add imports after the mod tests { line
            mod_line=$((test_line + 1))

            # Create a temporary file with the fixes
            awk -v modline="$mod_line" '
            NR == modline && /^mod tests \{/ {
                print $0
                print "    use candle_core::{Device, DType};"
                print "    use std::fs::File;"
                print "    use std::io::Write;"
                print "    use tempfile::tempdir;"
                next
            }
            {print}
            ' "$file" > "$file.tmp" && mv "$file.tmp" "$file"
        fi
    fi
done

echo "=== Done ==="