import os
import re


def patch_usb_ids():
    # Find the core's pins_arduino.h for ESP32-S3 DevKitC-1
    core_variants = os.path.expanduser("~/.platformio/packages/framework-arduinoespressif32/variants")
    if not os.path.exists(core_variants):
        # Try PlatformIO default location
        core_variants = os.path.join(os.environ.get("PLATFORMIO_PACKAGES_DIR", ""), "framework-arduinoespressif32", "variants")
    
    pins_file = os.path.join(core_variants, "esp32s3_devkitc", "pins_arduino.h")
    
    if not os.path.exists(pins_file):
        # Try alternative variant name
        pins_file = os.path.join(core_variants, "esp32s3", "pins_arduino.h")
    
    if os.path.exists(pins_file):
        print(f"Patching {pins_file} with JOC Midi USB descriptors")
        
        # Read and replace
        with open(pins_file, 'r') as f:
            content = f.read()
        
        # Replace USB_VID and USB_PID
        content = content.replace(
            '#define USB_VID 0x303a',
            '#define USB_VID 0x1235'
        )
        content = content.replace(
            '#define USB_PID 0x1001',
            '#define USB_PID 0x0103'
        )
        # Add or replace manufacturer/product
        if 'USB_MANUFACTURER' not in content:
            content = content.replace(
                '#define USB_PID 0x0103',
                '#define USB_PID 0x0103\n#define USB_MANUFACTURER "JOC"\n#define USB_PRODUCT "JOC Midi"'
            )
        else:
            content = content.replace('USB_MANUFACTURER "Novation"', 'USB_MANUFACTURER "JOC"')
            content = content.replace('USB_PRODUCT "Launchpad X"', 'USB_PRODUCT "JOC Midi"')
        
        with open(pins_file, 'w') as f:
            f.write(content)
        print("USB VID/PID patched successfully")
    else:
        print(f"WARNING: Could not find pins_arduino.h at {pins_file}")


def patch_sdkconfig_bt():
    # Find the precompiled sdkconfig.h for this board variant
    sdk_dir = os.path.expanduser(
        "~/.platformio/packages/framework-arduinoespressif32/tools/sdk/esp32s3/qio_opi/include"
    )
    alt_name = os.path.join(os.environ.get("PLATFORMIO_PACKAGES_DIR", ""),
                            "framework-arduinoespressif32/tools/sdk/esp32s3/qio_opi/include")
    if not os.path.exists(sdk_dir):
        sdk_dir = alt_name

    sdkconfig_file = os.path.join(sdk_dir, "sdkconfig.h")
    
    if not os.path.exists(sdkconfig_file):
        print(f"WARNING: sdkconfig.h not found at {sdkconfig_file}")
        return

    print(f"Patching {sdkconfig_file} with BLE/NimBLE support")
    
    with open(sdkconfig_file, 'r') as f:
        content = f.read()

    defines = {
        "CONFIG_BT_ENABLED": 1,
        "CONFIG_BT_NIMBLE_ENABLED": 1,
        "CONFIG_BTDM_CTRL_MODE_BLE_ONLY": 1,
        "CONFIG_BT_NIMBLE_MAX_CONNECTIONS": 1,
        "CONFIG_BT_NIMBLE_TASK_STACK_SIZE": 6144,
        "CONFIG_BT_NIMBLE_ROLE_PERIPHERAL": 1,
        "CONFIG_BT_NIMBLE_ROLE_BROADCASTER": 1,
    }

    changed = False
    for name, value in defines.items():
        pattern = re.compile(
            r'^[#/]*\s*#?\s*(define\s+' + re.escape(name) + r'\b).*$',
            re.MULTILINE
        )
        if pattern.search(content):
            # Already defined, replace the line
            content = pattern.sub(r'#define ' + name + ' ' + str(value), content)
            changed = True
        elif f'#define {name}' not in content:
            # Not present at all, add it
            content += f'\n#define {name} {value}'
            changed = True

    if changed:
        with open(sdkconfig_file, 'w') as f:
            f.write(content)
        print("BLE/NimBLE sdkconfig patched successfully")
    else:
        print("BLE/NimBLE already configured in sdkconfig")


patch_usb_ids()
patch_sdkconfig_bt()