// AssetManager.cpp
#include "AssetManager.h"
#include "assets/Texture2D.h"
#include "Logger.h"

#include <stb/stb_image.h>
#include <GL/glew.h>
#include <yaml-cpp/yaml.h>
#include <fstream>
#include <ranges>

#include "Profiler.h"

namespace fs = std::filesystem;

namespace OX
{
    // statics
    std::unordered_map<std::string, std::shared_ptr<Asset> > AssetManager::s_LoadedAssets;
    std::unordered_map<std::string, AssetManager::AssetMetadata> AssetManager::s_MetadataMap;

    std::atomic<size_t> OX::AssetManager::s_TotalTexturesToLoad{0};
    std::atomic<size_t> OX::AssetManager::s_LoadedTexturesCount{0};


    // Map from an actual file path to virtual ID
    std::unordered_map<std::string, std::string> AssetManager::s_PathToID;
    std::mutex AssetManager::s_AssetMutex;
    std::mutex AssetManager::s_TreeMutex;

    std::shared_ptr<ResourceTreeNode> AssetManager::s_FileTree = std::make_shared<ResourceTreeNode>();

    std::filesystem::path AssetManager::s_ProjectRoot;
    std::atomic<bool> AssetManager::s_Scanning{true};
    std::thread AssetManager::s_ScanThread;

    std::queue<AssetManager::PendingTexture> AssetManager::s_TextureQueue;
    std::mutex AssetManager::s_QueueMutex;

    void AssetManager::Init(const std::string &projectRoot)
    {
        s_ProjectRoot = fs::absolute(projectRoot);
        s_Scanning = true;

        s_FileTree = std::make_shared<ResourceTreeNode>();
        s_FileTree->name = "res://";
        s_FileTree->path = "res://";
        s_FileTree->isDirectory = true;
        s_TotalTexturesToLoad.store(0);
        s_LoadedTexturesCount.store(0);
        s_ScanThread = std::thread(BackgroundScan);
    }

    void AssetManager::Shutdown()
    {
        s_Scanning = false;

        if (s_ScanThread.joinable()) {
            s_ScanThread.join();
        }
    }

    void AssetManager::Rescan()
    {
        if (s_Scanning) return;
        s_Scanning = true;

        s_FileTree = std::make_shared<ResourceTreeNode>();
        s_FileTree->name = "res://";
        s_FileTree->path = "res://";
        s_FileTree->isDirectory = true;
        s_TotalTexturesToLoad.store(0);
        s_LoadedTexturesCount.store(0);
        s_ScanThread = std::thread(BackgroundScan);




    }

    void AssetManager::BackgroundScan()
    {
        namespace fs = std::filesystem;

        size_t totalTextures = 0;
        for (auto &entry: fs::recursive_directory_iterator(s_ProjectRoot)) {
            if (!s_Scanning) break;
            if (!entry.is_directory()) {
                if (DetectAssetType(entry.path()) == "texture2D")
                    ++totalTextures;
            }
        }
        s_TotalTexturesToLoad.store(totalTextures, std::memory_order_relaxed);
        s_LoadedTexturesCount.store(0, std::memory_order_relaxed);


        while (s_Scanning) {
            std::vector<fs::directory_entry> allEntries;
            allEntries.reserve(1024);
            for (auto &entry: fs::recursive_directory_iterator(s_ProjectRoot)) {
                if (!s_Scanning) break;
                allEntries.push_back(entry);
            }

            std::vector<std::pair<std::string, AssetMetadata> > newMetadata;
            newMetadata.reserve(256);


            for (auto &entry: allEntries) {
                if (!s_Scanning) break;

                const auto &path = entry.path();
                std::string resPath = MakeVirtualPath(path);

                if (entry.is_directory()) {
                    AddToTree(resPath, true);
                    continue;
                } {
                    std::lock_guard<std::mutex> lock(s_AssetMutex);
                    if (s_MetadataMap.count(resPath))
                        continue;
                    s_MetadataMap.emplace(resPath, AssetMetadata{"", ""});
                }
                AddToTree(resPath, false);

                if (DetectAssetType(path) == "texture2D") {
                    int w, h, ch;
                    stbi_set_flip_vertically_on_load(1);
                    unsigned char *data = stbi_load(path.string().c_str(), &w, &h, &ch, 0);
                    if (!data) continue;

                    newMetadata.emplace_back(resPath, AssetMetadata{"texture2D", path.string()}); {
                        std::lock_guard<std::mutex> qlock(s_QueueMutex);
                        s_TextureQueue.push({resPath, w, h, ch, data});
                    }
                }
            }

            if (!newMetadata.empty()) {
                std::lock_guard<std::mutex> lock(s_AssetMutex);
                for (auto &kv: newMetadata)
                    s_MetadataMap[kv.first] = kv.second;
            }

            if (s_Scanning)
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }

    }


    void AssetManager::Tick()
    {
        OX_PROFILE_FUNCTION();
        std::lock_guard<std::mutex> qlock(s_QueueMutex);
        while (!s_TextureQueue.empty()) {
            auto pending = s_TextureQueue.front();
            s_TextureQueue.pop();
            s_LoadedTexturesCount.fetch_add(1, std::memory_order_relaxed);


            GLuint texID;
            glGenTextures(1, &texID);
            glBindTexture(GL_TEXTURE_2D, texID);
            GLenum fmt = (pending.channels == 4 ? GL_RGBA : GL_RGB);
            glTexImage2D(GL_TEXTURE_2D, 0, fmt, pending.width, pending.height, 0, fmt, GL_UNSIGNED_BYTE, pending.data);
            glGenerateMipmap(GL_TEXTURE_2D);
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR);
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
            stbi_image_free(pending.data);

            auto tex = std::make_shared<Texture2D>();
            tex->SetFromGL(texID, pending.width, pending.height);

            std::lock_guard<std::mutex> lock(s_AssetMutex);
            // Store asset by ID
            s_LoadedAssets[pending.id] = tex;
            // Also map original file path to this ID
            auto meta = s_MetadataMap[pending.id];
            s_PathToID[meta.absolutePath] = pending.id;
        }
    }

    // Get asset by virtual ID or file path
    std::shared_ptr<Asset> AssetManager::Get(const std::string &keyOrPath)
    {
        std::string key = keyOrPath;
        // Normalize accidental triple slashes in resource path: "res:///..." -> "res://..."
        const std::string triple = "res:///";
        if (key.rfind(triple, 0) == 0) {
            key = std::string("res://") + key.substr(triple.size());
        }

        std::lock_guard<std::mutex> lock(s_AssetMutex);

        // 1) Direct ID lookup
        auto it = s_LoadedAssets.find(key);
        if (it != s_LoadedAssets.end()) {
            return it->second;
        }
        // 2) Path-to-ID mapping
        auto pit = s_PathToID.find(key);
        if (pit != s_PathToID.end()) {
            auto ait = s_LoadedAssets.find(pit->second);
            if (ait != s_LoadedAssets.end()) {
                return ait->second;
            }
        }
        // 3) Convert filesystem path to virtual and lookup
        try {
            std::string virt = MakeVirtualPath(fs::absolute(key));
            // Normalize if virt had leading slash
            if (!virt.empty() && virt.front() == '/') virt.erase(0, 1);
            auto vit = s_LoadedAssets.find(virt);
            if (vit != s_LoadedAssets.end()) {
                return vit->second;
            }
        } catch (...) {
            // ignore
        }
        return nullptr;
    }


    std::shared_ptr<ResourceTreeNode> AssetManager::GetFileTree()
    {
        return s_FileTree;
    }

    void AssetManager::SaveAssetPack(const std::string &outputPath)
    {
        YAML::Emitter out;
        out << YAML::BeginMap << YAML::Key << "assets" << YAML::Value << YAML::BeginSeq; {
            std::lock_guard<std::mutex> lock(s_AssetMutex);
            for (auto &[id, meta]: s_MetadataMap) {
                fs::path rel = fs::relative(meta.absolutePath, s_ProjectRoot);
                out << YAML::BeginMap
                        << YAML::Key << "id" << YAML::Value << id
                        << YAML::Key << "type" << YAML::Value << meta.type
                        << YAML::Key << "path" << YAML::Value << rel.generic_string()
                        << YAML::EndMap;
            }
        }
        out << YAML::EndSeq << YAML::EndMap;

        std::ofstream fout(outputPath);
        fout << out.c_str();
        Logger::LogInfo("Saved asset pack: %s", outputPath.c_str());
    }

    std::string AssetManager::MakeVirtualPath(const fs::path &full)
    {
        auto rel = fs::relative(full, s_ProjectRoot);
        return "res://" + rel.generic_string();
    }

    void AssetManager::AddToTree(const std::string &resPath, bool isDir)
    {
        std::vector<std::string> parts;
        size_t pos = 0, next;
        while ((next = resPath.find('/', pos)) != std::string::npos) {
            if (next != pos) parts.push_back(resPath.substr(pos, next - pos));
            pos = next + 1;
        }
        if (pos < resPath.size()) parts.push_back(resPath.substr(pos));
        std::lock_guard<std::mutex> lock(AssetManager::s_TreeMutex);

        auto current = s_FileTree;
        for (size_t i = 1; i < parts.size(); ++i) {
            auto &name = parts[i];
            auto it = std::find_if(current->children.begin(), current->children.end(),
                                   [&](auto &c) { return c->name == name; });
            if (it == current->children.end()) {
                auto node = std::make_shared<ResourceTreeNode>();
                node->name = name;
                node->path = current->path + "/" + name;
                node->isDirectory = (i + 1 < parts.size()) || isDir;
                current->children.push_back(node);
                current = node;
            } else {
                current = *it;
            }
        }
    }

    std::string AssetManager::DetectAssetType(const fs::path &path)
    {
        auto e = path.extension().string();
        if (e == ".png" || e == ".jpg" || e == ".jpeg") return "texture2D";
        return "";
    }
} // namespace OX
