#include "model_download_manager.h"
#include "user_data_manager.h"
#include "hid_handler.h"
#include "utils.h"

#include <windows.h>
#include <shlobj.h>
#include <string>
#include <filesystem>
#include <iostream>

ModelDownloadManager modelManager;

std::wstring GetRoamingAppData()
{
    PWSTR path = nullptr;
    std::wstring result;

    if (SUCCEEDED(SHGetKnownFolderPath(FOLDERID_RoamingAppData, 0, NULL, &path)))
    {
        result = path;
        CoTaskMemFree(path);
    }

    return result;
}

// Build a master list of models from local folders and fetched remote models
void ModelDownloadManager::RebuildModelList(std::vector<std::string> trainedModels)
{
    UI_Run_State* rs = ui.get_runtime_state();
    std::vector<std::string> excludedList = rs->et_removed_models;

    masterModelList.clear();

    for (std::string t_model : trainedModels)
    {
        std::string excludedModel = "";

        if (rs)
        {
            UI_Run_State& rsRef = *rs;
            excludedModel = utils.et_format_model_name(rsRef, t_model, true);
        }

        if (std::find(excludedList.begin(), excludedList.end(), excludedModel) != excludedList.end())
        {
            continue;
        }

        ETModel newModel = {
            std::wstring(t_model.begin(), t_model.end()),
            modelPath,
            (int)Model_Types::Remote,
        };

        masterModelList.push_back(newModel);
    }

    std::filesystem::path legPath = GetRoamingAppData() + legacyModelPath;

    // Validate legacy directory
    if (std::filesystem::exists(legPath) && std::filesystem::is_directory(legPath))
    {
        for (const auto& entry : std::filesystem::recursive_directory_iterator(legPath))
        {
            if (entry.is_regular_file() && entry.path().extension() == ".model")
            {
                std::string excludedModel = "";

                if (rs)
                {
                    UI_Run_State& rsRef = *rs;
                    excludedModel = utils.et_format_model_name(rsRef, entry.path().filename().string(), true);
                }

                if (std::find(excludedList.begin(), excludedList.end(), excludedModel) != excludedList.end())
                {
                    continue;
                }

                ETModel newModel = {
                    entry.path().filename().wstring(),
                    legPath,
                    (int)Model_Types::Legacy,
                };

                masterModelList.push_back(newModel);
            }
        }
    }

    // Validate relative directory
    if (std::filesystem::exists(modelPath) && std::filesystem::is_directory(modelPath))
    {
        for (const auto& entry : std::filesystem::recursive_directory_iterator(modelPath))
        {
            if (entry.is_regular_file() && entry.path().extension() == ".model")
            {
                std::string excludedModel = "";

                if (rs)
                {
                    UI_Run_State& rsRef = *rs;
                    excludedModel = utils.et_format_model_name(rsRef, entry.path().filename().string(), true);
                }

                if (std::find(excludedList.begin(), excludedList.end(), excludedModel) != excludedList.end())
                {
                    continue;
                }

                ETModel newModel = {
                    entry.path().filename().wstring(),
                    modelPath,
                    (int)Model_Types::Local,
                };

                masterModelList.push_back(newModel);
            }
        }
    }
}

ModelCheck ModelDownloadManager::ModelExists(std::wstring modelName)
{
    namespace fs = std::filesystem;

    const fs::path original(modelName);
    const std::string fileName = original.filename().string();
    const std::string parentFolderName =
        original.has_parent_path() ? original.parent_path().filename().string() : std::string{};

    const std::string downloadName =
        parentFolderName.empty() ? fileName : (parentFolderName + "_" + fileName);

    const fs::path modelsLocalPath(modelPath);
    const fs::path localPath = modelsLocalPath / downloadName;

    const fs::path modelsLegacyPath(GetRoamingAppData() + legacyModelPath);
    const fs::path legacyPath = modelsLegacyPath / downloadName;

    bool localExists = fs::exists(localPath);
    bool legacyExists = fs::exists(legacyPath);

    return { localExists, legacyExists, localPath.string(), legacyPath.string() };
}

void ModelDownloadManager::OnModelSelected(ETModel model)
{
    std::thread([this, model]() {
        // Always make model directory
        std::filesystem::create_directories(modelPath);

        hid_handler.log << "Model selected: " << model.Name.c_str() << std::endl;
        userManager.SetModelLoadingError("");

        namespace fs = std::filesystem;

        // If path is already a valid local file, just signal ready.
        if ((Model_Types)model.Type != Model_Types::Remote)
        {
            if (!fs::exists(model.Location + L"/" + model.Name))
            {
                hid_handler.log << "Failed to find local model: " << (model.Location + L"/" + model.Name).c_str() << std::endl;
            }

            if (OnModelReady)
            {
                OnModelReady(utils.string_wide_to_narrow(model.Location + L"/" + model.Name), false);
            }
            return;
        }

        ModelCheck exists = ModelExists(model.Name);

        if (!exists.LocalExists && !exists.LegacyExists)
        {
            hid_handler.log << "Model not found locally, downloading..." << std::endl;
            userManager.DownloadModel(utils.string_wide_to_narrow(model.Name), exists.Path);

            // Update model list and load into runtime
            OnModelReady(exists.Path, true);
        }
        else
        {
            OnModelReady(exists.LocalExists ? exists.Path : exists.LegacyPath, false);
        }
    }).detach();
}