using Statistics
import Base.Meta: isexpr

"""
    ret = @freshexec [setup_ex] ex

Runs `ex` in an external process and gets back the final result (, which is supposed to be
  such a simple Julia object that we can restore it from its string representation).
Running  in external process can be useful for testing JET analysis, because:
- the first time analysis is not affected by native code cache and JET's global report cache
- we can do something that might break a Julia process without wondering it breaks the
  original test process later on
The optional positional argument `setup_ex` runs before each execution of `ex` and defaults
  to `JET_LOAD_EX`, which just loads JET into the process.
"""
macro freshexec(args...)
    args = map(a->Expr(:quote,a), args)
    return Expr(:escape, Expr(:call, GlobalRef(@__MODULE__, :freshexec), __module__, args...))
end
freshexec(mod, ex) = freshexec(mod, JET_LOAD_EX, ex)
freshexec(args...) = _freshexec(args..., collect_last_result)

function collect_last_result(exs)
    lines = string.(exs)
    lines[end] = "ret = $(lines[end])"
    return join(lines, '\n')
end

"""
    stats = @benchmark_freshexec [setup_ex] ex

Runs `ex` in an external process and collects execution statistics from [`@timed`](@ref).
This is particularly useful for testing the performance of first-time analysis, where
  the native code cache and JET's global report cache have no effect for its performance.
The optional positional argument `setup_ex` runs before each execution of `ex` and its
  execution statistis are not included in the benchmark result; it defaults to
  `JET_WARMUP_EX`, which loads JET and runs a warm up analysis `@profile identity(nothing)`.
"""
macro benchmark_freshexec(args...)
    args = map(a->Expr(:quote,a), args)
    return Expr(:escape, Expr(:call, GlobalRef(@__MODULE__, :benchmark_freshexec), __module__, args...))
end
benchmark_freshexec(mod, ex) = benchmark_freshexec(mod, JET_WARMUP_EX, ex)
benchmark_freshexec(args...) = _freshexec(args..., collect_statistics)

function collect_statistics(exs)
    return """
    ret = @timed begin
        $(join(exs, '\n'))
        nothing # ensure `stats` can be parsed
    end
    """
end

"""
    stats = @nbenchmark_freshexec [ntimes = 5] [setup_ex] ex

Runs `ex` in an external process multiple times (which can be configured by the optional
  keyword argument `ntimes`), and collects execution statistics from [`@timed`](@ref).
  The statistics are generated by taking the `mean` of all the trials.
This is particularly useful for benchmarking the performance of first-time analysis, where
  the native code cache and JET's global report cache have no effect for its performance.
The optional positional argument `setup_ex` runs before each execution of `ex` and its
  execution statistis are not included in the benchmark result; it defaults to
  `JET_WARMUP_EX`, which loads JET and runs a warm up analysis `@profile identity(nothing)`.
"""
macro nbenchmark_freshexec(args...)
    isn(x) = isexpr(x, :(=)) && first(x.args) === :ntimes
    i = findfirst(isn, args)
    ntimes = i === nothing ? 5 : Expr(:quote, last(args[i].args))
    args = map(a->Expr(:quote,a), filter(!isn, args))
    return Expr(:escape, Expr(:call, GlobalRef(@__MODULE__, :nbenchmark_freshexec), ntimes, __module__, args...))
end
nbenchmark_freshexec(ntimes, mod, ex) = nbenchmark_freshexec(ntimes, mod, JET_WARMUP_EX, ex)
function nbenchmark_freshexec(ntimes, args...)
    stats = [_freshexec(args..., collect_statistics) for _ in 1:ntimes]
    return (; time   = mean(getproperty.(stats, :time)),
              bytes  = mean(getproperty.(stats, :bytes)),
              gctime = mean(getproperty.(stats, :gctime)),
              )
end

function _freshexec(mod, setup_ex, ex, exs2script)
    # we need to flatten block expression into a toplevel expression to correctly handle
    # e.g. macro expansions
    setup_exs = isexpr(setup_ex, :block) ? setup_ex.args : [setup_ex]
    setup_script = join(string.(setup_exs), '\n')
    exs = isexpr(ex, :block) ? ex.args : [ex]
    script = exs2script(exs)

    prog = """
    old = stdout
    rw, wr = redirect_stdout()
    $(setup_script)
    $(script)
    redirect_stdout(old)
    close(rw); close(wr)
    println(stdout, repr(ret))
    """

    cmd = Cmd([JULIA_BIN, "-e", prog])
    io = IOBuffer()
    run(pipeline(cmd; stdout = io))

    return Core.eval(mod, Meta.parse(String(take!(io))))
end

const JET_LOAD_EX = :(using JET)
const JET_WARMUP_EX = quote
    using JET
    @profile_call identity(nothing) # warm-up for JET analysis
end
const JULIA_BIN = normpath(Sys.BINDIR, "julia")
