#include "Profiler.h"
#include <sstream>
#include <iomanip>
#include <algorithm>
#include <functional>
#include <unordered_set>
namespace OX
{
    void Profiler::BeginFrame() {
        s_samples.clear();
        s_startTimes.clear();
        s_formatted.clear();
        s_savedTree.clear();
        s_stack.clear();
    }

    void Profiler::StartSample(const std::string& name) {
        s_startTimes[name] = std::chrono::high_resolution_clock::now();

        if (!s_stack.empty()) {
            const std::string& parent = s_stack.back();

            if (parent != name) { // Prevent self-cycle
                auto& parentSample = s_samples[parent];
                if (std::find(parentSample.children.begin(), parentSample.children.end(), name) == parentSample.children.end()) {
                    parentSample.children.push_back(name);
                }
            }
        }


        s_stack.push_back(name);
    }

    void Profiler::EndSample(const std::string& name) {
        auto end = std::chrono::high_resolution_clock::now();
        auto start = s_startTimes[name];
        double elapsed = std::chrono::duration<double, std::milli>(end - start).count();

        auto& sample = s_samples[name];
        sample.totalTime += elapsed;
        sample.callCount++;

        if (!s_stack.empty() && s_stack.back() == name) {
            s_stack.pop_back();
        }
    }

    void Profiler::EndFrame() {
        s_formatted.clear();
        s_savedTree.clear();

        std::unordered_set<std::string> childSet;
        for (const auto& [name, sample] : s_samples) {
            for (const auto& child : sample.children) {
                childSet.insert(child);
            }
        }

        std::vector<std::string> roots;
        for (const auto& [name, sample] : s_samples) {
            if (childSet.find(name) == childSet.end()) {
                roots.push_back(name);
            }
        }

        std::function<SavedSample(const std::string&, int)> recurse;
        recurse = [&](const std::string& name, int depth) {
            const auto& sample = s_samples[name];

            std::ostringstream ss;
            ss << std::string(depth * 2, ' ') << name << " (" << sample.callCount << "): "
               << std::fixed << std::setprecision(2) << sample.totalTime << " ms";
            s_formatted.push_back(ss.str());

            SavedSample node;
            node.name = name;
            node.totalTime = sample.totalTime;
            node.callCount = sample.callCount;

            for (const auto& child : sample.children) {
                node.children.push_back(recurse(child, depth + 1));
            }

            return node;
        };

        for (const auto& root : roots) {
            s_savedTree.push_back(recurse(root, 0));
        }

        s_lastFormatted = s_formatted;
        s_lastSavedTree = s_savedTree;
    }

    const std::vector<std::string>& Profiler::GetFormattedData() {
        return s_lastFormatted;
    }

    const std::vector<Profiler::SavedSample>& Profiler::GetSavedTree() {
        return s_lastSavedTree;
    }

    ProfileScope::ProfileScope(const std::string& name) : m_name(name) {
        Profiler::StartSample(name);
    }

    ProfileScope::~ProfileScope() {
        Profiler::EndSample(m_name);
    }
}