#include "ui/Application.h"
#include "utils/Logger.h"
#include <pxr/base/plug/registry.h>
#include <cstdio>
#include <exception>
#include <Windows.h>
#include <string>
#include <filesystem>

static void SetUsdPluginPath() {
    char exePath[MAX_PATH];
    GetModuleFileNameA(nullptr, exePath, MAX_PATH);
    std::string exeDir(exePath);
    size_t lastSlash = exeDir.find_last_of("\\/");
    if (lastSlash != std::string::npos) {
        exeDir = exeDir.substr(0, lastSlash);
    }

    std::string pluginPath = exeDir + "\\usd";
    SetEnvironmentVariableA("PXR_PLUGINPATH_NAME", pluginPath.c_str());

    // --- MaterialX / Arnold OSL include path -------------------------------
    // hdArnold compiles MaterialX shaders by generating OSL whose first line is
    //   #include "mx_funcs.h"
    // and compiling it as an "osl" node's `code` parameter.  Without an include
    // path that compile fails with:
    //   [osl] error: <buffer>:2:10: fatal error: 'mx_funcs.h' file not found
    //
    // The knob is Arnold's options.osl_includepath, which hdArnold exposes as the
    // HDARNOLD_osl_includepath env setting (render_delegate/config.cpp).
    //
    // THIS MUST BE SET BEFORE ANY PLUGIN DLL IS LOADED.  TF_DEFINE_ENV_SETTING
    // registers a TF_REGISTRY_FUNCTION that eagerly calls TfGetEnvSetting() and
    // caches the value permanently when hdArnold.dll loads -- which the
    // LoadLibraryA pre-flight below does.  Setting it later (e.g. from
    // Application::Initialize) is silently ignored.
    //
    // Set via both the CRT (_putenv_s -> getenv) and the Win32 block
    // (SetEnvironmentVariableA) since the two are not kept in sync on Windows.
    {
        namespace fs = std::filesystem;
        const std::string mtlxLibs       = exeDir + "\\libraries";
        const std::string mtlxOslInclude = mtlxLibs + "\\stdlib\\genosl\\include";

        if (fs::exists(mtlxOslInclude)) {
            if (!getenv("HDARNOLD_osl_includepath")) {
                _putenv_s("HDARNOLD_osl_includepath", mtlxOslInclude.c_str());
                SetEnvironmentVariableA("HDARNOLD_osl_includepath",
                                        mtlxOslInclude.c_str());
            }
            // Nodedef lookup (MATERIALX_NODE_DEFINITIONS).  Separate knob from
            // the OSL include path above -- it does NOT fix the include error.
            if (!getenv("PXR_MTLX_STDLIB_SEARCH_PATHS")) {
                _putenv_s("PXR_MTLX_STDLIB_SEARCH_PATHS", mtlxLibs.c_str());
                SetEnvironmentVariableA("PXR_MTLX_STDLIB_SEARCH_PATHS",
                                        mtlxLibs.c_str());
            }
        }
    }

    std::vector<std::string> pluginPaths;
    WIN32_FIND_DATAA findData;
    std::string searchPattern = pluginPath + "\\*";
    HANDLE hFind = FindFirstFileA(searchPattern.c_str(), &findData);
    if (hFind != INVALID_HANDLE_VALUE) {
        do {
            if (findData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
                std::string dirName(findData.cFileName);
                if (dirName != "." && dirName != "..") {
                    std::string plugInfoPath = pluginPath + "\\" + dirName + "\\resources\\plugInfo.json";
                    DWORD attrs = GetFileAttributesA(plugInfoPath.c_str());
                    if (attrs != INVALID_FILE_ATTRIBUTES && !(attrs & FILE_ATTRIBUTE_DIRECTORY)) {
                        pluginPaths.push_back(plugInfoPath);
                    }
                }
            }
        } while (FindNextFileA(hFind, &findData));
        FindClose(hFind);
    }

    // Pre-flight: attempt to load each plugin DLL before registering it.
    // This catches missing dependencies (e.g. sycl8.dll → Intel GPU drivers)
    // gracefully and shows our own warning rather than a cryptic TF_ERROR.
    std::vector<std::string> checkedPaths;
    for (const auto& infoPath : pluginPaths) {
        // Derive the DLL path from the plugInfo.json location:
        // <plugin>/resources/plugInfo.json  →  ../<plugin>.dll
        // e.g. usd/hdCycles/resources/plugInfo.json  →  usd/hdCycles.dll
        size_t resPos = infoPath.rfind("\\resources\\plugInfo.json");
        if (resPos != std::string::npos) {
            std::string pluginDir = infoPath.substr(0, resPos);
            size_t dirSlash = pluginDir.rfind('\\');
            std::string pluginName = (dirSlash != std::string::npos)
                                     ? pluginDir.substr(dirSlash + 1) : pluginDir;
            std::string dllPath = pluginPath + "\\" + pluginName + ".dll";
            DWORD dllAttrs = GetFileAttributesA(dllPath.c_str());
            if (dllAttrs != INVALID_FILE_ATTRIBUTES) {
                // DLL file exists — verify it can actually be loaded.
                // No LOAD_WITH_ALTERED_SEARCH_PATH: keep the application
                // directory (exe dir, where embree4.dll etc. live) in the
                // DLL search path so transitive deps resolve correctly.
                HMODULE h = LoadLibraryA(dllPath.c_str());
                if (!h) {
                    DWORD err = GetLastError();
                    char msg[256] = {};
                    FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                                   nullptr, err, 0, msg, static_cast<DWORD>(sizeof(msg) - 1), nullptr);
                    std::string errMsg(msg);
                    while (!errMsg.empty() && (errMsg.back() == '\r' || errMsg.back() == '\n'))
                        errMsg.pop_back();
                    char hexErr[12] = {};
                    sprintf_s(hexErr, "0x%08X", static_cast<unsigned>(err));
                    LOG_WARNING("Skipping plugin " + pluginName
                                + " (DLL load failed " + hexErr + "): " + errMsg);
                    continue; // Skip registering this plugin
                }
                // Intentionally do NOT FreeLibrary here. Releasing the handle
                // drops the refcount to 0 and unloads the DLL; when USD's plug
                // system later loads it again DllMain re-fires, re-registering
                // any TfEnvSettings and causing "duplicate definition" warnings.
                // Keeping the handle loaded means USD's LoadLibrary is a no-op
                // (refcount bump only, no DllMain call), preventing duplicates.
                // The process holds these handles for its entire lifetime anyway.
            }
        }
        checkedPaths.push_back(infoPath);
    }

    if (!checkedPaths.empty()) {
        auto& registry = pxr::PlugRegistry::GetInstance();
        auto registered = registry.RegisterPlugins(checkedPaths);
        LOG_INFO("Registered " + std::to_string(registered.size()) + " USD plugins from " + pluginPath);
    } else {
        LOG_WARNING("No USD plugins found at " + pluginPath);
    }
}

int main(int /*argc*/, char* /*argv*/[]) {
    try {
        UsdLayerManager::Logger::Instance().SetLogLevel(UsdLayerManager::LogLevel::Info);
        // Temporary: capture LOG_INFO to a file so the [NodeGraph]/[ed] debug
        // traces are readable after the fact (this is a WIN32-subsystem app,
        // so stdout isn't visible when launched normally). Remove once the
        // node-editor drag regression is diagnosed.
        if (const char* appData = getenv("APPDATA")) {
            std::filesystem::path logDir = std::filesystem::path(appData) / "UsdLayerManager";
            std::filesystem::create_directories(logDir);
            UsdLayerManager::Logger::Instance().SetLogFile((logDir / "debug.log").string());
        }

        LOG_INFO("=== USD Layer Manager Starting ===");

        SetUsdPluginPath();
        
        UsdLayerManager::Application app;
        if (!app.Initialize("USD Layer Manager", 1280, 720)) {
            LOG_ERROR("Failed to initialize application");
            return 1;
        }
        
        app.Run();
        app.Shutdown();
        
        LOG_INFO("=== USD Layer Manager Exiting ===");
        return 0;
        
    } catch (const std::exception& e) {
        LOG_ERROR(std::string("Unhandled exception: ") + e.what());
        return 1;
    } catch (...) {
        LOG_ERROR("Unknown exception occurred");
        return 1;
    }
}
