#!/usr/bin/env python3
"""
ESP32 USB Device Detection Script
Detects and lists ESP32 devices connected via USB
"""

import serial.tools.list_ports
import subprocess
import sys


def detect_chip_type(port):
    """
    Try to detect the ESP32 chip type using esptool.py
    Returns chip type string or 'Unknown'
    """
    try:
        result = subprocess.run(
            ['esptool.py', '--port', port, 'chip_id'],
            capture_output=True,
            text=True,
            timeout=10
        )
        
        output = result.stdout + result.stderr
        
        if 'ESP32-S3' in output:
            return 'ESP32-S3'
        elif 'ESP32-S2' in output:
            return 'ESP32-S2'
        elif 'ESP32-C3' in output:
            return 'ESP32-C3'
        elif 'ESP32-C6' in output:
            return 'ESP32-C6'
        elif 'ESP32-H2' in output:
            return 'ESP32-H2'
        elif 'ESP32' in output:
            return 'ESP32'
        
    except Exception:
        pass
    
    return 'Unknown'


def guess_chip_type_from_usb(vid, pid):
    """
    Make an educated guess about chip type based on USB VID:PID
    Returns chip type string or None
    """
    if vid == 0x303A:  # Espressif USB
        if pid == 0x1001:
            return 'ESP32-S3/C3'  # Built-in USB, could be either
        elif pid == 0x1002:
            return 'ESP32-S3/C3'
    
    # For CP210x, CH340, etc., we can't tell from USB alone
    return None


def detect_esp32_devices():
    """
    Detect ESP32 devices connected via USB.
    Returns a list of tuples containing (port, description, hwid)
    """
    esp32_devices = []
    
    # Common USB vendor IDs used by ESP32 boards
    ESP32_VID_PIDS = [
        ('10C4', 'EA60'),  # Silicon Labs CP210x
        ('1A86', '7523'),  # CH340
        ('1A86', '55D4'),  # CH9102
        ('0403', '6001'),  # FTDI FT232
        ('0403', '6010'),  # FTDI FT2232
        ('303A', '1001'),  # Espressif USB JTAG/serial debug unit (ESP32-C3, ESP32-S3)
        ('303A', '1002'),  # Espressif USB JTAG/serial debug unit
    ]
    
    # Keywords that might appear in ESP32 device descriptions
    ESP32_KEYWORDS = [
        'CP210',
        'CH340',
        'CH9102',
        'UART',
        'USB-SERIAL',
        'USB SERIAL',
        'JTAG',
        'ESP32',
    ]
    
    # Get all available ports
    ports = serial.tools.list_ports.comports()
    
    for port in ports:
        # Check by VID:PID
        if port.vid is not None and port.pid is not None:
            vid = f"{port.vid:04X}"
            pid = f"{port.pid:04X}"
            
            if any((vid == v and pid == p) for v, p in ESP32_VID_PIDS):
                esp32_devices.append(port)
                continue
        
        # Check by description keywords
        description_upper = port.description.upper()
        if any(keyword in description_upper for keyword in ESP32_KEYWORDS):
            esp32_devices.append(port)
    
    return esp32_devices


def main():
    import argparse
    
    parser = argparse.ArgumentParser(description='ESP32 USB Device Detection')
    parser.add_argument('--no-probe', action='store_true',
                       help='Skip chip type probing (faster, less accurate)')
    args = parser.parse_args()
    
    print("=" * 60)
    print("ESP32 USB Device Detection")
    print("=" * 60)
    print()
    
    # Detect ESP32 devices
    esp32_devices = detect_esp32_devices()
    
    if esp32_devices:
        print(f"Found {len(esp32_devices)} ESP32 device(s):\n")
        
        for idx, device in enumerate(esp32_devices, 1):
            print(f"Device {idx}:")
            print(f"  Port:        {device.device}")
            print(f"  Description: {device.description}")
            print(f"  Hardware ID: {device.hwid}")
            
            if device.vid is not None and device.pid is not None:
                print(f"  VID:PID:     {device.vid:04X}:{device.pid:04X}")
            
            # Chip type detection (probe by default)
            chip_type = None
            if not args.no_probe:
                print(f"  Probing chip type...", end='', flush=True)
                chip_type = detect_chip_type(device.device)
                print(f"\r  Chip Type:   {chip_type}")
            else:
                # Try to guess from USB VID:PID
                if device.vid is not None and device.pid is not None:
                    chip_type = guess_chip_type_from_usb(device.vid, device.pid)
                    if chip_type:
                        print(f"  Chip Guess:  {chip_type} (remove --no-probe for exact type)")
                    else:
                        print(f"  Chip Type:   Unknown (remove --no-probe to detect)")
            
            if device.manufacturer:
                print(f"  Manufacturer: {device.manufacturer}")
            
            if device.serial_number:
                print(f"  Serial:      {device.serial_number}")
            
            print()
        
        print("=" * 60)
        print(f"Total ESP32 devices detected: {len(esp32_devices)}")
        if args.no_probe:
            print("Tip: Remove --no-probe flag to detect exact chip types")
        print("=" * 60)
    else:
        print("No ESP32 devices detected.")
        print()
        print("Available USB serial devices:")
        all_ports = serial.tools.list_ports.comports()
        
        if all_ports:
            for port in all_ports:
                print(f"  - {port.device}: {port.description}")
                if port.vid is not None and port.pid is not None:
                    print(f"    VID:PID = {port.vid:04X}:{port.pid:04X}")
        else:
            print("  No USB serial devices found.")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\nDetection interrupted by user.")
    except Exception as e:
        print(f"\nError: {e}")
        print("\nMake sure pyserial is installed: pip install pyserial")