import io

import buildconfig
import yaml
from mozbuild.preprocessor import Preprocessor

HEADER_TEMPLATE = """\
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#ifndef %(includeguard)s
#define %(includeguard)s

/* This file is generated by jit/GenerateABIFunctionType.py. Do not edit! */

%(contents)s

#endif // %(includeguard)s
"""


def generate_header(c_out, includeguard, contents):
    c_out.write(
        HEADER_TEMPLATE
        % {
            "includeguard": includeguard,
            "contents": contents,
        }
    )


def load_yaml(yaml_path):
    # First invoke preprocessor.py so that we can use #ifdef JS_SIMULATOR in
    # the YAML file.
    pp = Preprocessor()
    pp.context.update(buildconfig.defines["ALLDEFINES"])
    pp.out = io.StringIO()
    pp.do_filter("substitution")
    pp.do_include(yaml_path)
    contents = pp.out.getvalue()
    return yaml.safe_load(contents)


def cpp_arg_type(arg_type):
    if arg_type == "General":
        return "intptr_t"
    elif arg_type == "Int32":
        return "int32_t"
    elif arg_type == "Int64":
        return "int64_t"
    elif arg_type == "Float32":
        return "float"
    elif arg_type == "Float64":
        return "double"
    elif arg_type == "Void":
        return "void"
    else:
        raise ValueError(arg_type)


def func_type_name(func_type):
    if "name" in func_type:
        return func_type["name"]

    # Autogenerate a name like `General_GeneralGeneral` if none is specified
    return f"{func_type['ret']}_{''.join(func_type['args'])}"


def func_type_has_floats(func_type):
    for arg in func_type["args"]:
        if arg in {"Float32", "Float64"}:
            return True
    return False


# Generate the ARM32 argument loading for a func type for when soft-FP is enabled
def arm32_soft_fp_args(func_type):
    # This must match ABIArgGenerator::softNext() in Assembler-arm.cpp
    contents = ""
    numIntArgRegs = 4
    intRegIndex = 0
    stackOffset = 0
    for i, arg in enumerate(func_type["args"]):
        if i != 0:
            contents += ", "

        if arg == "General":
            if intRegIndex == numIntArgRegs:
                contents += f"stack_pointer[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"a{intRegIndex}"
                intRegIndex += 1
        elif arg == "Int32":
            if intRegIndex == numIntArgRegs:
                contents += f"stack_pointer[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"a{intRegIndex}"
                intRegIndex += 1
        elif arg == "Int64":
            intRegIndex += intRegIndex % 2
            if intRegIndex == numIntArgRegs:
                stackOffset += stackOffset % 2
                contents += f"MakeInt64(stack_pointer[{stackOffset}], stack_pointer[{stackOffset + 1}])"
                stackOffset += 2
            else:
                contents += f"MakeInt64(a{intRegIndex}, a{intRegIndex + 1})"
                intRegIndex += 2
        elif arg == "Float32":
            if intRegIndex == numIntArgRegs:
                contents += f"mozilla::BitwiseCast<float>(stack_pointer[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"mozilla::BitwiseCast<float>(a{intRegIndex})"
                intRegIndex += 1
        elif arg == "Float64":
            intRegIndex += intRegIndex % 2
            if intRegIndex == numIntArgRegs:
                stackOffset += stackOffset % 2
                contents += f"mozilla::BitwiseCast<double>(MakeInt64(stack_pointer[{stackOffset}], stack_pointer[{stackOffset + 1}]))"
                stackOffset += 2
            else:
                contents += f"mozilla::BitwiseCast<double>(MakeInt64(a{intRegIndex}, a{intRegIndex + 1}))"
                intRegIndex += 2
    assert intRegIndex <= numIntArgRegs
    return contents


# Generate the ARM32 argument loading for a func type for when hard-FP is enabled
def arm32_hard_fp_args(func_type):
    # This must match ABIArgGenerator::hardNext() in Assembler-arm.cpp
    contents = ""
    numIntArgRegs = 4
    numFloatArgRegs = 16
    intRegIndex = 0
    floatRegIndex = 0
    stackOffset = 0
    for i, arg in enumerate(func_type["args"]):
        if i != 0:
            contents += ", "

        if arg == "General":
            if intRegIndex == numIntArgRegs:
                contents += f"stack_pointer[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"a{intRegIndex}"
                intRegIndex += 1
        elif arg == "Int32":
            if intRegIndex == numIntArgRegs:
                contents += f"stack_pointer[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"a{intRegIndex}"
                intRegIndex += 1
        elif arg == "Int64":
            intRegIndex += intRegIndex % 2
            if intRegIndex == numIntArgRegs:
                stackOffset += stackOffset % 2
                contents += f"MakeInt64(stack_pointer[{stackOffset}], stack_pointer[{stackOffset + 1}])"
                stackOffset += 2
            else:
                contents += f"MakeInt64(a{intRegIndex}, a{intRegIndex + 1})"
                intRegIndex += 2
        elif arg == "Float32":
            if floatRegIndex == numFloatArgRegs:
                contents += f"mozilla::BitwiseCast<float>(stack_pointer[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"s{floatRegIndex}"
                floatRegIndex += 1
        elif arg == "Float64":
            floatRegIndex += floatRegIndex % 2
            if floatRegIndex == numFloatArgRegs:
                stackOffset += stackOffset % 2
                contents += f"mozilla::BitwiseCast<double>(MakeInt64(stack_pointer[{stackOffset}], stack_pointer[{stackOffset + 1}]))"
                stackOffset += 2
            else:
                contents += f"d{round(floatRegIndex / 2)}"
                floatRegIndex += 2
    assert intRegIndex <= numIntArgRegs
    assert floatRegIndex <= numFloatArgRegs
    return contents


def arm32_simulator_dispatch(func_types):
    contents = ""
    for func_type in func_types:
        hard_fp_args = arm32_hard_fp_args(func_type)
        soft_fp_args = arm32_soft_fp_args(func_type)
        ret = func_type["ret"]
        has_ret = ret != "Void"
        ret_setter = "ret = " if has_ret else ""

        contents += f"case js::jit::Args_{func_type_name(func_type)}: {{\\\n"
        contents += f"  auto target = reinterpret_cast<Prototype_{func_type_name(func_type)}>(external);\\\n"
        if has_ret:
            contents += f"  {cpp_arg_type(ret)} ret;\\\n"
        if func_type_has_floats(func_type):
            contents += "  if (ARMFlags::UseHardFpABI()) {\\\n"
            contents += f"    {ret_setter}target({hard_fp_args});\\\n"
            contents += "  } else {\\\n"
            contents += f"    {ret_setter}target({soft_fp_args});\\\n"
            contents += "  }\\\n"
        else:
            # No float args means we don't need to check the float ABI and
            # either generated args will do.
            contents += f"  {ret_setter}target({soft_fp_args});\\\n"
        contents += "  scratchVolatileRegisters((void*)target);\\\n"
        if ret in {"General", "Int32", "Int64"}:
            contents += "  setCallResult(ret);\\\n"
        elif ret == "Float32":
            contents += "  setCallResultFloat(ret);\\\n"
        elif ret == "Float64":
            contents += "  setCallResultDouble(ret);\\\n"
        contents += "  break;\\\n"
        contents += "}\\\n"
    return contents


# Generate the ARM64 argument loading for a func type
def arm64_args(func_type):
    # This must match ABIArgGenerator::next() in Assembler-arm64.cpp
    contents = ""
    numIntArgRegs = 8
    numFloatArgRegs = 8
    intRegIndex = 0
    floatRegIndex = 0
    stackOffset = 0
    for i, arg in enumerate(func_type["args"]):
        if i != 0:
            contents += ", "

        if arg == "General":
            if intRegIndex == numIntArgRegs:
                contents += f"sp[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"x{intRegIndex}"
                intRegIndex += 1
        elif arg == "Int32":
            if intRegIndex == numIntArgRegs:
                contents += f"sp[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"x{intRegIndex}"
                intRegIndex += 1
        elif arg == "Int64":
            if intRegIndex == numIntArgRegs:
                contents += f"sp[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"x{intRegIndex}"
                intRegIndex += 1
        elif arg == "Float32":
            if floatRegIndex == numFloatArgRegs:
                contents += f"mozilla::BitwiseCast<float>(sp[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"s{floatRegIndex}"
                floatRegIndex += 1
        elif arg == "Float64":
            if floatRegIndex == numFloatArgRegs:
                contents += f"mozilla::BitwiseCast<double>(sp[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"d{floatRegIndex}"
                floatRegIndex += 1
    assert intRegIndex <= numIntArgRegs
    assert floatRegIndex <= numFloatArgRegs
    return contents


def arm64_simulator_dispatch(func_types):
    contents = ""
    for func_type in func_types:
        args = arm64_args(func_type)
        contents += f"case js::jit::Args_{func_type_name(func_type)}: {{\\\n"
        contents += f"  auto target = reinterpret_cast<Prototype_{func_type_name(func_type)}>(nativeFn);\\\n"
        ret = func_type["ret"]
        if ret == "Void":
            contents += f"  target({args});\\\n"
        else:
            contents += f"  auto ret = target({args});\\\n"
        if ret == "General":
            contents += "  setGPR64Result(ret);\\\n"
        elif ret == "Int32":
            contents += "  setGPR32Result(ret);\\\n"
        elif ret == "Int64":
            contents += "  setGPR64Result(ret);\\\n"
        elif ret == "Float32":
            contents += "  setFP32Result(ret);\\\n"
        elif ret == "Float64":
            contents += "  setFP64Result(ret);\\\n"
        contents += "  break;\\\n"
        contents += "}\\\n"
    return contents


# Generate the LoongArch64 argument loading for a func type
def loongarch64_args(func_type):
    # This must match ABIArgGenerator::next() in Assembler-loong64.cpp
    contents = ""
    numIntArgRegs = 8
    numFloatArgRegs = 8
    intRegIndex = 0
    floatRegIndex = 0
    stackOffset = 0
    for i, arg in enumerate(func_type["args"]):
        if i != 0:
            contents += ", "

        if arg == "General":
            if intRegIndex == numIntArgRegs:
                contents += f"sp_[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"a{intRegIndex}_"
                intRegIndex += 1
        elif arg == "Int32":
            if intRegIndex == numIntArgRegs:
                contents += f"I32(sp_[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"I32(a{intRegIndex}_)"
                intRegIndex += 1
        elif arg == "Int64":
            if intRegIndex == numIntArgRegs:
                contents += f"sp_[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"a{intRegIndex}_"
                intRegIndex += 1
        elif arg == "Float32":
            if floatRegIndex == numFloatArgRegs:
                contents += f"*mozilla::BitwiseCast<float*>(sp_[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"f{floatRegIndex}_s"
                floatRegIndex += 1
        elif arg == "Float64":
            if floatRegIndex == numFloatArgRegs:
                contents += f"mozilla::BitwiseCast<double>(sp_[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"f{floatRegIndex}_d"
                floatRegIndex += 1
    assert intRegIndex <= numIntArgRegs
    assert floatRegIndex <= numFloatArgRegs
    return contents


def loongarch64_simulator_dispatch(func_types):
    contents = ""
    for func_type in func_types:
        args = loongarch64_args(func_type)
        contents += f"case js::jit::Args_{func_type_name(func_type)}: {{\\\n"
        contents += f"  auto target = reinterpret_cast<Prototype_{func_type_name(func_type)}>(nativeFn);\\\n"
        ret = func_type["ret"]
        if ret == "Void":
            contents += f"  target({args});\\\n"
        else:
            contents += f"  auto ret = target({args});\\\n"
        if ret == "General":
            contents += "  setCallResult(ret);\\\n"
        elif ret == "Int32":
            contents += "  setCallResult(I64(ret));\\\n"
        elif ret == "Int64":
            contents += "  setCallResult(ret);\\\n"
        elif ret == "Float32":
            contents += "  setCallResultFloat(ret);\\\n"
        elif ret == "Float64":
            contents += "  setCallResultDouble(ret);\\\n"
        contents += "  break;\\\n"
        contents += "}\\\n"
    return contents


# Generate the MIPS64 argument loading for a func type
def mips64_args(func_type):
    # This must match ABIArgGenerator::next() in Assembler-mips64.cpp
    contents = ""
    numIntArgRegs = 8
    numFloatArgRegs = 8
    regIndex = 0
    stackOffset = 0
    for i, arg in enumerate(func_type["args"]):
        if i != 0:
            contents += ", "

        if arg == "General":
            if regIndex == numIntArgRegs:
                contents += f"sp_[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"a{regIndex}_"
                regIndex += 1
        elif arg == "Int32":
            if regIndex == numIntArgRegs:
                contents += f"I32(sp_[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"I32(a{regIndex}_)"
                regIndex += 1
        elif arg == "Int64":
            if regIndex == numIntArgRegs:
                contents += f"sp_[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"a{regIndex}_"
                regIndex += 1
        elif arg == "Float32":
            if regIndex == numFloatArgRegs:
                contents += f"*mozilla::BitwiseCast<float*>(sp_[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"f{regIndex + 12}_s"
                regIndex += 1
        elif arg == "Float64":
            if regIndex == numFloatArgRegs:
                contents += f"mozilla::BitwiseCast<double>(sp_[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"f{regIndex + 12}_d"
                regIndex += 1
    assert regIndex <= numIntArgRegs
    assert numIntArgRegs == numFloatArgRegs
    return contents


def mips64_simulator_dispatch(func_types):
    contents = ""
    for func_type in func_types:
        args = mips64_args(func_type)
        contents += f"case js::jit::Args_{func_type_name(func_type)}: {{\\\n"
        contents += f"  auto target = reinterpret_cast<Prototype_{func_type_name(func_type)}>(nativeFn);\\\n"
        ret = func_type["ret"]
        if ret == "Void":
            contents += f"  target({args});\\\n"
        else:
            contents += f"  auto ret = target({args});\\\n"
        if ret == "General":
            contents += "  setCallResult(ret);\\\n"
        elif ret == "Int32":
            contents += "  setCallResult(I64(ret));\\\n"
        elif ret == "Int64":
            contents += "  setCallResult(ret);\\\n"
        elif ret == "Float32":
            contents += "  setCallResultFloat(ret);\\\n"
        elif ret == "Float64":
            contents += "  setCallResultDouble(ret);\\\n"
        contents += "  break;\\\n"
        contents += "}\\\n"
    return contents


def riscv64_args(func_type):
    # This must match ABIArgGenerator::next() in Assembler-riscv64.cpp
    contents = ""
    numIntArgRegs = 8
    numFloatArgRegs = 8
    regIndex = 0
    floatRegIndex = 0
    stackOffset = 0
    for i, arg in enumerate(func_type["args"]):
        if i != 0:
            contents += ", "

        if arg == "General":
            if regIndex == numIntArgRegs:
                contents += f"sp_[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"arg{regIndex}"
                regIndex += 1
        elif arg == "Int32":
            if regIndex == numIntArgRegs:
                contents += f"I32(sp_[{stackOffset}])"
                stackOffset += 1
            else:
                contents += f"I32(arg{regIndex})"
                regIndex += 1
        elif arg == "Int64":
            if regIndex == numIntArgRegs:
                contents += f"sp_[{stackOffset}]"
                stackOffset += 1
            else:
                contents += f"arg{regIndex}"
                regIndex += 1
        elif arg == "Float32":
            if floatRegIndex == numFloatArgRegs:
                contents += f"mozilla::BitwiseCast<float>(static_cast<uint32_t>(sp_[{stackOffset}]))"
                stackOffset += 1
            else:
                contents += f"getFpuRegisterFloat(fa{floatRegIndex})"
                floatRegIndex += 1
        elif arg == "Float64":
            if floatRegIndex == numFloatArgRegs:
                contents += f"mozilla::BitwiseCast<double>(static_cast<uint64_t>(sp_[{stackOffset}]))"
                stackOffset += 1
            else:
                contents += f"getFpuRegisterDouble(fa{floatRegIndex})"
                floatRegIndex += 1
        else:
            raise ValueError(f"Unknown arg type: {arg}")
    return contents


def riscv64_simulator_dispatch(func_types):
    contents = ""
    for func_type in func_types:
        args = riscv64_args(func_type)
        contents += f"case js::jit::Args_{func_type_name(func_type)}: {{\\\n"
        contents += f"  auto target = reinterpret_cast<Prototype_{func_type_name(func_type)}>(nativeFn);\\\n"
        ret = func_type["ret"]
        if ret == "Void":
            contents += f"  target({args});\\\n"
        else:
            contents += f"  auto ret = target({args});\\\n"
        if ret == "Void":
            pass
        elif ret == "General":
            contents += "  setCallResult(ret);\\\n"
        elif ret == "Int32":
            contents += "  setCallResult(I64(ret));\\\n"
        elif ret == "Int64":
            contents += "  setCallResult(ret);\\\n"
        elif ret == "Float32":
            contents += "  setCallResultFloat(ret);\\\n"
        elif ret == "Float64":
            contents += "  setCallResultDouble(ret);\\\n"
        else:
            raise ValueError(f"Unknown ret type: {ret}")
        contents += "  break;\\\n"
        contents += "}\\\n"
    return contents


def main(c_out, yaml_path):
    func_types = load_yaml(yaml_path)

    # Define the ABIFunctionType enum
    contents = "#define ABI_FUNCTION_TYPE_ENUM \\\n"
    for func_type in func_types:
        name = "Args_" + func_type_name(func_type)
        args = ", ".join(f"ABIType::{p}" for p in func_type["args"])
        ret = f"ABIType::{func_type['ret']}"

        contents += f"    {name} = detail::MakeABIFunctionType({ret}, {{{args}}}),\\\n"
    contents += "\n"

    # Define the prototypes of the types
    contents += "#define ABI_FUNCTION_TYPE_SIM_PROTOTYPES \\\n"
    for func_type in func_types:
        name = "Prototype_" + func_type_name(func_type)
        args = ", ".join(cpp_arg_type(p) for p in func_type["args"])
        ret = cpp_arg_type(func_type["ret"])

        contents += f"    typedef {ret} (*{name})({args});\\\n"
    contents += "\n"

    contents += "#define ABI_FUNCTION_TYPE_ARM64_SIM_DISPATCH \\\n"
    contents += arm64_simulator_dispatch(func_types)
    contents += "\n"

    contents += "#define ABI_FUNCTION_TYPE_ARM32_SIM_DISPATCH \\\n"
    contents += arm32_simulator_dispatch(func_types)
    contents += "\n"

    contents += "#define ABI_FUNCTION_TYPE_LOONGARCH64_SIM_DISPATCH \\\n"
    contents += loongarch64_simulator_dispatch(func_types)
    contents += "\n"

    contents += "#define ABI_FUNCTION_TYPE_MIPS64_SIM_DISPATCH \\\n"
    contents += mips64_simulator_dispatch(func_types)
    contents += "\n"

    contents += "#define ABI_FUNCTION_TYPE_RISCV64_SIM_DISPATCH \\\n"
    contents += riscv64_simulator_dispatch(func_types)
    contents += "\n"

    generate_header(c_out, "jit_ABIFunctionTypeGenerated_h", contents)