"""
**IJulia** is a [Julia-language](http://julialang.org/) backend
combined with the [Jupyter](http://jupyter.org/) interactive
environment (also used by [IPython](http://ipython.org/)).  This
combination allows you to interact with the Julia language using
Jupyter/IPython's powerful [graphical
notebook](http://ipython.org/notebook.html), which combines code,
formatted text, math, and multimedia in a single document.

The `IJulia` module is used in three ways

* Typing `using IJulia; notebook()` will launch the Jupyter notebook
  interface in your web browser.  This is an alternative to launching
  `jupyter notebook` directly from your operating-system command line.
* In a running notebook, the `IJulia` module is loaded and `IJulia.somefunctions`
  can be used to interact with the running IJulia kernel:

  - `IJulia.load(filename)` and `IJulia.load_string(s)` load the contents
    of a file or a string, respectively, into a notebook cell.
  - `IJulia.clear_output()` to clear the output from the notebook cell,
    useful for simple animations.
  - `IJulia.clear_history()` to clear the history variables `In` and `Out`.
  - `push_X_hook(f)` and `pop_X_hook(f)`, where `X` is either
    `preexecute`, `postexecute`, or `posterror`.  This allows you to
    insert a "hook" function into a list of functions to execute
    when notebook cells are evaluated.
  - `IJulia.set_verbose()` enables verbose output about what IJulia
    is doing internally; this is mainly used for debugging.

* It is used internally by the IJulia kernel when talking
  to the Jupyter server.
"""
module IJulia
export notebook, jupyterlab, nbclassic, installkernel

import SHA
using ZMQ
import Base: invokelatest, RefValue
import Dates
using Dates: now, format, UTC, ISODateTimeFormat
import Random
import Random: seed!
using Base64: Base64EncodePipe
import REPL
import Logging

# InteractiveUtils is not used inside IJulia, but loaded in src/kernel.jl
# and this import makes it possible to load InteractiveUtils from the IJulia namespace
import InteractiveUtils

const depfile = joinpath(dirname(@__FILE__), "..", "deps", "deps.jl")
isfile(depfile) || error("IJulia not properly installed. Please run Pkg.build(\"IJulia\")")
include(depfile) # generated by Pkg.build("IJulia")

# use our own random seed for msg_id so that we
# don't alter the user-visible random state (issue #336)
const IJulia_RNG = seed!(Random.MersenneTwister(0))
import UUIDs
uuid4() = string(UUIDs.uuid4(IJulia_RNG))

"""
IPython message struct.
"""
mutable struct Msg
    idents::Vector{String}
    header::Dict
    content::Dict
    parent_header::Dict
    metadata::Dict
    buffers::Vector{Vector{UInt8}}

    function Msg(idents, header::Dict, content::Dict,
                 parent_header=Dict{String,Any}(), metadata=Dict{String,Any}(),
                 buffers=Vector{UInt8}[])
        new(idents, header, content, parent_header, metadata, buffers)
    end
end

mutable struct Comm{target}
    id::String
    primary::Bool
    on_msg::Function
    on_close::Function
    function (::Type{Comm{target}})(id, primary, on_msg, on_close, kernel) where {target}
        comm = new{target}(id, primary, on_msg, on_close)
        kernel.comms[id] = comm
        return comm
    end
end

# similar to Pkg.REPLMode.MiniREPL, a minimal REPL-like emulator
# for use with Pkg.do_cmd.  We have to roll our own to
# make sure it uses the redirected stdout, and because
# we don't have terminal support.
struct MiniREPL <: REPL.AbstractREPL
    display::TextDisplay
end
REPL.REPLDisplay(repl::MiniREPL) = repl.display

@kwdef mutable struct Kernel
    verbose::Bool = IJULIA_DEBUG
    inited::Bool = false
    current_module::Module = Main
    shutting_down::Threads.Atomic{Bool} = Threads.Atomic{Bool}(false)

    # These fields are special and are mirrored to their corresponding global
    # variables.
    In::Dict{Int, String} = Dict{Int, String}()
    Out::Dict{Int, Any} = Dict{Int, Any}()
    ans::Any = nothing
    n::Int = 0

    capture_stdout::Bool = true
    capture_stderr::Bool = !IJULIA_DEBUG
    capture_stdin::Bool = true

    minirepl::Union{MiniREPL, Nothing} = nothing

    # This dict holds a map from CommID to Comm so that we can
    # pick out the right Comm object when messages arrive
    # from the front-end.
    comms::Dict{String, Comm} = Dict{String, Comm}()

    shutdown::Function = start_shutdown

    # the following constants need to be initialized in init().
    publish::RefValue{Socket} = Ref{Socket}()
    raw_input::RefValue{Socket} = Ref{Socket}()
    requests::RefValue{Socket} = Ref{Socket}()
    control::RefValue{Socket} = Ref{Socket}()
    heartbeat::RefValue{Socket} = Ref{Socket}()
    zmq_context::RefValue{Context} = Ref{Context}()
    profile::Dict{String, Any} = Dict{String, Any}()
    connection_file::Union{String, Nothing} = nothing
    read_stdout::RefValue{Base.PipeEndpoint} = Ref{Base.PipeEndpoint}()
    read_stderr::RefValue{Base.PipeEndpoint} = Ref{Base.PipeEndpoint}()
    socket_locks::Dict{Socket, ReentrantLock} = Dict{Socket, ReentrantLock}()
    sha_ctx::RefValue{SHA.SHA_CTX} = Ref{SHA.SHA_CTX}()
    hmac_key::Vector{UInt8} = UInt8[]

    stop_event::Base.Event = Base.Event()
    waitloop_task::RefValue{Task} = Ref{Task}()

    # This looks a bit strange, but it's not uncommon for loading packages to
    # invalidate IJulia code. This can contribute significantly to compilation
    # latency so after each execute_request we launch tasks to precompile
    # certain request handlers asynchronously in the hope that it cuts down on
    # perceived latency by the user. We store the tasks here to prevent any
    # lingering tasks upon shutdown that could cause hangs.
    execute_precompile_task::RefValue{Task} = Ref{Task}()
    completion_precompile_task::RefValue{Task} = Ref{Task}()

    requests_task::RefValue{Task} = Ref{Task}()
    watch_stdout_task::RefValue{Task} = Ref{Task}()
    watch_stderr_task::RefValue{Task} = Ref{Task}()
    flush_stdout_task::RefValue{Task} = Ref{Task}()
    flush_stderr_task::RefValue{Task} = Ref{Task}()

    # name=>iobuffer for each stream ("stdout","stderr") so they can be sent in flush
    bufs::Dict{String, IOBuffer} = Dict{String, IOBuffer}()
    bufs_locks::Dict{String, ReentrantLock} = Dict{String, ReentrantLock}()
    # max output per code cell is 512 kb by default
    max_output_per_request::RefValue{Int} = Ref(1 << 19)

    # Variable so that display can be done in the correct Msg context
    execute_msg::Msg = Msg(["julia"], Dict("username"=>"jlkernel", "session"=>uuid4()), Dict())
    # Variable tracking the number of bytes written in the current execution request
    stdio_bytes::Int = 0
    # Use an array to accumulate "payloads" for the execute_reply message
    execute_payloads::Vector{Dict} = Dict[]

    heartbeat_threadid::Vector{Int} = zeros(Int, 128) # sizeof(uv_thread_t) <= 8 on Linux, OSX, Win

    # queue of objects to display at end of cell execution
    displayqueue::Vector{Any} = Any[]
end

function Base.setproperty!(kernel::Kernel, name::Symbol, x)
    # These fields need to be assigned explicitly to their global counterparts
    if name ∈ (:ans, :n, :In, :Out, :inited)
        setproperty!(IJulia, name, x)
    end

    setfield!(kernel, name, x)
end

function Base.wait(kernel::Kernel)
    if isassigned(kernel.waitloop_task)
        wait(kernel.waitloop_task[])
    end
end

function start_shutdown(kernel::Kernel)
    kernel.shutting_down[] = true
    kernel.inited = false

    # First we call zmq_ctx_shutdown() to close the context and stop all sockets
    # from working. We don't call ZMQ.close(::Context) directly because that
    # currently isn't threadsafe:
    # https://github.com/JuliaInterop/ZMQ.jl/issues/256
    ZMQ.lib.zmq_ctx_shutdown(kernel.zmq_context[])

    # Wait for the heartbeat thread to stop
    @ccall uv_thread_join(kernel.heartbeat_threadid::Ptr{Int})::Cint

    # Now all the sockets should have been cancelled and the eventloop tasks
    # should be ready to shutdown.
    notify(kernel.stop_event)
end

function Base.close(kernel::Kernel)
    for task_ref in (kernel.execute_precompile_task, kernel.completion_precompile_task)
        if isassigned(task_ref)
            wait(task_ref[])
        end
    end

    # Reset the IO streams first so that any later errors get printed
    if kernel.capture_stdout
        redirect_stdout(orig_stdout[])
        close(kernel.read_stdout[])
        wait(kernel.flush_stdout_task[])
        wait(kernel.watch_stdout_task[])
    end
    if kernel.capture_stderr
        redirect_stderr(orig_stderr[])
        close(kernel.read_stderr[])
        wait(kernel.flush_stderr_task[])
        wait(kernel.watch_stderr_task[])
    end
    if kernel.capture_stdin
        redirect_stdin(orig_stdin[])
    end

    # Reset the logger so that @log statements work and pop the InlineDisplay
    if isassigned(orig_logger)
        # orig_logger seems to not be set during precompilation
        Logging.global_logger(orig_logger[])
    end
    popdisplay()

    start_shutdown(kernel)
    wait(kernel)

    # Close all sockets
    close(kernel.publish[])
    close(kernel.raw_input[])
    close(kernel.requests[])
    close(kernel.control[])
    close(kernel.heartbeat[])
    close(kernel.zmq_context[])

    # Reset global variables
    IJulia.n = 0
    IJulia.ans = nothing
    IJulia.In = Dict{Int, String}()
    IJulia.Out = Dict{Int, Any}()
    IJulia._default_kernel = nothing
    IJulia.CommManager.comms = Dict{String, CommManager.Comm}()
    IJulia.profile = Dict{String, Any}()
end

function Kernel(f::Function, profile::Union{String, Dict}; kwargs...)
    kernel = Kernel(; kwargs...)
    if profile isa Dict
        init([], kernel, profile)
    else
        init([profile], kernel)
    end

    try
        f(kernel)
    finally
        close(kernel)
    end
end

_default_kernel::Union{Kernel, Nothing} = nothing

function get_kernel_or_error(kernel=_default_kernel)
    if isnothing(kernel)
        error("Kernel has not been initialized.")
    end
    return kernel
end

"""
    set_verbose(v=true)

This function enables (or disables, for `set_verbose(false)`) verbose
output from the IJulia kernel, when called within a running notebook.
This consists of log messages printed to the terminal window where
`jupyter` was launched, displaying information about every message sent
or received by the kernel.   Used for debugging IJulia.
"""
function set_verbose(v::Bool=true, kernel=_default_kernel)
    kernel = get_kernel_or_error(kernel)
    kernel.verbose = v
end

"""
`inited` is a global variable that is set to `true` if the IJulia
kernel is running, i.e. in a running IJulia notebook.  To test
whether you are in an IJulia notebook, therefore, you can check
`isdefined(Main, :IJulia) && IJulia.inited`.
"""
inited::Bool = false

function set_current_module(m::Module; kernel=_default_kernel)
    kernel = get_kernel_or_error(kernel)
    kernel.current_module = m
end

#######################################################################
include(joinpath("..", "deps", "kspec.jl"))
include("jupyter.jl")
#######################################################################

"""
    load_string(s, replace=false)

Load the string `s` into a new input code cell in the running IJulia notebook,
somewhat analogous to the `%load` magics in IPython. If the optional argument
`replace` is `true`, then `s` replaces the *current* cell rather than creating
a new cell.
"""
function load_string(s::AbstractString, replace::Bool=false, kernel=_default_kernel)
    push!(kernel.execute_payloads, Dict(
        "source"=>"set_next_input",
        "text"=>s,
        "replace"=>replace
    ))
    return nothing
end

"""
    load(filename, replace=false)

Load the file given by `filename` into a new input code cell in the running
IJulia notebook, analogous to the `%load` magics in IPython.
If the optional argument `replace` is `true`, then the file contents
replace the *current* cell rather than creating a new cell.
"""
load(filename::AbstractString, replace::Bool=false, kernel=_default_kernel) =
    load_string(read(filename, String), replace, kernel)

#######################################################################
# History: global In/Out and other exported history variables
"""
`In` is a global dictionary of input strings, where `In[n]`
returns the string for input cell `n` of the notebook (as it was
when it was *last evaluated*).
"""
In::Dict{Int, String} = Dict{Int, String}()
"""
`Out` is a global dictionary of output values, where `Out[n]`
returns the output from the last evaluation of cell `n` in the
notebook.
"""
Out::Dict{Int, Any} = Dict{Int, Any}()
"""
`ans` is a global variable giving the value returned by the last
notebook cell evaluated.
"""
ans::Any = nothing

# execution counter
"""
`IJulia.n` is the (integer) index of the last-evaluated notebook cell.
"""
n::Int = 0

#######################################################################
# methods to clear history or any subset thereof

function clear_history(indices; kernel=_default_kernel)
    for n in indices
        delete!(kernel.In, n)
        if haskey(kernel.Out, n)
            delete!(kernel.Out, n)
        end
    end
end

# since a range could be huge, intersect it with 1:n first
clear_history(r::AbstractRange{<:Integer}) =
    invoke(clear_history, Tuple{Any}, intersect(r, 1:n))

function clear_history(; kernel=_default_kernel)
    empty!(kernel.In)
    empty!(kernel.Out)
    kernel.ans = nothing
end

"""
    clear_history([indices])

The `clear_history()` function clears all of the input and output
history stored in the running IJulia notebook.  This is sometimes
useful because all cell outputs are remember in the `Out` global variable,
which prevents them from being freed, so potentially this could
waste a lot of memory in a notebook with many large outputs.

The optional `indices` argument is a collection of indices indicating
a subset of cell inputs/outputs to clear.
"""
clear_history

#######################################################################
# methods to print history or any subset thereof
function history(io::IO, indices::AbstractVector{<:Integer}; kernel=_default_kernel)
    for n in intersect(indices, 1:kernel.n)
      if haskey(kernel.In, n)
        println(io, kernel.In[n])
      end
    end
end

history(io::IO, x::Union{Integer,AbstractVector{<:Integer}}...; kernel=_default_kernel) = history(io, vcat(x...); kernel)
history(x...; kernel=_default_kernel) = history(stdout, x...; kernel)
history(io::IO, x...; kernel=_default_kernel) = throw(MethodError(history, (io, x...,)))
history(; kernel=_default_kernel) = history(1:kernel.n; kernel)
"""
    history([io], [indices...])

The `history()` function prints all of the input history stored in
the running IJulia notebook in a format convenient for copying.

The optional `indices` argument is one or more indices or collections
of indices indicating a subset input cells to print.

The optional `io` argument is for specifying an output stream. The default
is `stdout`.
"""
history

#######################################################################
# Similar to the ipython kernel, we provide a mechanism by
# which modules can register thunk functions to be called after
# executing an input cell, e.g. to "close" the current plot in Pylab.

const _preexecute_hooks = Function[]
const _postexecute_hooks = Function[]
const _posterror_hooks = Function[]
# This is an internal hooks list to be called when stdio is being flushed by
# execute_request(). Used by the PythonCall extension.
const _flush_hooks = Function[]

function _pop_hook!(f, hooks)
    hook_idx = findlast(isequal(f), hooks)
    if isnothing(hook_idx)
        error("Could not find hook: $(f)")
    else
        splice!(hooks, hook_idx)
    end
end

"""
    push_postexecute_hook(f::Function)

Push a function `f()` onto the end of a list of functions to
execute after executing any notebook cell.
"""
push_postexecute_hook(f::Function) = push!(IJulia._postexecute_hooks, f)

"""
    pop_postexecute_hook(f::Function)

Remove a function `f()` from the list of functions to
execute after executing any notebook cell.
"""
pop_postexecute_hook(f::Function) = _pop_hook!(f, IJulia._postexecute_hooks)

"""
    push_preexecute_hook(f::Function)

Push a function `f()` onto the end of a list of functions to
execute before executing any notebook cell.
"""
push_preexecute_hook(f::Function) = push!(IJulia._preexecute_hooks, f)

"""
    pop_preexecute_hook(f::Function)

Remove a function `f()` from the list of functions to
execute before executing any notebook cell.
"""
pop_preexecute_hook(f::Function) = _pop_hook!(f, IJulia._preexecute_hooks)

# similar, but called after an error (e.g. to reset plotting state)
"""
    pop_posterror_hook(f::Function)

Remove a function `f()` from the list of functions to
execute after an error occurs when a notebook cell is evaluated.
"""
push_posterror_hook(f::Function) = push!(IJulia._posterror_hooks, f)

"""
    pop_posterror_hook(f::Function)

Remove a function `f()` from the list of functions to
execute after an error occurs when a notebook cell is evaluated.
"""
pop_posterror_hook(f::Function) = _pop_hook!(f, IJulia._posterror_hooks)

#######################################################################

# The user can call IJulia.clear_output() to clear visible output from the
# front end, useful for simple animations.  Using wait=true clears the
# output only when new output is available, for minimal flickering.
"""
    clear_output(wait=false)

Call `clear_output()` to clear visible output from the current notebook
cell.  Using `wait=true` clears the output only when new output is
available, which reduces flickering and is useful for simple animations.
"""
function clear_output(wait=false, kernel=_default_kernel)
    # flush pending stdio
    flush_all()
    empty!(kernel.displayqueue) # discard pending display requests
    send_ipython(kernel.publish[], kernel, msg_pub(kernel.execute_msg::Msg, "clear_output",
                                           Dict("wait" => wait)))
    kernel.stdio_bytes = 0 # reset output throttling
    return nothing
end


"""
    set_max_stdio(max_output::Integer)

Sets the maximum number of bytes, `max_output`, that can be written to stdout and
stderr before getting truncated. A large value here allows a lot of output to be
displayed in the notebook, potentially bogging down the browser.
"""
function set_max_stdio(max_output::Integer; kernel=_default_kernel)
    kernel.max_output_per_request[] = max_output
end

"""
    reset_stdio_count()

Reset the count of the number of bytes written to stdout/stderr. See
[`set_max_stdio`](@ref) for more details.
"""
function reset_stdio_count(kernel=_default_kernel)
    kernel.stdio_bytes = 0
end

#######################################################################

# These are stubs for the PythonCall extension

"""
    init_ipywidgets()

Initialize the integration with ipywidgets by setting up the right hooks to
allow ipywidgets to use IJulia comms.
"""
function init_ipywidgets end

"""
    init_ipython()

Initialize the integration with IPython by overriding its display system to call
Julia's `display()` instead.
"""
function init_ipython end

"""
    init_matplotlib(backend="module://ipympl.backend_nbagg")

Initialize the integration with matplotlib.
"""
function init_matplotlib end

include("init.jl")
include("hmac.jl")
include("eventloop.jl")
include("stdio.jl")
include("msg.jl")
include("display.jl")
include("magics.jl")
include("comm_manager.jl")
include("execute_request.jl")
include("handlers.jl")
include("heartbeat.jl")
include("inline.jl")
include("kernel.jl")
include("precompile.jl")

end # IJulia
