using Cassette, InteractiveUtils

Cassette.@context TraceurCtx

struct Trace
  seen::Set
  stack::Vector{Call}
  warn
  maxdepth::Int
end

function Trace(w; maxdepth=typemax(Int))
  Trace(Set(), Vector{Call}(), w, maxdepth)
end

isprimitive(f) = f isa Core.Builtin || f isa Core.IntrinsicFunction

const ignored_methods = Set([@which((1,2)[1])])
const ignored_functions = Set([getproperty, setproperty!])

function Cassette.prehook(ctx::TraceurCtx, f, args...)
  tra = ctx.metadata
  C = DynamicCall(f, args...)
  push!(tra.stack, C)
  nothing
end

function Cassette.posthook(ctx::TraceurCtx, out, f, args...)
  tra = ctx.metadata
  C = tra.stack[end]
  f = C.f
  T = typeof.((f, args))
  if !(f ∈ ignored_functions || T ∈ tra.seen || isprimitive(f) ||
       method(C) ∈ ignored_methods || method(C).module ∈ (Core, Core.Compiler))
    push!(tra.seen, T)
    analyse((a...) -> tra.warn(Warning(a..., copy(tra.stack))), C)
  end
  pop!(tra.stack)
  nothing
end

function Cassette.overdub(ctx::TraceurCtx, f, args...)
  tra = ctx.metadata
  if (length(tra.stack) > tra.maxdepth)
    return f(args...)
  else
    invoke(Cassette.overdub, Tuple{Cassette.Context, typeof(f), typeof.(args)...}, ctx, f, args...)
  end
end

trace(w, f; kwargs...) = Cassette.recurse(TraceurCtx(metadata=Trace(w; kwargs...)), f)

function warntrace(f; modules=[], kwargs...)
  trace(warning_printer(modules), f; kwargs...)
end

"""
    warnings(f; kwargs...)::Vector{Traceur.Warnings}

Collect all warnings generated by Traceur's analysis of the execution of the
no-arg function `f` and return them.

Possible keyword arguments:
- `maxdepth=typemax(Int)` constrols how far Traceur recurses through the call stack.
- If `modules` is nonempty, only warnings for methods defined in one of the modules specified will be printed.
"""
function warnings(f; modules=[], kwargs...)
  warnings = Warning[]
  trace(w -> push!(warnings, w), f; kwargs...)
  if !isempty(modules)
    filter!(x -> getmod(x) in modules, warnings)
  end
  return warnings
end

"""
    @trace(functioncall(args...), maxdepth=2, modules=[])

Analyse `functioncall(args...)` for common performance problems and print them to
the terminal.

Optional arguments:
- `maxdepth` constrols how far Traceur recurses through the call stack.
- If `modules` is nonempty, only warnings for methods defined in one of the modules specified will be printed.
"""
macro trace(ex, args...)
  quote
    warntrace(() -> $(esc(ex)); $(map(esc, args)...))
  end
end
