/**
 * \file imperative/src/impl/profiler/chrome_timeline.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#if defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
#include <unistd.h>
#elif defined(_WIN32)
#include <process.h>
#else
#error Unsupported platform
#endif

#include "nlohmann/json.hpp"

#include "megbrain/utils/debug.h"

#include "./formats.h"
#include "./states.h"

namespace mgb::imperative::profiler {


class ChromeTraceEvent {
public:
    ChromeTraceEvent& name(std::string name) {
        m_name = std::move(name);
        return *this;
    }
    ChromeTraceEvent& tid(uint64_t tid) {
        m_tid = std::move(tid);
        return *this;
    }
    ChromeTraceEvent& cat(std::string cat) {
        m_cat = std::move(cat);
        return *this;
    }
    ChromeTraceEvent& scope(std::string scope) {
        m_scope = std::move(scope);
        return *this;
    }
    ChromeTraceEvent& pid(uint64_t pid) {
        m_pid = pid;
        return *this;
    }
    ChromeTraceEvent& id(uint64_t id) {
        m_id = id;
        return *this;
    }
    ChromeTraceEvent& idx(uint64_t idx) {
        m_idx = idx;
        return *this;
    }
    template <typename TDuration>
    ChromeTraceEvent& ts(TDuration ts) {
        m_ts = std::chrono::duration_cast<std::chrono::duration<double, std::micro>>(ts).count();
        return *this;
    }
    ChromeTraceEvent& dur(uint64_t dur) {
        m_dur = dur;
        return *this;
    }
    ChromeTraceEvent& ph(char ph) {
        m_ph = ph;
        return *this;
    }
    ChromeTraceEvent& bp(char bp) {
        m_bp = bp;
        return *this;
    }
    ChromeTraceEvent& args(nlohmann::json args) {
        m_args = std::move(args);
        return *this;
    }
    ChromeTraceEvent& arg(std::string key, std::string value) {
        m_args[key] = value;
        return *this;
    }
    ChromeTraceEvent& arg(std::string key, double value) {
        m_args[key] = value;
        return *this;
    }
    ChromeTraceEvent& stack(Trace trace) {
        m_stack = std::move(trace);
        return *this;
    }

    nlohmann::json to_json() const {
        nlohmann::json result;
        auto prop_str = [&](auto key, auto value) {
            if (value.empty()) {
                return;
            }
            result[key] = value;
        };
        auto prop_num = [&](auto key, auto value) {
            if (!value) {
                return;
            }
            result[key] = value.value();
        };
        auto prop_char = [&](auto key, auto value) {
            if (!value) {
                return;
            }
            result[key] = std::string{} + value.value();
        };
        prop_str("name", m_name);
        prop_str("cat", m_cat);
        prop_str("scope", m_scope);
        prop_num("tid", m_tid);
        prop_num("pid", m_pid);
        prop_num("id", m_id);
        prop_num("idx", m_idx);
        prop_num("ts", m_ts);
        prop_num("dur", m_dur);
        prop_char("ph", m_ph);
        prop_char("bp", m_bp);
        if (!m_args.empty()) {
            result["args"] = m_args;
        }
        if (m_stack) {
            nlohmann::json stack;
            for (auto&& frame: m_stack->frames()) {
                stack.push_back(ssprintf("%s%ld", frame.node->name().c_str(), frame.version));
            }
            std::reverse(stack.begin(), stack.end());
            result["stack"] = stack;
        }
        return result;
    }
private:
    std::string m_name;
    std::string m_cat;
    std::string m_scope;

    std::optional<uint64_t> m_tid;
    std::optional<uint64_t> m_pid;
    std::optional<uint64_t> m_id;
    std::optional<uint64_t> m_idx;
    std::optional<double> m_ts;
    std::optional<uint64_t> m_dur;
    std::optional<char> m_ph;
    std::optional<char> m_bp;
    nlohmann::json m_args;
    std::optional<Trace> m_stack;
};

class ChromeTraceEvents {
public:
    ChromeTraceEvent& new_event() {
        m_content.emplace_back();
        return m_content.back();
    }

    std::string& metadata(std::string key) {
        return m_metadata[key];
    }

    nlohmann::json to_json() const {
        nlohmann::json result;
        nlohmann::json event_list;
        nlohmann::json metadata;
        for (auto&& event: m_content) {
            event_list.push_back(event.to_json());
        }
        for (auto&& [key, value]: m_metadata) {
            metadata[key] = value;
        }
        result["traceEvents"] = event_list;
        result["metadata"] = metadata;
        return result;
    }

    std::string to_string() const {
        auto json = to_json();
        return "{" "\"traceEvents\":" + nlohmann::to_string(json["traceEvents"]) + ","
            "\"metadata\":" + nlohmann::to_string(json["metadata"]) + "}";
    }
private:
    std::vector<ChromeTraceEvent> m_content;
    std::unordered_map<std::string, std::string> m_metadata;
};


struct ChromeTimelineEventVisitor: EventVisitor<ChromeTimelineEventVisitor> {
    ChromeTraceEvents trace_events;
    decltype(getpid()) pid = getpid();
    std::string pid_str = std::to_string(pid);

    ChromeTraceEvent& new_event(std::string name, char ph, size_t tid, profiler::HostTime time) {
        return trace_events.new_event().name(name).ph(ph).pid(pid).tid(tid).ts(since_start(time));
    }

    ChromeTraceEvent& new_host_event(std::string name, char ph) {
        return trace_events.new_event().name(name).ph(ph).pid(pid).tid(to_tid(current->tid)).ts(since_start(current->time));
    };

    ChromeTraceEvent& new_device_event(std::string name, char ph, CompNode device) {
        using namespace std::literals::chrono_literals;
        auto time = since_start(to_device_time(current->time, device));
        return trace_events.new_event().name(name).ph(ph).pid(pid).tid(to_tid(device)).ts(time);
    };

    const char* to_cstr(TensorCommandKind kind) {
        switch(kind) {
            case TensorCommandKind::Put:
                return "Put";
            case TensorCommandKind::Drop:
                return "Drop";
            case TensorCommandKind::Del:
                return "Del";
            case TensorCommandKind::SwapIn:
                return "SwapIn";
            case TensorCommandKind::SwapOut:
                return "SwapOut";
            case TensorCommandKind::RecFree:
                return "RecFree";
            case TensorCommandKind::ReGen:
                return "ReGen";
            case TensorCommandKind::GetValue:
                return "GetValue";
        }
        return "UnknownCommand";
    }

    template <typename TEvent>
    void visit_event(const TEvent &event) {
        if constexpr (std::is_same_v<TEvent, OpDispatchEvent>) {
            new_host_event("OpDispatch", 'B');
            new_host_event(pid_str, 's')
                    .cat("OpDispatch")
                    .id(event.op_id)
                    .scope(pid_str);
            new_host_event("OpDispatch", 'E').args(current_op->detail());
        } else if constexpr (std::is_same_v<TEvent, OpExecuteEvent>) {
            mgb_assert(event.op_id != 0);
            current_op->execute_begin = current->time;
            new_host_event(current_op->name, 'B');
            new_host_event(pid_str, 't')
                    .cat("OpDispatch")
                    .id(current_op->id)
                    .scope(pid_str);
        } else if constexpr (std::is_same_v<TEvent, OpExecuteFinishEvent>) {
            current_op->execute_end = current->time;
            new_host_event(current_op->name, 'E')
                    .args(current_op->detail());
        } else if constexpr (std::is_same_v<TEvent, KernelLaunchEvent>) {
            new_host_event(pid_str, 's')
                    .id(event.kernel_id)
                    .cat("KernelLaunch")
                    .scope(pid_str);
            new_device_event(current_op->name, 'B', event.device)
                    .cat("Kernel");
            new_device_event(pid_str, 'f', event.device)
                    .id(event.kernel_id)
                    .bp('e')
                    .cat("KernelLaunch")
                    .scope(pid_str);
        } else if constexpr (std::is_same_v<TEvent, KernelLaunchFinishEvent>) {
            new_device_event(current_op->name, 'E', event.device)
                    .cat("Kernel")
                    .args(current_op->detail());
        } else if constexpr (std::is_same_v<TEvent, TensorProduceEvent>) {
            if (current_tensor->living_time != profiler::Duration::zero()) {
                new_host_event(pid_str, 's')
                    .id(event.tensor_id)
                    .cat("TensorLink")
                    .scope(pid_str);
            } else {
                new_host_event(pid_str, 't')
                    .id(event.tensor_id)
                    .cat("TensorLink")
                    .scope(pid_str);
            }
        } else if constexpr (std::is_same_v<TEvent, TensorUsageEvent>) {
            new_host_event(pid_str, 't')
                    .id(event.tensor_id)
                    .cat("TensorLink")
                    .scope(pid_str);
        } else if constexpr (std::is_same_v<TEvent, TensorReleaseEvent>) {
            current_tensor->living_time += current->time - current_tensor->produced;
            current_tensor->produced = {};
            new_host_event(pid_str, 't')
                    .id(event.tensor_id)
                    .cat("TensorLink")
                    .scope(pid_str);
        } else if constexpr (std::is_same_v<TEvent, TensorEraseEvent>) {
            if (current_tensor->living_time != profiler::Duration::zero()) {
                new_host_event(pid_str, 'f')
                        .id(event.tensor_id)
                        .bp('e')
                        .cat("TensorLink")
                        .scope(pid_str);
            }
        } else if constexpr (std::is_same_v<TEvent, TensorGetPropEvent>) {
            new_host_event("TensorGetProp", 'X')
                    .dur(0).args(current_tensor->detail(current->time));
        } else if constexpr (std::is_same_v<TEvent, TensorWaitPropEvent>) {
            new_host_event("TensorWaitProp", 'B');
        } else if constexpr (std::is_same_v<TEvent, TensorWaitPropFinishEvent>) {
            new_host_event(pid_str, 'f')
                    .id(event.tensor_id)
                    .bp('e')
                    .cat("TensorProp")
                    .scope(pid_str);
            new_host_event("TensorWaitProp", 'E')
                    .args(current_tensor->detail(current->time));
        } else if constexpr (std::is_same_v<TEvent, TensorNotifyPropEvent>) {
            new_host_event(pid_str, 's')
                    .id(event.tensor_id)
                    .cat("TensorProp")
                    .scope(pid_str);
        } else if constexpr (std::is_same_v<TEvent, SampleDeviceFinishEvent>) {
            std::string device_name = event.device.locator().to_string();
            new_host_event("memory", 'C')
                    .arg(ssprintf("%s_alloc_mem", device_name.c_str()), event.total_memory - event.free_memory);
        } else if constexpr (std::is_same_v<TEvent, TensorCommandEvent>) {
            new_host_event(ssprintf("%s %zu", to_cstr(event.kind), event.tensor_id), 'B');
        } else if constexpr (std::is_same_v<TEvent, TensorCommandFinishEvent>) {
            new_host_event(ssprintf("%s %zu", to_cstr(event.kind), event.tensor_id), 'E')
                    .args(current_tensor->detail(current->time));
        } else if constexpr (std::is_same_v<TEvent, ScopeEvent>) {
            new_host_event(event.name, 'B');
        } else if constexpr (std::is_same_v<TEvent, ScopeFinishEvent>) {
            new_host_event(event.name, 'E');
        } else if constexpr (std::is_same_v<TEvent, OpInputEvent>) {
            new_host_event(ssprintf("Input %zu", event.tensor_id), 'B')
                    .args(current_tensor->detail(current->time));
        } else if constexpr (std::is_same_v<TEvent, OpInputFinishEvent>) {
            new_host_event(ssprintf("Input %zu", event.tensor_id), 'E')
                    .args(current_tensor->detail(current->time));
        } else if constexpr (std::is_same_v<TEvent, OpOutputEvent>) {
            new_host_event(ssprintf("Output %zu", event.tensor_id), 'B')
                    .args(current_tensor->detail(current->time));
        } else if constexpr (std::is_same_v<TEvent, OpOutputFinishEvent>) {
            new_host_event(ssprintf("Output %zu", event.tensor_id), 'E')
                    .args(current_tensor->detail(current->time));
        } else if constexpr (std::is_same_v<TEvent, StartProfileEvent>) {
            new_host_event("StartProfile", 'B');
        } else if constexpr (std::is_same_v<TEvent, StartProfileFinishEvent>) {
            new_host_event("StartProfile", 'E');
        } else if constexpr (std::is_same_v<TEvent, StopProfileEvent>) {
            new_host_event("StopProfile", 'B');
        } else if constexpr (std::is_same_v<TEvent, StopProfileFinishEvent>) {
            new_host_event("StopProfile", 'E');
        } else if constexpr (std::is_same_v<TEvent, CustomEvent>) {
            new_host_event(event.title, 'B');
        } else if constexpr (std::is_same_v<TEvent, CustomFinishEvent>) {
            new_host_event(event.title, 'E');
        } else if constexpr (std::is_same_v<TEvent, AutoEvictEvent>) {
            new_host_event("AutoEvict", 'B');
        } else if constexpr (std::is_same_v<TEvent, AutoEvictFinishEvent>) {
            new_host_event("AutoEvict", 'E');
        } else if constexpr (std::is_same_v<TEvent, HostToDeviceEvent>) {
            new_device_event("HostToDevice", 'B', event.device);
        } else if constexpr (std::is_same_v<TEvent, HostToDeviceFinishEvent>) {
            new_device_event("HostToDevice", 'E', event.device)
                    .arg("shape", event.layout.TensorShape::to_string())
                    .arg("dtype", event.layout.dtype.name())
                    .arg("nr_elements", event.layout.total_nr_elems())
                    .arg("device", event.device.to_string());
        }
    }

    void notify_counter(std::string key, int64_t old_val, int64_t new_val) {
        new_host_event(key, 'C').arg("value", new_val);
    }

    void name_threads(Profiler::thread_dict_t thread_dict) {
        for (auto&& [tid, tname]: thread_dict) {
            trace_events.new_event()
                    .name("thread_name")
                    .pid('M')
                    .tid(to_tid(tid))
                    .arg("name", tname);
        }
    }
};


void dump_chrome_timeline(std::string filename, Profiler::bundle_t result){
    ChromeTimelineEventVisitor visitor;
    visitor.process_events(result);
    visitor.name_threads(result.thread_dict);
    auto trace_events = std::move(visitor.trace_events);
    trace_events.metadata("localTime") = std::to_string(result.start_at.time_since_epoch().count());
    std::string json_repr = trace_events.to_string();
    mgb::debug::write_to_file(filename.c_str(), json_repr);
}

}