using JET
using Test

onfail(::Function, ::Test.Pass) = nothing
onfail(f::Function, ::Test.Fail) = f()

"""
    JETFrameFingerprint(jet_report_type, module_name, method_name, jet_signature)
    JETFrameFingerprint(frame, jet_report_type, virtual_module_name_dot)

To filter out error reports generated by JET, we need a way to refer their contents.
Each error report contains virtualized stack trace leading to the error, each
element of the stack is a frame containing information about a `MethodInstance`, but
since JET analyzes the code in virtual modules, it is not straightforward to map these
`MethodInstances` to the types in the original module. The virtual modules are not persistent
across JET runs, which makes making references to the frames difficult.

To overcome this issue we summarize each frame by 4 values:
* `jet_report_type` -- the type of the corresponding error report
* `module_name` -- name of the module of the frames' captured `MethodInstance`
* `method_name` -- name of the method in the frames' captured `MethodInstance`
* `jet_signature` -- string representation of the frames' `JET.Signature` with references
to the virtual module removed.

You can generate all the `JETFrameFingerprint`s for current error reports using
`get_reports_filtered` with `verbose=true`.
"""
struct JETFrameFingerprint
    jet_report_type
    module_name::Symbol
    method_name::Symbol
    jet_signature::String
end

function _singnature_string(linfo::Core.MethodInstance, virtual_module_name_dot::String)
    io = IOBuffer()
    m = linfo.def
    if m isa Module
        Base.show_mi(io, linfo, #=from_stackframe=#true)
    else
        Base.StackTraces.show_spec_sig(IOContext(io, :backtrace=>true, :limit=>true), m, linfo.specTypes)
    end
    out = String(take!(io))
    return replace(out, virtual_module_name_dot => "")
end

function JETFrameFingerprint(
    frame::JET.VirtualFrame,
    jet_report_type,
    virtual_module_name_dot::String
)
    @nospecialize
    module_name = Symbol(replace(string(frame.linfo.def.module), virtual_module_name_dot => ""))
    method_name = frame.linfo.def.name
    # For UndefKeywordErrors, the code gets lowered in a way that `frame.linfo.def.name` is
    # `Type##kw`, one way to get the name of the relevant method is to look it up
    # from the call parameters.
    if method_name == Symbol("Type##kw")
        method_name = Symbol(split(string(frame.linfo.specTypes.types[end].parameters[1]), '.')[end])
    end

    # NOTE(tomasd): We extract signature information from frame.linfo.specTypes.
    # We create a string representation of that signature and we remove the references to the virtual module
    # in which the types and methods were defined. This makes this signature consistent
    # across JET runs.
    signature_string = _singnature_string(frame.linfo, virtual_module_name_dot)
    return JETFrameFingerprint(jet_report_type, module_name, method_name, signature_string)
end

# NOTE(tomasd): adapted from https://github.com/aviatesk/JET.jl/blob/v0.5.14/src/JET.jl#L629
"""
    get_reports_filtered(result::JET.JETToplevelResult; ignored_modules, jet_frames_to_skip)

Given a JET error analysis result (i.e. the output of `JET.report_package`), return those reports
that do not come from modules specified in `ignored_modules`, or match `jet_frames_to_skip`.

The modules and signatures to skip should be strings and could be copied from the anaysis results.
"""
function get_reports_filtered(
    result::JET.JETToplevelResult;
    jet_frames_to_skip = (),
    verbose = false,
)
    check_call_stack = !isempty(jet_frames_to_skip)
    # Handle API change in JET 0.11+: actual2virtual is now nothing instead of a Pair
    virtual_module_name_dot = if isnothing(result.res.actual2virtual)
        ""  # JET 0.11+: no virtual module prefix to strip
    else
        string(result.res.actual2virtual.second, '.')  # JET 0.9-0.10
    end
    reports = JET.get_reports(result)

    return filter!(reports) do report
        verbose && println(report)
        if check_call_stack || verbose
            for (depth, frame) in enumerate(report.vst)
                sig = JETFrameFingerprint(
                    frame, typeof(report), virtual_module_name_dot
                )
                if verbose
                    print("    ")
                    @show depth, sig
                end
                sig in jet_frames_to_skip && return false
            end
        end
        return true
    end
end

"""
    jet_test_package(package; ignored_modules = nothing, jet_frames_to_skip = ())

Run `JET.report_package` on `package` and test that there are no error reports
other than those matched by modules or frames provided through kwargs.

See [JETFrameFingerprint](@ref) and [get_reports_filtered](@ref) for more information.
"""
function jet_test_package(package; ignored_modules = nothing, jet_frames_to_skip = (), mode = :typo)
    result = JET.report_package(package; ignored_modules, toplevel_logger=nothing, mode)
    reports = get_reports_filtered(result; jet_frames_to_skip)
    onfail(@test(isempty(reports))) do
        # Handle API change in JET 0.11+
        if isnothing(result.res.actual2virtual)
            # JET 0.11+: simplified API, print result directly
            JET.print_reports(stdout, result)
        else
            # JET 0.9-0.10: use PostProcessor
            JET.print_reports(
                stdout,
                reports,
@static if pkgversion(JET) < v"0.10.7"
                JET.gen_postprocess(result.res.actual2virtual)
else
                JET.PostProcessor(result.res.actual2virtual)
end
            )
        end
        println()
        println("Remaining error reports and their `JETFrameFingerprint`s:")
        get_reports_filtered(result; jet_frames_to_skip, verbose=true)
        return
    end
    return
end

"""
    jet_test_file(file_path; ignored_modules = nothing, jet_frames_to_skip = ())

Run `JET.report_file` at `file_path` and test that there are no error reports
other than those matched by modules or frames provided through kwargs.

See [JETFrameFingerprint](@ref) and [get_reports_filtered](@ref) for more information.
"""
function jet_test_file(file_path; ignored_modules = nothing, jet_frames_to_skip = (), mode = :basic)
    result = JET.report_file(file_path; ignored_modules, toplevel_logger=nothing, mode)
    reports = get_reports_filtered(result; jet_frames_to_skip)
    onfail(@test(isempty(reports))) do
        # Handle API change in JET 0.11+
        if isnothing(result.res.actual2virtual)
            # JET 0.11+: simplified API, print result directly
            JET.print_reports(stdout, result)
        else
            # JET 0.9-0.10: use PostProcessor
            JET.print_reports(
                stdout,
                reports,
@static if pkgversion(JET) < v"0.10.7"
                JET.gen_postprocess(result.res.actual2virtual)
else
                JET.PostProcessor(result.res.actual2virtual)
end
            )
        end
        println()
        println("Remaining error reports and their `JETFrameFingerprint`s:")
        get_reports_filtered(result; jet_frames_to_skip, verbose=true)
        return
    end
    return
end
