
import glob
import os
from invoke import Collection, task
import shutil

ROOT_DIR = os.path.dirname(__file__)

DEFAULT_BUILD_DIR = "build"
DEFAULT_BOARD = "nrf5340dk/nrf5340/cpuapp/ns"
DEFAULT_TARGET = "app"

# Constants for builds
ZEPHYR_BASE = os.getenv('ZEPHYR_BASE')

# Build a specified target
@task
def build(ctx, target=DEFAULT_TARGET, board=DEFAULT_BOARD, debug=False, options=None):

    # Get any modules dependencies

    # Configs
    module_configs = ''
    # Drivers
    for config in glob.glob(f'{ROOT_DIR}/src/drivers/**/*.conf', recursive=True):
        module_configs += f'{config} '
    # Libs
    for config in glob.glob(f'{ROOT_DIR}/src/lib/**/*.conf', recursive=True):
        module_configs += f'{config} '

    # DTS bindings
    module_bindings = ''
    for binding_path in glob.glob(f'{ROOT_DIR}/src/drivers/**'):

        # Add the additional bindings if they exist within the module
        if os.path.exists(f'{binding_path}/dts/bindings'):
            module_bindings += f'{binding_path} '

    # Set the build command
    build_str = f'west build -b {board} --build-dir build__{board.replace("/", "_")} --sysbuild -p -- -DCONF_FILE="{module_configs} prj.conf"'

    # Add any addtional options passed
    if options:
        build_str += options

    # Invoke the actual build command
    with ctx.cd(ROOT_DIR):
        ctx.run(build_str)

@task
def clean(ctx):

    # Get all the generated build directories
    dirs = []
    for dir in glob.glob(f'{ROOT_DIR}/build*'):
        dirs.append(dir)

    # Clean all the build directories
    for dir in dirs:
        shutil.rmtree(f'{dir}')

@task
def flash(ctx, target=DEFAULT_TARGET, board=DEFAULT_BOARD):

    if target == "all":
        targets = ["mcuboot", "app", "radio"]
    else:
        targets = [target]

    # Flash all images requested
    for type in targets:
        print(f"Flashing {type}...")

        match type:
            case "app":
                with ctx.cd(ROOT_DIR):
                    ctx.run(f'west flash --skip-rebuild --build-dir build__{board.replace("/", "_")} --hex-file build__{board.replace("/", "_")}/firefly_node/zephyr/firefly.signed.hex')

            case "mcuboot":
                with ctx.cd(ROOT_DIR):
                    ctx.run(f'west flash --skip-rebuild --build-dir build__{board.replace("/", "_")} --hex-file build__{board.replace("/", "_")}/mcuboot/zephyr/mcuboot.hex')

            case "radio":
                with ctx.cd(ROOT_DIR):
                    ctx.run(f'west flash --skip-rebuild --build-dir build__{board.replace("/", "_")} --hex-file build__{board.replace("/", "_")}/merged_CPUNET.hex')

            case default:
                with ctx.cd(ROOT_DIR):
                    ctx.run(f'west flash --skip-rebuild --build-dir build__{board.replace("/", "_")} --hex-file build__{board.replace("/", "_")}/firefly_node/zephyr/firefly.signed.hex')
