#!/usr/bin/env python3
"""
Flash uniform firmware to all ESP32 devices, then reconfigure via serial
This is much faster than building unique firmwares for each device
"""

import os
import sys
import subprocess
import argparse
import glob
from pathlib import Path
import time

def detect_devices():
    """Detect all ESP32 devices"""
    devices = sorted(glob.glob('/dev/ttyUSB*'))
    return devices

def build_firmware(project_dir):
    """Build the firmware once"""
    print("=" * 60)
    print("Building firmware (one time)...")
    print("=" * 60)
    
    result = subprocess.run(
        ['idf.py', 'build'],
        cwd=project_dir,
        capture_output=False
    )
    
    if result.returncode != 0:
        print("✗ Build failed!")
        return False
    
    print("✓ Build complete")
    return True

def flash_device(port, project_dir):
    """Flash a single device"""
    try:
        result = subprocess.run(
            ['idf.py', '-p', port, 'flash'],
            cwd=project_dir,
            capture_output=True,
            text=True,
            timeout=120
        )
        return result.returncode == 0
    except Exception as e:
        return False

def flash_all_devices(devices, project_dir):
    """Flash all devices with the same firmware"""
    print(f"\nFlashing {len(devices)} devices...")
    
    success_count = 0
    for idx, dev in enumerate(devices, 1):
        print(f"[{idx:2d}/{len(devices)}] Flashing {dev}...", end='', flush=True)
        
        if flash_device(dev, project_dir):
            print(" ✓")
            success_count += 1
        else:
            print(" ✗")
    
    print(f"\nFlashed: {success_count}/{len(devices)}")
    return success_count

def reconfigure_devices(ssid, password, start_ip, gateway="192.168.1.1"):
    """Reconfigure devices using the reconfig script"""
    script_path = os.path.join(os.path.dirname(__file__), 'reconfig_simple.py')
    
    if not os.path.exists(script_path):
        print(f"Error: {script_path} not found!")
        return False
    
    print("\n" + "=" * 60)
    print("Reconfiguring WiFi settings via serial...")
    print("=" * 60)
    
    cmd = [
        'python3', script_path,
        '-s', ssid,
        '-p', password,
        '--start-ip', start_ip,
        '-g', gateway
    ]
    
    result = subprocess.run(cmd)
    return result.returncode == 0

def main():
    parser = argparse.ArgumentParser(
        description='Flash and configure all ESP32 devices',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
This script:
  1. Builds firmware ONCE
  2. Flashes the SAME firmware to all devices (fast!)
  3. Reconfigures each device via serial with unique IP

Much faster than building 32 different firmwares!

Examples:
  # Basic usage
  %(prog)s --ssid MyWiFi --password mypass

  # Custom IP range
  %(prog)s --ssid MyWiFi --password mypass --start-ip 192.168.1.100

  # Build only (no flash)
  %(prog)s --build-only

  # Reconfigure only (no flash)
  %(prog)s --reconfig-only --ssid MyWiFi --password mypass
        """
    )
    
    parser.add_argument('--ssid', help='WiFi SSID')
    parser.add_argument('--password', help='WiFi password')
    parser.add_argument('--start-ip', default='192.168.1.50',
                        help='Starting IP address (default: 192.168.1.50)')
    parser.add_argument('--gateway', default='192.168.1.1',
                        help='Gateway IP (default: 192.168.1.1)')
    parser.add_argument('--project-dir', default=None,
                        help='ESP32 project directory (default: current dir)')
    parser.add_argument('--build-only', action='store_true',
                        help='Only build, do not flash')
    parser.add_argument('--reconfig-only', action='store_true',
                        help='Only reconfigure, do not build/flash')
    parser.add_argument('--skip-build', action='store_true',
                        help='Skip build, use existing firmware')
    
    args = parser.parse_args()
    
    # Determine project directory
    if args.project_dir:
        project_dir = args.project_dir
    else:
        project_dir = os.path.dirname(os.path.abspath(__file__))
    
    if not os.path.exists(os.path.join(project_dir, 'CMakeLists.txt')):
        print(f"Error: Not an ESP-IDF project directory: {project_dir}")
        return 1
    
    # Detect devices
    devices = detect_devices()
    if not devices and not args.build_only:
        print("Error: No devices found!")
        return 1
    
    print(f"Found {len(devices)} device(s)")
    
    # Reconfigure only mode
    if args.reconfig_only:
        if not args.ssid or not args.password:
            print("Error: --ssid and --password required for reconfigure mode")
            return 1
        
        return 0 if reconfigure_devices(args.ssid, args.password, args.start_ip, args.gateway) else 1
    
    # Build firmware
    if not args.skip_build:
        if not build_firmware(project_dir):
            return 1
    
    if args.build_only:
        print("\n Build complete. Use --skip-build to flash later.")
        return 0
    
    # Flash all devices
    flash_count = flash_all_devices(devices, project_dir)
    
    if flash_count == 0:
        print("✗ No devices flashed successfully")
        return 1
    
    # Reconfigure if credentials provided
    if args.ssid and args.password:
        print("\nWaiting for devices to boot...")
        time.sleep(5)
        
        if not reconfigure_devices(args.ssid, args.password, args.start_ip, args.gateway):
            print("✗ Reconfiguration failed")
            return 1
    else:
        print("\n" + "=" * 60)
        print("Flashing complete!")
        print("=" * 60)
        print("\nTo configure WiFi settings, run:")
        print(f"  python3 reconfig_simple.py -s YourSSID -p YourPassword --start-ip {args.start_ip}")
    
    print("\n✓ Done!")
    return 0

if __name__ == '__main__':
    sys.exit(main())