import os
import time
import shutil
import glob
from pathlib import Path
from string import ascii_uppercase
from os import path
import usb.core
import usb.backend.libusb1

# Device configurations
DEVICES = {
    'RP2040': {
        'vendor_id': 0x2E8A,
        'product_id': 0x000A,
        'name': 'RP2040'
    },
    'RP2350': {
        'vendor_id': 0x2E8A,
        'product_id': 0x0009,
        'name': 'RP2350'
    }
}

# USB control constants
RESET_INTERFACE = 2
REQUEST_BOOTSEL = 0x01
REQUEST_RESET = 0x02

# Timing constants
TIMEOUT_THRESHOLD_s = 10
TIMEOUT_INCREMENT_ms = 100

def detect_device():
    """
    Detect which Pico device is connected.
    Returns device type string or None if no device found.
    """
    print("Auto-detecting connected device...")
    
    for device_type, config in DEVICES.items():
        device = usb.core.find(idVendor=config['vendor_id'], idProduct=config['product_id'])
        if device is not None:
            print(f"{config['name']} detected")
            return device_type
    
    return None

def reset_device_to_bootloader(device_type):
    """
    Reset the specified device type to bootloader mode.
    """
    config = DEVICES[device_type]
    print(f"Resetting {config['name']} to bootloader mode...")
    
    device = usb.core.find(idVendor=config['vendor_id'], idProduct=config['product_id'])
    
    if device is None:
        print(f'No {config["name"]} device found')
        print(f'Make sure your {config["name"]} is connected and not in bootloader mode')
        return False
    
    print(f"{config['name']} device found, attempting to enter bootloader mode...")
    
    try:
        usb.util.claim_interface(device, RESET_INTERFACE)
        Request_Type = (0 << 6) | 1  # Standard Request Host-to-Device and to interface
        wIndex = RESET_INTERFACE
        wValue = 0  # No other data
        
        device.ctrl_transfer(Request_Type, REQUEST_BOOTSEL, wValue, wIndex, [])
        print("Bootloader reset command sent successfully!")
        print(f"{config['name']} should now appear as a USB drive for programming")
        return True
        
    except usb.core.USBError as e:
        # print(f"USB Error: {e}")
        # print("This is normal - the device resets and disconnects")
        return True  # This is actually expected behavior
    except Exception as e:
        print(f"Unexpected error: {e}")
        return False
    finally:
        try:
            usb.util.release_interface(device, RESET_INTERFACE)
        except:
            pass  # Device may have already disconnected

def find_uf2_drive():
    """
    Find the UF2 bootloader drive by looking for characteristic files.
    Both RP2040 and RP2350 use the same UF2 bootloader mechanism.
    """
    for drive in ascii_uppercase:
        if path.exists(drive + ":\\INDEX.HTM") and path.exists(drive + ":\\INFO_UF2.TXT"):
            return drive + ":\\"
    return ""

def copy_uf2_file(device_type):
    """
    Copy the UF2 file to the detected bootloader drive.
    """
    print("Looking for UF2 bootloader drive...")
    
    TARGET_DRIVE = find_uf2_drive()
    Timeout_Counter = 0
    
    # Wait for drive to appear
    while not TARGET_DRIVE:
        time.sleep(TIMEOUT_INCREMENT_ms/1000)
        TARGET_DRIVE = find_uf2_drive()
        
        Timeout_Counter = Timeout_Counter + TIMEOUT_INCREMENT_ms
        if Timeout_Counter >= TIMEOUT_THRESHOLD_s * 1000:
            print(f"Error: Drive for {device_type} could not be found")
            return False
    
    print(f"Found {device_type} drive at: {TARGET_DRIVE}")
    
    # Change to script directory and find UF2 files
    os.chdir(os.path.dirname(os.path.realpath(__file__)))
    
    uf2_files = glob.glob("./*.uf2")
    if len(uf2_files) == 0:
        print("Error: No uf2-file has been found")
        return False
    
    print(f"Copying {uf2_files[0]} to {device_type}...")
    
    try:
        shutil.copy(uf2_files[0], TARGET_DRIVE)
        print("Copy completed successfully!")
        return True
    except Exception as e:
        print(f"Error copying file: {e}")
        return False

def main():
    """
    Main function to orchestrate the programming process.
    """
    print("=== Pico Programmer ===")
    
    # Detect device
    device_type = detect_device()
    
    if device_type:
        # Device found - attempt reset to bootloader
        print(f"\n-----------------------")
        print(f"Resetting {device_type}...")
        print(f"-----------------------")
        
        # Reset to bootloader
        if not reset_device_to_bootloader(device_type):
            print("Failed to reset device to bootloader mode")
            os._exit(1)
        
        print(f"\n-----------------------")
        print(f"Copying uf2-file...")
        print(f"-----------------------")
        
        # Copy UF2 file
        if not copy_uf2_file(device_type):
            print("Failed to copy UF2 file")
            os._exit(1)
        
        print(f"\nProgramming {device_type} completed successfully!")
    
    else:
        # No device found - skip reset and try to find bootloader drive directly
        print("No supported Pico device found in normal mode")
        print("Skipping reset step and looking for existing bootloader drive...")
        
        print(f"\n-----------------------")
        print(f"Looking for bootloader drive...")
        print(f"-----------------------")
        
        # Try to find and copy to existing bootloader drive
        if not copy_uf2_file("RP2xxx"):  # Generic name since we don't know the specific type
            print("No bootloader drive found and no device to reset")
            print("Make sure your device is:")
            print("1. Connected and in normal mode (for auto-reset), OR")
            print("2. Already in bootloader mode (showing as USB drive)")
            os._exit(1)
        
        print(f"\nProgramming completed successfully!")

if __name__ == "__main__":
    main()