local c = require "skynet.core"
local skynet = require("skynet")
local new_span = require("tracing.span").new
local json = require("json")
local clua = require("clua")
local coroutine = coroutine
local tostring = tostring
local pcall = pcall
local next = next
local pairs = pairs
-------------------------------------------------------

local skynet_send = skynet.send
local skynet_call = skynet.call
local skynet_dispatch = skynet.dispatch
local skynet_register_protocol = skynet.register_protocol

skynet.PTYPE_MESSAGE_COLLOCTOR = 13 -- use for msg colloctor

-- messageColloctor
-- 消息追踪地址
local message_source_head = {}
local message_collector_head = {}
local message_running_head = {}
local message_self = false

local hook_thread
local hook_addr = {}
local hook_flag

local function set_message_source_head(source, trace_id, parent_id)
    message_source_head[source] = {
        trace_id = trace_id,
        parent_id = parent_id,
    }
end

local hook_cmd = {}

function hook_cmd.send_collect_info(source, trace_id, parent_id)
    set_message_source_head(source, trace_id, parent_id)
end

function hook_cmd.start_hook(_)
    if hook_thread == nil then
        hook_flag = true
    end
end

function hook_cmd.stop_hook(_)
    skynet.stop_hook_collector()
end

local function getSpan()
    local running_thread = tostring(coroutine.running())
    return message_running_head[running_thread]
end

skynet_register_protocol({
    name = "message_collector",
    id = skynet.PTYPE_MESSAGE_COLLOCTOR,
    pack = skynet.pack,
    unpack = skynet.unpack,
    dispatch = function(_, source, cmd, ...)
        local f = hook_cmd[cmd]
        if f then
            f(source, ...)
            -- else
            -- Log.d(cmd)
        end

    end,
})

function skynet.send(addr, ...)
    local span = getSpan()
    if span and not message_self then
        skynet.send_collect_info(addr, span)
    end
    if hook_thread == coroutine.running() then
        skynet.send_start_hook(addr)
    end
    return skynet_send(addr, ...)
end

function skynet.call(addr, ...)
    local span = getSpan()
    if span and not message_self then
        skynet.send_collect_info(addr, span)
    end

    if hook_thread == coroutine.running() then
        skynet.send_start_hook(addr)
    end
    return skynet_call(addr, ...)
end

function skynet.dispatch(typename, func)
    if typename ~= skynet.PTYPE_MESSAGE_COLLOCTOR then
        local message_collect_dipatch = function(session, source, cmd, ...)
            local running_thread = tostring(coroutine.running())
            local head = message_source_head[source]
            if head then
                message_source_head[source] = nil
                local span = new_span(cmd, head.trace_id, head.parent_id)
                message_running_head[running_thread] = span
            elseif (cmd and message_collector_head[cmd]) then
                local span = new_span(cmd)
                message_running_head[running_thread] = span
            end
            if hook_flag then
                skynet.start_hook_collector()
                hook_flag = false
            end
            func(session, source, cmd, ...)
            local span = message_running_head[running_thread]
            if span then
                message_running_head[running_thread] = nil
                local param = {...}
                if next(param) then
                    local ok, jsonValue = pcall(json.encode, param)
                    if ok then
                        span:set_tag("param", jsonValue)
                    end
                end
                span:finish()
            end
        end

        return skynet_dispatch(typename, message_collect_dipatch)
    end
end

function skynet.start_message_collector(funcName)
    message_collector_head[funcName] = true
end

function skynet.stop_message_collector(funcName)
    message_collector_head[funcName] = nil
end

function skynet.start_hook_collector(filename)
    if hook_thread == nil then
        -- Log.d("start_hook_collector", SERVICE_NAME)
        hook_thread = coroutine.running()
        filename = filename or string.format("%s_%d.cov", SERVICE_NAME, os.time())
        clua.start(hook_thread, filename)
    end
end

function skynet.stop_hook_collector()
    if hook_thread then
        -- Log.d("stop_hook_collector", SERVICE_NAME)
        clua.stop(hook_thread)
        hook_thread = nil
        for addr in pairs(hook_addr) do
            skynet.send_trace_msg(addr, "stop_hook")
        end
    end
end

function skynet.send_start_hook(addr)
    if not hook_addr[addr] then
        hook_addr[addr] = true
        -- Log.d("skynet.send_start_hook", SERVICE_NAME, coroutine.running(), addr)
        skynet.send_trace_msg(addr, "start_hook")
    end
end

function skynet.set_message_self()
    message_self = true
end

function skynet.send_collect_info(addr, span)
    skynet.send_trace_msg(addr, "send_collect_info", span.traceId, span.id)
end

function skynet.send_trace_msg(addr, ...)
    c.send(addr, skynet.PTYPE_MESSAGE_COLLOCTOR, 0, skynet.pack(...))
end

return skynet
