import os
import re

# pylint: disable=E0602
Import("env")  # noqa


def apply_memory_patches(content):
    """Apply IRAM, DRAM, and Flash patches to linker script content.

    Args:
        content: Linker script content as string

    Returns:
        Patched content as string
    """
    patches_applied = []

    # Replace IRAM size from 0x8000 (32KB) to 0x200000 (2MB)
    # The line looks like: iram1_0_seg : org = 0x40100000, len = 0x8000
    new_content = re.sub(
        r"(iram1_0_seg\s*:\s*org\s*=\s*0x40100000\s*,\s*len\s*=\s*)0x8000",
        r"\g<1>0x200000",
        content,
    )
    if new_content != content:
        patches_applied.append("IRAM: 32KB -> 2MB")
        content = new_content

    # Replace DRAM (BSS) size to allow larger uninitialized data sections
    # The line looks like: dram0_0_seg : org = 0x3FFE8000, len = 0x14000
    # Increase from 0x14000 (80KB) to 0x200000 (2MB)
    new_content = re.sub(
        r"(dram0_0_seg\s*:\s*org\s*=\s*0x3FFE8000\s*,\s*len\s*=\s*)0x14000",
        r"\g<1>0x200000",
        content,
    )
    if new_content != content:
        patches_applied.append("DRAM: 80KB -> 2MB")
        content = new_content

    # Replace Flash/irom0 size to allow larger code sections
    # The line looks like: irom0_0_seg : org = 0x40201010, len = 0xfeff0
    # Increase from 0xfeff0 (~1MB) to 0x2000000 (32MB) - fake huge flash for testing
    new_content = re.sub(
        r"(irom0_0_seg\s*:\s*org\s*=\s*0x40201010\s*,\s*len\s*=\s*)0x[0-9a-fA-F]+",
        r"\g<1>0x2000000",
        content,
    )
    if new_content != content:
        patches_applied.append("Flash: 1MB -> 32MB")
        content = new_content

    if patches_applied:
        print(f"  Patches applied: {', '.join(patches_applied)}")

    return content


def patch_linker_script_file(filepath, description):
    """Patch a single linker script file in place."""
    if not os.path.exists(filepath):
        print(f"ESPHome: {description} not found at {filepath}")
        return False

    print(f"ESPHome: Patching {description}...")
    with open(filepath, "r") as f:
        content = f.read()

    patched_content = apply_memory_patches(content)

    if patched_content != content:
        with open(filepath, "w") as f:
            f.write(patched_content)
        print(f"ESPHome: Successfully patched {description}")
        return True
    else:
        print(f"ESPHome: {description} already patched or no changes needed")
        return False


def patch_local_linker_script(source, target, env):
    """Patch the local.eagle.app.v6.common.ld in build directory for IRAM."""
    # Check if we're in testing mode
    build_flags = env.get("BUILD_FLAGS", [])
    testing_mode = any("-DESPHOME_TESTING_MODE" in flag for flag in build_flags)

    if not testing_mode:
        return

    # Patch the local linker script if it exists
    build_dir = env.subst("$BUILD_DIR")
    ld_dir = os.path.join(build_dir, "ld")
    if os.path.exists(ld_dir):
        local_ld = os.path.join(ld_dir, "local.eagle.app.v6.common.ld")
        if os.path.exists(local_ld):
            patch_linker_script_file(local_ld, "local.eagle.app.v6.common.ld")


# Check if we're in testing mode
build_flags = env.get("BUILD_FLAGS", [])
testing_mode = any("-DESPHOME_TESTING_MODE" in flag for flag in build_flags)

if testing_mode:
    # Create custom linker script immediately (before linker command is built)
    build_dir = env.subst("$BUILD_DIR")
    ldscript = env.GetProjectOption("board_build.ldscript", "")
    assert ldscript, "No linker script configured in board_build.ldscript"

    framework_dir = env.PioPlatform().get_package_dir("framework-arduinoespressif8266")
    assert framework_dir is not None, "Could not find framework-arduinoespressif8266 package"

    sdk_ld = os.path.join(framework_dir, "tools", "sdk", "ld", ldscript)
    custom_ld = os.path.join(build_dir, f"testing_{ldscript}")

    if os.path.exists(sdk_ld) and not os.path.exists(custom_ld):
        # Read and patch the SDK linker script
        with open(sdk_ld, "r") as f:
            content = f.read()

        patched_content = apply_memory_patches(content)

        # Write custom linker script
        with open(custom_ld, "w") as f:
            f.write(patched_content)

        print(f"ESPHome: Created custom linker script: {custom_ld}")

    # Tell the linker to use our custom script
    assert os.path.exists(custom_ld), f"Custom linker script not found: {custom_ld}"
    env.Replace(LDSCRIPT_PATH=custom_ld)
    print(f"ESPHome: Using custom linker script with patched memory limits")

    # Hook to patch local.eagle.app.v6.common.ld after it's created
    env.AddPreAction("$BUILD_DIR/${PROGNAME}.elf", patch_local_linker_script)