# Copyright 2014-present PlatformIO <contact@platformio.org>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import shutil
from os import SEEK_CUR, SEEK_END
from os.path import basename, isfile
from pathlib import Path

from SCons.Script import Builder

Import("env")

board = env.BoardConfig()
mcu = board.get("build.mcu", "esp32")
is_xtensa = mcu in ("esp32", "esp32s2", "esp32s3")

cmake_dir = str(env.PioPlatform().get_package_dir("tool-cmake"))
cmake_cmd = f'"{Path(cmake_dir) / "bin" / "cmake"}"'

idf_dir = str(env.PioPlatform().get_package_dir("framework-espidf"))
data_embed_script = (
    f'"{Path(idf_dir) / "tools" / "cmake" / "scripts" / "data_file_embed_asm.cmake"}"'
)

#
# Embedded files helpers
#


def extract_files(cppdefines, files_type):
    result = []
    files = env.GetProjectOption("board_build.%s" % files_type, "").splitlines()
    if files:
        result.extend([str(Path("$PROJECT_DIR") / f.strip()) for f in files if f.strip()])
    else:
        files_define = "COMPONENT_" + files_type.upper()
        for define in cppdefines:
            if files_define not in define:
                continue

            value = define[1]
            if not isinstance(define, tuple):
                print("Warning! %s macro cannot be empty!" % files_define)
                return []

            if not isinstance(value, str):
                print(
                    "Warning! %s macro must contain "
                    "a list of files separated by ':'" % files_define
                )
                return []

            for f in value.split(":"):
                f = f.strip()
                if not f:
                    continue
                result.append(str(Path("$PROJECT_DIR") / f))

    for f in result:
        if not isfile(env.subst(f)):
            print('Warning! Could not find file "%s"' % basename(f))

    return result


def remove_config_define(cppdefines, files_type):
    for define in cppdefines:
        if files_type in define:
            env.ProcessUnFlags("-D%s" % "=".join(str(d) for d in define))
            return


def prepare_file(source, target, env):
    filepath = source[0].get_abspath()
    shutil.copy(filepath, filepath + ".piobkp")

    with open(filepath, "rb+") as fp:
        fp.seek(0, SEEK_END)
        size = fp.tell()
        if size == 0:
            fp.write(b"\0")
        else:
            fp.seek(-1, SEEK_END)
            if fp.read(1) != b"\0":
                fp.write(b"\0")


def revert_original_file(source, target, env):
    filepath = source[0].get_abspath()
    if isfile(filepath + ".piobkp"):
        shutil.move(filepath + ".piobkp", filepath)


def embed_files(files, files_type):
    for f in files:
        filename = basename(f) + ".txt.o"
        file_target = env.TxtToBin(str(Path("$BUILD_DIR") / filename), f)
        env.Depends("$PIOMAINPROG", file_target)
        if files_type == "embed_txtfiles":
            env.AddPreAction(file_target, prepare_file)
            env.AddPostAction(file_target, revert_original_file)
        env.AppendUnique(PIOBUILDFILES=[env.File(str(Path("$BUILD_DIR") / filename))])


def transform_to_asm(target, source, env):
    asm_targets = [str(Path("$BUILD_DIR") / (s.name + ".S")) for s in source]
    return asm_targets, source

    
env.Append(
    BUILDERS=dict(
        TxtToBin=Builder(
            action=env.VerboseAction(
                " ".join(
                    [
                        "riscv32-esp-elf-objcopy"
                        if not is_xtensa
                        else f"xtensa-{mcu}-elf-objcopy",
                        "--input-target",
                        "binary",
                        "--output-target",
                        "elf32-littleriscv" if not is_xtensa else "elf32-xtensa-le",
                        "--binary-architecture",
                        "riscv" if not is_xtensa else "xtensa",
                        "--rename-section",
                        ".data=.rodata.embedded",
                        "$SOURCE",
                        "$TARGET",
                    ]
                ),
                "Converting $TARGET",
            ),
            suffix=".txt.o",
        ),
        FileToAsm=Builder(
            action=env.VerboseAction(
                " ".join(
                    [
                        cmake_cmd,
                        "-DDATA_FILE=$SOURCE",
                        "-DSOURCE_FILE=$TARGET",
                        "-DFILE_TYPE=$FILE_TYPE",
                        "-P",
                        data_embed_script,
                    ]
                ),
                "Generating assembly for $TARGET",
            ),
            emitter=transform_to_asm,
            single_source=True,
        ),
    )
)


flags = env.get("CPPDEFINES")
for files_type in ("embed_txtfiles", "embed_files"):
    if (
        "COMPONENT_" + files_type.upper() not in env.Flatten(flags)
        and "build." + files_type not in board
    ):
        continue

    files = extract_files(flags, files_type)
    if "espidf" in env.subst("$PIOFRAMEWORK"):
        env.Requires(
            str(Path("$BUILD_DIR") / "${PROGNAME}.elf"),
            env.FileToAsm(
                files,
                FILE_TYPE="TEXT" if files_type == "embed_txtfiles" else "BINARY",
            ),
        )
    else:
        embed_files(files, files_type)
        remove_config_define(flags, files_type)