#!/usr/bin/env python3
"""
ESP32 Batch Configuration Tool
Detects all connected ESP32s and configures them with sequential Static IPs.
Requires: detect_esp32.py and config_device.py in the same directory.
"""

import sys
import os
import argparse
import time
import ipaddress
import re

# Ensure we can import the other scripts
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

try:
    import detect_esp32
    import config_device
except ImportError as e:
    print(f"Error: Could not import required modules ({e}).")
    print("Make sure 'detect_esp32.py' and 'config_device.py' are in the same folder.")
    sys.exit(1)

def natural_sort_key(device_obj):
    """
    Sorts ports naturally (ttyUSB2 comes before ttyUSB10)
    """
    s = device_obj.device
    # Split string into a list of integers and non-integers
    return [int(text) if text.isdigit() else text.lower()
            for text in re.split('([0-9]+)', s)]

def main():
    parser = argparse.ArgumentParser(
        description='Batch Config: Detects all ESP32s and configures sequential IPs',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    # Arguments matching config_device.py options
    parser.add_argument('--start-ip', required=True,
                        help='Starting Static IP (e.g., 192.168.1.101). Will increment for each device.')
    parser.add_argument('-s', '--ssid', default='ClubHouse2G',
                        help='WiFi SSID')
    parser.add_argument('-P', '--password', default='ez2remember',
                        help='WiFi password')
    parser.add_argument('-g', '--gateway', default='192.168.1.1',
                        help='Gateway IP')
    parser.add_argument('-m', '--netmask', default='255.255.255.0',
                        help='Netmask')
    parser.add_argument('-b', '--band', default='2.4G', choices=['2.4G', '5G'],
                        help='WiFi band')
    parser.add_argument('-B', '--bandwidth', default='HT20', choices=['HT20', 'HT40', 'VHT80'],
                        help='Channel bandwidth')
    parser.add_argument('-ps', '--powersave', default='NONE',
                        choices=['NONE', 'MIN', 'MIN_MODEM', 'MAX', 'MAX_MODEM'],
                        help='Power save mode')
    parser.add_argument('-M', '--mode', default='STA', choices=['STA', 'MONITOR'],
                        help='Operating mode')
    parser.add_argument('-mc', '--monitor-channel', type=int, default=36,
                        help='Monitor mode channel')
    parser.add_argument('-r', '--no-reboot', action='store_true',
                        help='Do NOT reboot devices after configuration')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Enable verbose output')

    args = parser.parse_args()

    # 1. Detect Devices
    print(f"{'='*60}")
    print("Step 1: Detecting ESP32 Devices...")
    print(f"{'='*60}")

    devices = detect_esp32.detect_esp32_devices()

    if not devices:
        print("No ESP32 devices found! Check USB connections.")
        sys.exit(1)

    # Sort devices naturally so IPs are assigned in order (USB0, USB1, USB2...)
    devices.sort(key=natural_sort_key)

    print(f"Found {len(devices)} devices:")
    for d in devices:
        print(f"  - {d.device} ({d.description})")
    print()

    # 2. Parse Starting IP
    try:
        start_ip_obj = ipaddress.IPv4Address(args.start_ip)
    except ipaddress.AddressValueError:
        print(f"Error: Invalid IP address format: {args.start_ip}")
        sys.exit(1)

    # 3. Configure Each Device
    print(f"{'='*60}")
    print("Step 2: Configuring Devices Sequentially")
    print(f"{'='*60}")

    success_count = 0
    fail_count = 0
    failed_devices = []

    for index, device in enumerate(devices):
        # Calculate current IP
        current_ip = str(start_ip_obj + index)
        port = device.device

        print(f"\n[{index+1}/{len(devices)}] Configuring {port} with IP {current_ip}...")

        # Call the config function from your existing script
        result = config_device.config_device(
            port=port,
            ip=current_ip,
            ssid=args.ssid,
            password=args.password,
            gateway=args.gateway,
            netmask=args.netmask,
            band=args.band,
            bandwidth=args.bandwidth,
            powersave=args.powersave,
            mode=args.mode,
            monitor_channel=args.monitor_channel,
            reboot=not args.no_reboot,
            verbose=args.verbose
        )

        if result:
            print(f"✓ Success: {port} -> {current_ip}")
            success_count += 1
        else:
            print(f"✗ Failed: {port}")
            fail_count += 1
            failed_devices.append(port)

        # Small delay to prevent USB power spikes if multiple devices reboot simultaneously
        if not args.no_reboot and index < len(devices) - 1:
            time.sleep(1.0)

    # 4. Summary
    print(f"\n{'='*60}")
    print("Batch Configuration Complete")
    print(f"{'='*60}")
    print(f"Total Devices: {len(devices)}")
    print(f"Successful:    {success_count}")
    print(f"Failed:        {fail_count}")

    if failed_devices:
        print("\nFailed Ports:")
        for p in failed_devices:
            print(f"  - {p}")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nBatch process interrupted by user.")