module NonlinearSolveSpeedMappingExt

using SpeedMapping: speedmapping

using NonlinearSolveBase: NonlinearSolveBase, is_fw_wrapped, get_raw_f
using NonlinearSolve: NonlinearSolve, SpeedMappingJL
using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
using Setfield: @set

function SciMLBase.__solve(
        prob::NonlinearProblem, alg::SpeedMappingJL, args...;
        abstol = nothing, maxiters = 1000, alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = false),
        maxtime = nothing, store_trace::Val = Val(false),
        termination_condition = nothing, kwargs...
    )
    # Unwrap AutoSpecialize — external packages do their own AD
    if is_fw_wrapped(prob.f.f)
        prob = @set prob.f.f = get_raw_f(prob.f.f)
    end

    if haskey(kwargs, :alias_u0)
        alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = kwargs[:alias_u0])
    end
    alias_u0 = alias.alias_u0
    NonlinearSolveBase.assert_extension_supported_termination_condition(
        termination_condition, alg
    )

    m!, u,
        resid = NonlinearSolveBase.construct_extension_function_wrapper(
        prob; alias_u0, make_fixed_point = Val(true)
    )
    tol = NonlinearSolveBase.get_tolerance(abstol, eltype(u))

    time_limit = ifelse(maxtime === nothing, 1000, maxtime)

    sol = speedmapping(
        u; m!, abstol = tol, pnorm = Inf, maps_limit = maxiters,
        orders = Tuple(alg.orders),
        initial_learning_rate = alg.σ_min == 0.0 ? 1.0 : alg.σ_min,
        store_trace = store_trace isa Val{true},
        time_limit
    )
    res = prob.u0 isa Number ? first(sol.minimizer) : sol.minimizer
    resid = NonlinearSolveBase.Utils.evaluate_f(prob, res)

    return SciMLBase.build_solution(
        prob, alg, res, resid;
        original = sol, stats = SciMLBase.NLStats(sol.maps, 0, 0, 0, sol.maps),
        retcode = ifelse(sol.status == :first_order, ReturnCode.Success, ReturnCode.Failure)
    )
end

end
