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 = {
        # Do NOT set CONFIG_BT_ENABLED here — forces Arduino framework
        # to pre-init BT controller at startup, which conflicts with
        # TinyUSB on ESP32-S3 (causes reboot on USB MIDI send).
        # NimBLEDevice::init() handles its own init.
        "CONFIG_BTDM_CTRL_MODE_BLE_ONLY": 1,
        "CONFIG_BT_NIMBLE_MAX_CONNECTIONS": 1,
        "CONFIG_BT_NIMBLE_TASK_STACK_SIZE": 6144,
    }

    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")


def patch_nimconfig():
    # Find nimconfig.h in the NimBLE-Arduino library
    search_dirs = [
        os.path.join(os.getcwd(), ".pio", "libdeps"),
        os.path.expanduser("~/.platformio/lib"),
    ]
    for env_dir in ["esp32s3", ""]:
        candidate = os.path.join(
            os.getcwd(), ".pio", "libdeps",
            env_dir, "NimBLE-Arduino", "src", "nimconfig.h"
        ) if env_dir else ""
        if candidate and os.path.exists(candidate):
            nimconfig_path = candidate
            break
    else:
        # Walk search dirs as fallback
        nimconfig_path = None
        for base in search_dirs:
            if not os.path.exists(base):
                continue
            for root, _dirs, files in os.walk(base):
                if "nimconfig.h" in files:
                    nimconfig_path = os.path.join(root, "nimconfig.h")
                    break
            if nimconfig_path:
                break

    if not nimconfig_path:
        print("WARNING: nimconfig.h not found, skipping ROLE define patch")
        return

    print(f"Patching {nimconfig_path} with #ifndef guards for ROLE defines")

    with open(nimconfig_path, 'r') as f:
        content = f.read()

    for role in ["CENTRAL", "OBSERVER", "PERIPHERAL", "BROADCASTER"]:
        pattern = re.compile(
            '#ifndef CONFIG_BT_NIMBLE_ROLE_' + role + '_DISABLED\n'
            '#define CONFIG_BT_NIMBLE_ROLE_' + role + '\n'
            '#endif'
        )
        replacement = (
            '#ifndef CONFIG_BT_NIMBLE_ROLE_' + role + '_DISABLED\n'
            '#ifndef CONFIG_BT_NIMBLE_ROLE_' + role + '\n'
            '#define CONFIG_BT_NIMBLE_ROLE_' + role + '\n'
            '#endif\n'
            '#endif'
        )
        if pattern.search(content):
            content = pattern.sub(replacement, content)
            print(f"  Guarded CONFIG_BT_NIMBLE_ROLE_{role}")
        else:
            print(f"  NOTE: CONFIG_BT_NIMBLE_ROLE_{role} pattern not found (may already be guarded)")

    with open(nimconfig_path, 'w') as f:
        f.write(content)
    print("nimconfig.h patched successfully")


patch_usb_ids()
patch_sdkconfig_bt()
patch_nimconfig()