"""
    $(TYPEDSIGNATURES)

Given a function expression `expr`, return a callable version of it.

# Keyword arguments
- `eval_expression`: Whether to use `eval` to make `expr` callable. If `false`, uses
  RuntimeGeneratedFunctions.jl.
- `eval_module`: The module to `eval` the expression `expr` in. If `!eval_expression`,
  this is the cache and context module for the `RuntimeGeneratedFunction`.
"""
function eval_or_rgf(expr::Expr; eval_expression = false, eval_module = @__MODULE__)
    if eval_expression
        return eval_module.eval(expr)
    else
        return drop_expr(RuntimeGeneratedFunction(eval_module, eval_module, expr))
    end
end

eval_or_rgf(::Nothing; kws...) = nothing

"""
    $(TYPEDSIGNATURES)

Return the name for the `i`th argument in a function generated by `build_function_wrapper`.
"""
function generated_argument_name(i::Int)
    return Symbol(:__mtk_arg_, i)
end

function compute_array_variable_buffer_idxs(@nospecialize(args); ignore_vars = Set{SymbolicT}())
    _compute_array_variable_buffer_idxs(args isa Vector ? args : collect(args), ignore_vars)
end

function _compute_array_variable_buffer_idxs(args::Vector, ignore_vars)
    # map array symbolic to an identically sized array where each element is (buffer_idx, idx_in_buffer)
    var_to_arridxs = Dict{SymbolicT, Vector{Tuple{Int, Int}}}()
    for (i, arg) in enumerate(args)
        # filter out non-arrays
        # any element of args which is not an array is assumed to not contain a
        # scalarized array symbolic. This works because the only non-array element
        # is the independent variable
        arg isa Vector{SymbolicT} || continue

        # go through symbolics
        for (j, var) in enumerate(arg)
            var = unwrap(var)
            # filter out non-array-symbolics
            arrvar, isarr = split_indexed_var(var)
            isarr || continue
            arrvar in ignore_vars && continue
            # get and/or construct the buffer storing indexes
            idxbuffer = get(var_to_arridxs, arrvar, nothing)
            if idxbuffer === nothing
                idxbuffer = map(Returns((0, 0)), SU.stable_eachindex(arrvar))
            end
            idxbuffer[SU.as_linear_idx(SU.shape(arrvar), get_stable_index(var))] = (i, j)
            var_to_arridxs[arrvar] = idxbuffer
        end
    end

    return var_to_arridxs
end

function array_variable_buffer_idxs_to_assignments(
        var_to_arridxs::Dict{SymbolicT, Vector{Tuple{Int, Int}}};
        argument_name = generated_argument_name, buffer_offset = 0,
        filter_vars = nothing,
    )
    assignments = Assignment[]
    for (arrvar, idxs) in var_to_arridxs
        filter_vars === nothing || arrvar in filter_vars || continue
        # all elements of the array need to be present in `args` to form the
        # reconstructing assignment
        any(iszero ∘ first, idxs) && continue

        var_size_expr = Expr(:tuple)
        sharrvar = SU.shape(arrvar)
        for ax in sharrvar
            push!(var_size_expr.args, length(ax))
        end

        # if they are all in the same buffer, we can take a shortcut and `view` into it
        if allequal(Iterators.map(first, idxs))
            buffer_idx = first(first(idxs)) + buffer_offset
            idxs = map(last, idxs)
            shape = SU.ShapeVecT((1:length(idxs),))
            # if all the elements are contiguous and ordered, turn the array of indexes into a range
            # to help reduce allocations
            if first(idxs) < last(idxs) && vec(idxs) == first(idxs):last(idxs)
                expr = Expr(:call, view, argument_name(buffer_idx), first(idxs):last(idxs))
            elseif vec(idxs) == first(idxs):-1:last(idxs)
                expr = Expr(:call, view, argument_name(buffer_idx), first(idxs):-1:last(idxs))
            elseif length(idxs) <= 16
                expr = Expr(:call, SVector)
                append!(expr.args, idxs)
                expr = Expr(:call, view, argument_name(buffer_idx), expr)
            else
                expr = Expr(:call, view, argument_name(buffer_idx), idxs)
            end
        else
            if length(idxs) <= 16
                expr = Expr(:call, SVector)
            else
                expr = Expr(:vect)
            end
            for idx in idxs
                i, j = idx
                push!(expr.args, Expr(:ref, argument_name(i), j))
            end
        end
        expr = Expr(:call, reshape, expr, var_size_expr)
        if any(!isone ∘ first, sharrvar)
            expr_origin = Expr(:call, Origin)
            for ax in sharrvar
                push!(expr_origin.args, first(ax))
            end
            expr = Expr(:call, expr_origin, expr)
        end
        push!(assignments, Assignment(arrvar, expr))
    end

    return assignments
end

"""
    $(TYPEDSIGNATURES)

Given the arguments to `build_function_wrapper`, return a list of assignments which
reconstruct array variables if they are present scalarized in `args`.

# Keyword Arguments

- `argument_name` a function of the form `(::Int) -> Symbol` which takes the index of
  an argument to the generated function and returns the name of the argument in the
  generated function.
"""
function array_variable_assignments(
        @nospecialize(args...); ignore_vars = Set{SymbolicT}(), filter_vars = nothing,
        argument_name = generated_argument_name, buffer_offset = 0
    )
    var_to_arridxs = compute_array_variable_buffer_idxs(args; ignore_vars)
    return array_variable_buffer_idxs_to_assignments(var_to_arridxs; argument_name, buffer_offset, filter_vars)
end

"""
    $(TYPEDSIGNATURES)

Check if the variable `var` is a delayed variable, where `iv` is the independent
variable.
"""
function isdelay(var, iv)
    iv === nothing && return false
    if iscall(var) && ModelingToolkitBase.isoperator(var, Differential)
        return isdelay(arguments(var)[1], iv)
    end
    isvariable(var) || return false
    isparameter(var) && return false
    if iscall(var) && !ModelingToolkitBase.isoperator(var, Symbolics.Operator)
        args = arguments(var)
        length(args) == 1 || return false
        arg = args[1]
        isequal(arg, iv) && return false
        iscall(arg) || return true
        issym(operation(arg)) && !iscalledparameter(arg) && return false
        return true
    end
    return false
end

"""
The argument of generated functions corresponding to the history function.
"""
const DDE_HISTORY_FUN = SSym(:___history___; type = SU.FnType{Tuple{Any, <:Real}, Vector{Real}, Nothing}, shape = SU.Unknown(1))
const BVP_SOLUTION = SSym(:__sol__; type = Symbolics.FnType{Tuple{<:Real}, Vector{Real}, Nothing}, shape = SU.Unknown(1))
const DDE_AT_IDX_SYM = SSym(:__delayvar_idxₘₜₖ; type = Int, shape = UnitRange{Int}[])
const DDE_DELAY_SYM = SSym(:__delayxₘₜₖ; type = Real, shape = UnitRange{Int}[])

"""
    $(TYPEDSIGNATURES)

Turn delayed unknowns in `eqs` into calls to `DDE_HISTORY_FUNCTION`.

# Arguments

- `sys`: The system of DDEs.
- `eqs`: The equations to convert.

# Keyword Arguments

- `param_arg`: The name of the variable containing the parameter object.
- `histfn`: The history function to use for codegen, called as `histfn(p, t)`
"""
function delay_to_function(
        sys::AbstractSystem, eqs = full_equations(sys); param_arg = MTKPARAMETERS_ARG, histfn = DDE_HISTORY_FUN
    )
    return delay_to_function(
        eqs,
        get_iv(sys),
        Dict{Any, Int}(operation(s) => i for (i, s) in enumerate(unknowns(sys))),
        parameters(sys),
        histfn; param_arg
    )
end
function delay_to_function(eqs::Vector, iv, sts, ps, h; param_arg = MTKPARAMETERS_ARG)
    return delay_to_function.(eqs, (iv,), (sts,), (ps,), (h,); param_arg)
end
function delay_to_function(eq::Equation, iv, sts, ps, h; param_arg = MTKPARAMETERS_ARG)
    return delay_to_function(eq.lhs, iv, sts, ps, h; param_arg) ~ delay_to_function(
        eq.rhs, iv, sts, ps, h; param_arg
    )
end
function delay_to_function(expr, iv, sts, ps, h; param_arg = MTKPARAMETERS_ARG)
    if isdelay(expr, iv)
        v = operation(expr)
        time = arguments(expr)[1]
        idx = sts[v]
        return term(getindex, h(param_arg, time), idx, type = Real)
    elseif iscall(expr)
        return maketerm(
            typeof(expr),
            operation(expr),
            map(x -> delay_to_function(x, iv, sts, ps, h; param_arg), arguments(expr)),
            metadata(expr)
        )
    else
        return expr
    end
end

function __search_dervars_recurse(x::SymbolicT)
    return iscall(x) && Moshi.Match.@match x begin
        BSImpl.Term(; f) && if f isa Operator end => false
        _ => true
    end
end

struct ParameterArrayAssignments
    var_to_arridxs::Dict{SymbolicT, Vector{Tuple{Int, Int}}}
end

function should_invalidate_mutable_cache_entry(::Type{ParameterArrayAssignments}, patch::NamedTuple)
    return haskey(patch, :ps)
end

function find_arrvars_is_atomic(ex::SymbolicT)
    SU.default_is_atomic(ex) && Symbolics.isarraysymbolic(ex)
end

"""
    $(TYPEDSIGNATURES)

A wrapper around `build_function` which performs the necessary transformations for
code generation of all types of systems. `expr` is the expression returned from the
generated functions, and `args` are the arguments.

# Keyword Arguments

- `p_start`, `p_end`: Denotes the indexes in `args` where the buffers of the splatted
  `MTKParameters` object are present. These are collapsed into a single argument and
  destructured inside the function. `p_start` must also be provided for non-split systems
  since it is used by `wrap_delays`.
- `u_arg`: The index in `args` of the argument corresponding to `unknowns(sys)` (the `u`
  vector). If `-1` (the default), the u vector is not treated specially. Otherwise, the
  argument must be a `Vector` and is wrapped in a `DestructuredArgs` with the common
  identifier `MTKUNKNOWNS_ARG`, giving it the predictable name `___mtkunknowns___`.
- `compress_args`: A list of argument ranges that end before `p_start`.
  Each range will be compressed into a single argument to the function. For example,
  If there are 5 elements in `args` and `compress_args = [2:3]`, then the generated function
  will take 4 arguments, where the second should be an indexable collection of the second
  and third elements in `args`.
- `wrap_delays`: Whether to transform delayed unknowns of `sys` present in `expr` into
  calls to a history function. The history function is added to the list of arguments
  right before parameters, at the index `p_start`.
- `histfn`: The history function to use for transforming delayed terms. For any delayed
  term `x(expr)`, this is called as `histfn(p, expr)` where `p` is the parameter object.
- `histfn_symbolic`: The symbolic history function variable to add as an argument to the
  generated function.
- `wrap_code`: Forwarded to `build_function`.
- `add_observed`: Whether to add assignment statements for observed equations in the
  generated code.
- `obsidxs_to_use`: An iterable of `Int` of the specific observed equations to use in this
  function. This elides automatic observed detection.
- `create_bindings`: Whether to explicitly destructure arrays of symbolics present in
  `args` in the generated code. If `false`, all usages of the individual symbolics will
  instead call `getindex` on the relevant argument. This is useful if the generated
  function writes to one of its arguments and expects subsequent code to use the new
  values. Note that the collapsed `MTKParameters` argument will always be explicitly
  destructured regardless of this keyword argument.
- `output_type`: The type of the output buffer. If `mkarray` (see below) is `nothing`,
  this will be passed to the `similarto` argument of `build_function`. If `output_type`
  is `Tuple`, `expr` will be wrapped in `SymbolicUtils.Code.MakeTuple` (regardless of
  whether it is scalar or an array).
- `mkarray`: A function which accepts `expr` and `output_type` and returns a code
  generation object similar to `MakeArray` or `MakeTuple` to be used to generate
  code for `expr`.
- `wrap_mtkparameters`: Whether to collapse parameter buffers for a split system into a
  argument.
- `extra_assignments`: Extra `Assignment` statements to prefix to `expr`, after all other
  assignments.

All other keyword arguments are forwarded to `build_function`.
"""
Base.@nospecializeinfer function build_function_wrapper(
        sys::AbstractSystem, @nospecialize(expr), @nospecialize(args...); p_start = 2,
        p_end = is_time_dependent(sys) ? length(args) - 1 : length(args), compress_args = UnitRange{Int}[],
        non_standard_param_layout = false, u_arg::Integer = -1,
        wrap_delays = is_dde(sys), histfn = DDE_HISTORY_FUN, histfn_symbolic = histfn, wrap_code = identity,
        add_observed = true, obsidxs_to_use = nothing,
        create_bindings = false, output_type = nothing, mkarray = nothing,
        wrap_mtkparameters = true, extra_assignments = Assignment[], cse = true,
        optimize = nothing, kwargs...
    )
    isscalar = !(expr isa AbstractArray || symbolic_type(expr) == ArraySymbolic())
    obs = observed(sys)
    args = Vector{Any}(collect(args))
    assignments = Assignment[]

    if u_arg != -1
        args[u_arg] isa AbstractVector ||
            throw(ArgumentError("argument at u_arg = $u_arg must be a Vector, got $(typeof(args[u_arg]))"))
    end

    u_argument_name = if u_arg == -1
        generated_argument_name
    else
        push!(assignments, Assignment(generated_argument_name(u_arg), :___mtkunknowns___))
        i -> i == u_arg ? :___mtkunknowns___ : generated_argument_name(i)
    end
    # turn delayed unknowns into calls to the history function
    if wrap_delays
        param_arg = is_split(sys) ? MTKPARAMETERS_ARG : generated_argument_name(p_start)
        delayfn_expr = STerm(
            Base.Fix1, SArgsT((histfn_symbolic, param_arg));
            type = SU.FnType{Tuple, Vector{Real}, Any}, shape = SU.Unknown(1)
        )
        for (i, var) in enumerate(unknowns(sys))
            Moshi.Match.@match var begin
                BSImpl.Term(; f) && if f isa SymbolicT end => begin
                    rhs = STerm(
                        ComposedFunction, SArgsT((Base.Fix2(getindex, i), delayfn_expr));
                        type = SU.FnType{Tuple, Real, Any}, shape = UnitRange{Int}[]
                    )
                    push!(assignments, Assignment(f, rhs))
                end
                _ => nothing
            end
        end

        # add extra argument
        insert!(args, p_start, histfn_symbolic)
        p_start += 1
        p_end += 1
    end

    ir_info = get_ir_info(sys)
    expr = ir_info.obs_subber(expr)

    ir = get_irstructure(sys)
    required_arrvars = Set{SymbolicT}()
    search_buffer = SU.IRStructureSearchBuffer(ir, required_arrvars)
    SU.search_variables!(search_buffer, expr; is_atomic = find_arrvars_is_atomic, recurse = !SU.default_is_atomic)

    # assignments for reconstructing scalarized array symbolics
    if non_standard_param_layout
        append!(assignments, array_variable_assignments(args...; filter_vars = required_arrvars, argument_name = u_argument_name))
    else
        cached = check_mutable_cache(sys, ParameterArrayAssignments, ParameterArrayAssignments, nothing)
        if cached isa ParameterArrayAssignments
            param_var_to_arridxs = cached.var_to_arridxs
        else
            param_var_to_arridxs = compute_array_variable_buffer_idxs(args[p_start:p_end])
            store_to_mutable_cache!(sys, ParameterArrayAssignments, ParameterArrayAssignments(param_var_to_arridxs))
        end
        append!(
            assignments, array_variable_buffer_idxs_to_assignments(
                param_var_to_arridxs; buffer_offset = p_start - 1, filter_vars = required_arrvars
            )
        )
        other_assigns = array_variable_assignments(args...; ignore_vars = keys(param_var_to_arridxs), filter_vars = required_arrvars, argument_name = u_argument_name)
        append!(assignments, other_assigns)
    end
    append!(assignments, extra_assignments)

    for (i, arg) in enumerate(args)
        # Make sure to use the proper names for arguments
        args[i] = if u_arg != -1 && i == u_arg
            DestructuredArgs(arg, MTKUNKNOWNS_ARG; create_bindings)
        elseif symbolic_type(arg) == NotSymbolic() && arg isa AbstractArray
            DestructuredArgs(arg, generated_argument_name(i); create_bindings)
        else
            arg
        end
    end

    # wrap into a single MTKParameters argument
    if is_split(sys) && wrap_mtkparameters
        if p_start > p_end
            # In case there are no parameter buffers, still insert an argument
            args = [args[1:(p_start - 1)]; MTKPARAMETERS_ARG; args[(p_end + 1):end]]
        else
            # cannot apply `create_bindings` here since it doesn't nest
            args = [args[1:(p_start - 1)]; DestructuredArgs(collect(args[p_start:p_end]), MTKPARAMETERS_ARG); args[(p_end + 1):end]]
        end
    end

    sort!(compress_args; by = first)
    reverse!(compress_args)
    for (i, range) in enumerate(compress_args)
        compressed = DestructuredArgs(args[range], Symbol(:__compressed, i))
        deleteat!(args, range)
        insert!(args, first(range), compressed)
    end

    # add preface assignments
    if has_preface(sys) && (pref = preface(sys)) !== nothing
        append!(assignments, pref)
    end

    wrap_code = wrap_code .∘ wrap_assignments(false, assignments)

    # handling of `output_type` and `mkarray`
    similarto = nothing
    if output_type === Tuple
        expr = MakeTuple(Tuple(expr))
    elseif mkarray === nothing
        similarto = output_type
    else
        expr = mkarray(expr, output_type)
    end

    optimize = resolve_optimize_option(optimize)
    return Symbolics.codegen_function(get_irstructure(sys), expr, args; wrap_code, similarto, cse, optimize, kwargs...)
end

resolve_optimize_option(x) = x
resolve_optimize_option(::Nothing) = nothing

"""
    $(TYPEDEF)

A wrapper around a generated in-place and out-of-place function. The type-parameter `P`
must be a 3-tuple where the first element is the index of the parameter object in the
arguments, the second is the expected number of arguments in the out-of-place variant
of the function, and the third is a boolean indicating whether the generated functions
are for a split system. For scalar functions, the inplace variant can be `nothing`.
"""
struct GeneratedFunctionWrapper{P, O, I} <: Function
    f_oop::O
    f_iip::I
end

function GeneratedFunctionWrapper{P}(foop::O, fiip::I) where {P, O, I}
    return GeneratedFunctionWrapper{P, O, I}(foop, fiip)
end

function GeneratedFunctionWrapper{P}(::Type{Val{true}}, foop, fiip; kwargs...) where {P}
    return :($(GeneratedFunctionWrapper{P})($foop, $fiip))
end

function GeneratedFunctionWrapper{P}(::Type{Val{false}}, foop, fiip; kws...) where {P}
    return GeneratedFunctionWrapper{P}(eval_or_rgf(foop; kws...), eval_or_rgf(fiip; kws...))
end

function (gfw::GeneratedFunctionWrapper)(args...)
    return _generated_call(gfw, args...)
end

function SciMLBase.numargs(::GeneratedFunctionWrapper{P}) where {P}
    n_oop = P[2]
    return (n_oop, n_oop + 1)
end

@generated function _generated_call(gfw::GeneratedFunctionWrapper{P}, args...) where {P}
    paramidx, nargs, issplit = P
    iip = false
    # IIP case has one more argument
    if length(args) == nargs + 1
        nargs += 1
        paramidx += 1
        iip = true
    end
    if length(args) != nargs
        throw(ArgumentError("Expected $nargs arguments, got $(length(args))."))
    end

    # the function to use
    f = iip ? :(gfw.f_iip) : :(gfw.f_oop)
    # non-split systems just call it as-is
    if !issplit
        return :($f(args...))
    end
    if args[paramidx] <: Union{Tuple, MTKParameters} &&
            !(args[paramidx] <: Tuple{Vararg{Number}})
        # for split systems, call it as-is if the parameter object is a tuple or MTKParameters
        # but not if it is a tuple of numbers
        return :($f(args...))
    else
        # The user provided a single buffer/tuple for the parameter object, so wrap that
        # one in a tuple
        fargs = ntuple(Val(length(args))) do i
            i == paramidx ? :((args[$i], nothing)) : :(args[$i])
        end
        return :($f($(fargs...)))
    end
end

"""
    $(TYPEDSIGNATURES)

Optionally compile a method and optionally wrap it in a `GeneratedFunctionWrapper` on the
basis of `expression` `wrap_gfw`, both of type `Union{Type{Val{true}}, Type{Val{false}}}`.
`gfw_args` is the first type parameter of `GeneratedFunctionWrapper`. `f` is a tuple of
function expressions of the form `(oop, iip)` or a single out-of-place function expression.
Keyword arguments are forwarded to `eval_or_rgf`.
"""
function maybe_compile_function(
        expression, wrap_gfw::Type{Val{true}},
        gfw_args::Tuple{Int, Int, Bool}, f::NTuple{2, Expr}; kwargs...
    )
    return GeneratedFunctionWrapper{gfw_args}(expression, f...; kwargs...)
end

function maybe_compile_function(
        expression::Type{Val{false}}, wrap_gfw::Type{Val{false}},
        gfw_args::Tuple{Int, Int, Bool}, f::NTuple{2, Expr}; kwargs...
    )
    return eval_or_rgf.(f; kwargs...)
end

function maybe_compile_function(
        expression::Type{Val{true}}, wrap_gfw::Type{Val{false}},
        gfw_args::Tuple{Int, Int, Bool}, f::Union{Expr, NTuple{2, Expr}}; kwargs...
    )
    return f
end

function maybe_compile_function(
        expression, wrap_gfw::Type{Val{true}},
        gfw_args::Tuple{Int, Int, Bool}, f::Expr; kwargs...
    )
    return GeneratedFunctionWrapper{gfw_args}(expression, f, nothing; kwargs...)
end

function maybe_compile_function(
        expression::Type{Val{false}}, wrap_gfw::Type{Val{false}},
        gfw_args::Tuple{Int, Int, Bool}, f::Expr; kwargs...
    )
    return eval_or_rgf(f; kwargs...)
end
