module PkgCompiler

using Pkg
using Libdl
using ProgressMeter

pkg_compiler_dir = "$(homedir())/.julia/pkgcompiler/v$(VERSION.major).$(VERSION.minor)"

global_snoop_file =  joinpath(pkg_compiler_dir, "snoop-file.jl")
snoop_file = nothing
snoop_file_io = nothing

function __init__()
    if haskey(ENV, "PkgCompilerEnabled") && ENV["PkgCompilerEnabled"] == "1"
        start_snooping()
    end
end

function start_snooping()
    global snoop_file_io
    if snoop_file_io === nothing
        global snoop_file = "$(tempname())-snoop.jl"
        mkpath(dirname(snoop_file))
        @info("Snooping -> $(snoop_file)")
        global snoop_file_io = open(snoop_file, "w")
        ccall(:jl_dump_compiles, Cvoid, (Ptr{Cvoid},), snoop_file_io.handle)
        atexit(finish_snooping)
    else
        @warn("Snooping is already running -> $(snoop_file)")
    end
end

# This function is for testing purposes
function set_global_snoop_file(file_name::AbstractString)
    global global_snoop_file = file_name
end

# This function is for testing purposes
function set_pkg_compiler_dir(new_diectory::AbstractString)
    global pkg_compiler_dir = new_diectory
end

function _flush()
    if snoop_file_io !== nothing
        flush(snoop_file_io)
    end
end

function _update_startup_file(content::AbstractString)
    begin_text = "Begin of PkgCompiler section"
    end_text = "End of PkgCompiler section"

    file_name = "$(homedir())/.julia/config/startup.jl"
    # Create new file if not existing
    if !isfile(file_name)
        @info("No startup file ($(file_name) will be created)")
        mkpath(dirname(file_name))
        touch(file_name)
    end

    lines = readlines(open(file_name))
    l = length(lines)

    # Test if PkgCompiler is already registered in ~/.julia/config/startup.jl
    start_idx = findfirst(x -> occursin(begin_text, x), lines)
    end_idx = findlast(x -> occursin(end_text, x), lines)

    if start_idx === nothing && end_idx === nothing
        # new text will be added at the end of the startup file
        start_idx = end_idx = l+1
    end
    if start_idx === nothing || end_idx === nothing
        @error("The startup file is malformed ($(file_name))")
        return
    end

    # Replace the old part from PkgCompiler
    open(file_name, "w") do file
        for i in 1:min(start_idx-1, l)
            write(file, lines[i], "\n")
        end
        if start_idx == l+1
            write(file, "\n")
        end
        write(file, "# $(begin_text) #####\n")
        write(file, content)
        write(file, "# $(end_text) #####\n")
        for i in (end_idx+1):l
            write(file, lines[i], "\n")
        end
    end
end

function register_snooping()
    content  = "ENV[\"PkgCompilerEnabled\"] = 1\n"
    content *= "try ; using PkgCompiler ; catch e ; end\n"
    _update_startup_file(content)
    __init__()
end

function deregister_snooping()
    _update_startup_file("\n")
    global snoop_file_io
    ccall(:jl_dump_compiles, Cvoid, (Ptr{Cvoid},), C_NULL)
    close(snoop_file_io)
    snoop_file_io = nothing
end

function finish_snooping()
    global snoop_file_io
    try
        if snoop_file_io === nothing
            @warn("No active snooping file")
            return 
        end
        ccall(:jl_dump_compiles, Cvoid, (Ptr{Cvoid},), C_NULL)
        close(snoop_file_io)
        snoop_file_io = nothing
        
        save_statements()
        rm("$(snoop_file)")
    catch e
        @warn(e)
    end
end

function save_statements()
    try
        _flush()       
        
        lines = readlines("$(snoop_file)")
        open("$(snoop_file).txt", "w") do file
            for line in lines
                sp = split(line, "\t")
                if length(sp) == 2
                    function_str = sp[2][2:end-1]
                    write(file, function_str * "\n")
                end          
            end
        end

        @info("Copy snooped function into $(global_snoop_file)")
        if isfile(global_snoop_file)
            run(`sort $(snoop_file).txt $(global_snoop_file) -o $(global_snoop_file).tmp`)
        else
            mkpath(pkg_compiler_dir)
            run(`sort $(snoop_file).txt -o $(global_snoop_file).tmp`)
        end
        run(pipeline(`uniq $(global_snoop_file).tmp`, stdout="$(global_snoop_file)"))
        rm("$(snoop_file).txt")
        rm("$(global_snoop_file).tmp")
    catch e
        @warn(e)
    end
end


# inspired by https://nextjournal.com/sdanisch/packagecompiler-talk
function build_system_image()
    save_statements()

    lib = abspath(joinpath(Sys.BINDIR, Base.PRIVATE_LIBDIR))
    img = unsafe_string(Base.JLOptions().image_file)   

    systemp = joinpath(pkg_compiler_dir, "sys.a")
    sysout = joinpath(pkg_compiler_dir, "sys.$(Libdl.dlext)")
    script = joinpath(@__DIR__, "compile.jl")

    # compile image
    @info("Compile $(@__DIR__)/compile.jl -> $(systemp)")
    @info("Using snoop file: $(global_snoop_file)")
    @time run(`julia -J$(img) --color=yes --startup-file=no -O3 -g 2 --output-o $(systemp) $(script) $(global_snoop_file)`)

    # From PackageCompiler.jl
    # `cc -shared '-DJULIAC_PROGRAM_LIBNAME="/home/petr/.julia/dev/PackageCompiler/sysimg/sys.so"' -o /home/petr/.julia/dev/PackageCompiler/sysimg/sys.so -Wl,--whole-archive /home/petr/.julia/dev/PackageCompiler/sysimg/sys.a -Wl,--no-whole-archive -std=gnu99 -I/home/petr/software/julia/julia-1.2.0/include/julia -DJULIA_ENABLE_THREADING=1 -fPIC -L/home/petr/software/julia/julia-1.2.0/lib -Wl,--export-dynamic -Wl,-rpath,/home/petr/software/julia/julia-1.2.0/lib -Wl,-rpath,/home/petr/software/julia/julia-1.2.0/lib/julia -ljulia -m64 -O3`

    # Link it
    @info("Link system image")
    # Prevent compiler from stripping all symbols from the shared lib.
    if Sys.isapple()
        temp_part = `-all_load $systemp`
    else
        temp_part = `-Wl,--whole-archive $systemp -Wl,--no-whole-archive`
    end
    part_end = ``
    if Sys.isapple()
        part_end = ` -Wl,-install_name,@rpath/$sysout`
    elseif Sys.iswindows()
        part_end = ` -Wl,--export-all-symbols`
    end
    @time run(`c++ -m64 -shared -DJULIAC_PROGRAM_LIBNAME=\"$sysout\" -fPIC -L$(lib) -L$(dirname(lib)) -o $(sysout) $(temp_part) $(part_end) `);
    @info("system image is prepared, you can run it by:\njulia -J $(sysout)")
end

end