#include "Gui.h"
#include <wincodec.h>
#include "shield_image_data.h"

static ID3D11Device* g_pd3dDevice = nullptr;
static ID3D11DeviceContext* g_pd3dDeviceContext = nullptr;
static IDXGISwapChain* g_pSwapChain = nullptr;
static bool                     g_SwapChainOccluded = false;
static UINT                     g_ResizeWidth = 0, g_ResizeHeight = 0;
static ID3D11RenderTargetView* g_mainRenderTargetView = nullptr;

static HWND                      g_OverlayHwnd = nullptr;


static ID3D11ShaderResourceView* g_ShieldTexture = nullptr;
static int                       g_ShieldWidth = 0;
static int                       g_ShieldHeight = 0;
static double                    g_LoaderStartTime = -1.0;
static bool                      g_ComInitialized = false;

struct DecodedImage
{
    int width = 0;
    int height = 0;
    std::vector<unsigned char> rgba;
};

// Forward declare message handler from imgui_impl_win32.cpp
extern IMGUI_IMPL_API LRESULT ImGui_ImplWin32_WndProcHandler(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);

static float Clamp01(float value)
{
    return value < 0.0f ? 0.0f : (value > 1.0f ? 1.0f : value);
}

static float SmoothStep01(float value)
{
    value = Clamp01(value);
    return value * value * (3.0f - 2.0f * value);
}

static bool DecodePngFromMemory(const unsigned char* data, size_t data_size, DecodedImage& image)
{
    IWICImagingFactory* factory = nullptr;
    IWICStream* stream = nullptr;
    IWICBitmapDecoder* decoder = nullptr;
    IWICBitmapFrameDecode* frame = nullptr;
    IWICFormatConverter* converter = nullptr;
    UINT width = 0;
    UINT height = 0;
    bool result = false;

    if (FAILED(CoCreateInstance(CLSID_WICImagingFactory, nullptr, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&factory))))
        goto cleanup;

    if (FAILED(factory->CreateStream(&stream)))
        goto cleanup;

    if (FAILED(stream->InitializeFromMemory(const_cast<BYTE*>(data), static_cast<DWORD>(data_size))))
        goto cleanup;

    if (FAILED(factory->CreateDecoderFromStream(stream, nullptr, WICDecodeMetadataCacheOnLoad, &decoder)))
        goto cleanup;

    if (FAILED(decoder->GetFrame(0, &frame)))
        goto cleanup;


    if (FAILED(frame->GetSize(&width, &height)))
        goto cleanup;

    if (FAILED(factory->CreateFormatConverter(&converter)))
        goto cleanup;

    if (FAILED(converter->Initialize(
        frame,
        GUID_WICPixelFormat32bppRGBA,
        WICBitmapDitherTypeNone,
        nullptr,
        0.0,
        WICBitmapPaletteTypeCustom)))
    {
        goto cleanup;
    }

    image.width = static_cast<int>(width);
    image.height = static_cast<int>(height);
    image.rgba.resize(static_cast<size_t>(width) * static_cast<size_t>(height) * 4);

    if (FAILED(converter->CopyPixels(
        nullptr,
        width * 4,
        static_cast<UINT>(image.rgba.size()),
        image.rgba.data())))
    {
        goto cleanup;
    }

    result = true;

cleanup:
    if (converter) converter->Release();
    if (frame) frame->Release();
    if (decoder) decoder->Release();
    if (stream) stream->Release();
    if (factory) factory->Release();
    return result;
}

static bool CreateTextureFromRgba(const DecodedImage& image, ID3D11ShaderResourceView** out_srv)
{
    if (!g_pd3dDevice || image.rgba.empty() || image.width <= 0 || image.height <= 0)
        return false;

    D3D11_TEXTURE2D_DESC desc;
    ZeroMemory(&desc, sizeof(desc));
    desc.Width = static_cast<UINT>(image.width);
    desc.Height = static_cast<UINT>(image.height);
    desc.MipLevels = 1;
    desc.ArraySize = 1;
    desc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    desc.SampleDesc.Count = 1;
    desc.Usage = D3D11_USAGE_DEFAULT;
    desc.BindFlags = D3D11_BIND_SHADER_RESOURCE;

    D3D11_SUBRESOURCE_DATA sub_resource;
    ZeroMemory(&sub_resource, sizeof(sub_resource));
    sub_resource.pSysMem = image.rgba.data();
    sub_resource.SysMemPitch = static_cast<UINT>(image.width * 4);

    ID3D11Texture2D* texture = nullptr;
    if (FAILED(g_pd3dDevice->CreateTexture2D(&desc, &sub_resource, &texture)))
        return false;

    D3D11_SHADER_RESOURCE_VIEW_DESC srv_desc;
    ZeroMemory(&srv_desc, sizeof(srv_desc));
    srv_desc.Format = desc.Format;
    srv_desc.ViewDimension = D3D11_SRV_DIMENSION_TEXTURE2D;
    srv_desc.Texture2D.MipLevels = 1;

    HRESULT hr = g_pd3dDevice->CreateShaderResourceView(texture, &srv_desc, out_srv);
    texture->Release();

    return SUCCEEDED(hr);
}

static void RemoveNearBlackBackground(DecodedImage& image)
{
    for (size_t i = 0; i + 3 < image.rgba.size(); i += 4)
    {
        unsigned char r = image.rgba[i + 0];
        unsigned char g = image.rgba[i + 1];
        unsigned char b = image.rgba[i + 2];
        unsigned char max_channel = r > g ? r : g;
        max_channel = max_channel > b ? max_channel : b;

        if (r < 32 && g < 32 && b < 32)
        {
            image.rgba[i + 3] = 0;
        }
        else if (max_channel < 58)
        {
            float t = (static_cast<float>(max_channel) - 32.0f) / 26.0f;
            t = Clamp01(t);
            image.rgba[i + 3] = static_cast<unsigned char>(static_cast<float>(image.rgba[i + 3]) * t);
        }
    }
}

static bool LoadShieldTextureFromGlobalMemory()
{
    DecodedImage image;
    if (!DecodePngFromMemory(g_ShieldPngData, sizeof(g_ShieldPngData), image))
        return false;

    RemoveNearBlackBackground(image);

    if (!CreateTextureFromRgba(image, &g_ShieldTexture))
        return false;

    g_ShieldWidth = image.width;
    g_ShieldHeight = image.height;
    return true;
}

static void SetupTransparentOverlayWindow(HWND hwnd)
{
    LONG_PTR ex_style = ::GetWindowLongPtrW(hwnd, GWL_EXSTYLE);
    ex_style |= WS_EX_LAYERED | WS_EX_TOPMOST | WS_EX_TOOLWINDOW;
    ::SetWindowLongPtrW(hwnd, GWL_EXSTYLE, ex_style);

    // Pure black pixels become transparent. Global alpha is used for fade in/out
    // to avoid blending image edges into the black backbuffer.
    ::SetLayeredWindowAttributes(hwnd, RGB(0, 0, 0), 255, LWA_COLORKEY | LWA_ALPHA);
}

static void CleanupShieldTexture()
{
    if (g_ShieldTexture)
    {
        g_ShieldTexture->Release();
        g_ShieldTexture = nullptr;
    }

    g_ShieldWidth = 0;
    g_ShieldHeight = 0;
}

static bool DrawAceStyleLoader()
{
    ImGuiViewport* viewport = ImGui::GetMainViewport();
    ImDrawList* draw = ImGui::GetBackgroundDrawList();

    ImVec2 vp_pos = viewport->Pos;
    ImVec2 vp_size = viewport->Size;

    if (!g_ShieldTexture || g_ShieldWidth <= 0 || g_ShieldHeight <= 0)
        return false;

    if (g_LoaderStartTime < 0.0)
        g_LoaderStartTime = ImGui::GetTime();

    const float fade_in_seconds = 1.0f;
    const float glow_seconds = 3.0f;
    const float hold_seconds = 1.0f;
    const float fade_out_seconds = 1.0f;
    const float fade_out_start = fade_in_seconds + glow_seconds + hold_seconds;
    const float close_time = fade_out_start + fade_out_seconds;

    const float elapsed = static_cast<float>(ImGui::GetTime() - g_LoaderStartTime);
    const float fade_in = SmoothStep01(elapsed / fade_in_seconds);
    const float fade_out = elapsed <= fade_out_start ? 1.0f : 1.0f - SmoothStep01((elapsed - fade_out_start) / fade_out_seconds);
    const float fade = Clamp01(fade_in * fade_out);
    const BYTE window_alpha = static_cast<BYTE>(255.0f * fade);

    if (g_OverlayHwnd)
        ::SetLayeredWindowAttributes(g_OverlayHwnd, RGB(0, 0, 0), window_alpha, LWA_COLORKEY | LWA_ALPHA);

    const ImVec2 uv0(0.0f, 0.0f);
    const ImVec2 uv1(1.0f, 1.0f);
    const float texture_aspect = static_cast<float>(g_ShieldHeight) / static_cast<float>(g_ShieldWidth);
    const float wanted_width = vp_size.x * 0.34f;
    float image_width = wanted_width < 420.0f ? wanted_width : 420.0f;
    if (image_width < 300.0f)
        image_width = 300.0f;
    const float image_height = image_width * texture_aspect;

    RECT work_area;
    work_area.left = 0;
    work_area.top = 0;
    work_area.right = static_cast<LONG>(vp_size.x);
    work_area.bottom = static_cast<LONG>(vp_size.y);
    ::SystemParametersInfoW(SPI_GETWORKAREA, 0, &work_area, 0);

    const float visible_right = 568.0f / 601.0f;
    const float visible_bottom = 215.0f / 254.0f;
    const float margin_x = 20.0f;
    const float margin_y = 8.0f;

    ImVec2 image_pos(
        vp_pos.x + static_cast<float>(work_area.right) - image_width * visible_right - margin_x,
        vp_pos.y + static_cast<float>(work_area.bottom) - image_height * visible_bottom - margin_y);
    ImVec2 image_end(image_pos.x + image_width, image_pos.y + image_height);

    draw->AddImage(
        g_ShieldTexture,
        image_pos,
        image_end,
        uv0,
        uv1,
        IM_COL32(255, 255, 255, 255));

    const float glow_elapsed = elapsed - fade_in_seconds;
    if (glow_elapsed < 0.0f)
        return false;

    const float sweep_t = SmoothStep01(glow_elapsed / glow_seconds);
    const int glow_alpha = glow_elapsed <= glow_seconds ? static_cast<int>(95.0f * fade_out) : 0;
    const int edge_alpha = glow_elapsed <= glow_seconds ? static_cast<int>(178.0f * fade_out) : 0;
    const int core_alpha = glow_elapsed <= glow_seconds ? static_cast<int>(238.0f * fade_out) : 0;

    // Restrict the ACE-style glow to the blue rectangular notification area.
    const ImVec2 panel_pos(
        image_pos.x + image_width * (183.0f / 601.0f),
        image_pos.y + image_height * (50.0f / 254.0f));
    const ImVec2 panel_end(
        image_pos.x + image_width * (568.0f / 601.0f),
        image_pos.y + image_height * (200.0f / 254.0f));
    const float panel_width = panel_end.x - panel_pos.x;
    const float panel_height = panel_end.y - panel_pos.y;

    const float beam_width = panel_width * 0.18f;
    const float beam_x = panel_pos.x - beam_width * 1.15f + (panel_width + beam_width * 2.30f) * sweep_t;
    const float beam_slant = panel_width * 0.17f;

    draw->PushClipRect(panel_pos, panel_end, true);

    const ImVec2 glow_top(beam_x - beam_width * 0.30f, panel_pos.y);
    const ImVec2 glow_bottom(beam_x + beam_slant - beam_width * 0.30f, panel_end.y);
    draw->AddLine(glow_top, glow_bottom, IM_COL32(45, 220, 255, glow_alpha), 32.0f);
    draw->AddLine(glow_top, glow_bottom, IM_COL32(115, 240, 255, edge_alpha), 13.0f);

    const ImVec2 core_top(beam_x, panel_pos.y);
    const ImVec2 core_bottom(beam_x + beam_slant, panel_end.y);
    draw->AddLine(core_top, core_bottom, IM_COL32(245, 255, 255, core_alpha), 3.0f);

    const ImVec2 side_top(beam_x + beam_width * 0.18f, panel_pos.y);
    const ImVec2 side_bottom(beam_x + beam_slant + beam_width * 0.18f, panel_end.y);
    draw->AddLine(side_top, side_bottom, IM_COL32(175, 250, 255, edge_alpha), 7.0f);

    draw->PopClipRect();

    return elapsed >= close_time;
}




namespace Gui
{
    auto MessageDispatch() -> void;

    bool CreateDeviceD3D(HWND hWnd);
    void CleanupDeviceD3D();
    void CreateRenderTarget();
    void CleanupRenderTarget();

    LRESULT WINAPI WndProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);

    WNDCLASSEXW wc = { sizeof(wc), CS_CLASSDC, WndProc, 0L, 0L, GetModuleHandle(nullptr), nullptr, nullptr, nullptr, nullptr, L"Esp Style Loader", nullptr };
   
    auto CreateGuiWindow() -> BOOL
    {
        HRESULT com_hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
        g_ComInitialized = SUCCEEDED(com_hr);

        ImGui_ImplWin32_EnableDpiAwareness();
        float main_scale = ImGui_ImplWin32_GetDpiScaleForMonitor(::MonitorFromPoint(POINT{ 0, 0 }, MONITOR_DEFAULTTOPRIMARY));

        ::RegisterClassExW(&wc);
        int screen_width = ::GetSystemMetrics(SM_CXSCREEN);
        int screen_height = ::GetSystemMetrics(SM_CYSCREEN);

        HWND hwnd = ::CreateWindowExW(
            WS_EX_TOPMOST | WS_EX_LAYERED | WS_EX_TRANSPARENT,
            wc.lpszClassName,
            L"Esp Style Loader",
            WS_POPUP,
            0,
            0,
            screen_width,
            screen_height,
            nullptr,
            nullptr,
            wc.hInstance,
            nullptr);

        g_OverlayHwnd = hwnd;

        ::SetWindowPos(
            hwnd,
            HWND_TOPMOST,
            0,
            0,
            screen_width,
            screen_height,
            SWP_NOACTIVATE | SWP_SHOWWINDOW);

        // Initialize Direct3D
        if (!CreateDeviceD3D(hwnd))
        {
            CleanupDeviceD3D();
            ::UnregisterClassW(wc.lpszClassName, wc.hInstance);

            if (g_ComInitialized)
                CoUninitialize();

            return FALSE;
        }

        SetupTransparentOverlayWindow(hwnd);

        if (!LoadShieldTextureFromGlobalMemory())
        {
            CleanupDeviceD3D();
            ::DestroyWindow(hwnd);
            ::UnregisterClassW(wc.lpszClassName, wc.hInstance);
            if (g_ComInitialized)
                CoUninitialize();
            return 1;
        }

        // Pure black pixels become transparent. Global alpha is used for fade in/out
        // to avoid blending image edges into the black backbuffer.


        // Show the window
        ::ShowWindow(hwnd, SW_SHOW);
        ::UpdateWindow(hwnd);

        // Setup Dear ImGui context
        //IMGUI_CHECKVERSION();
        ImGui::CreateContext();
        ImGuiIO& io = ImGui::GetIO(); (void)io;
        io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
        io.ConfigFlags |= ImGuiConfigFlags_NavEnableGamepad;

        // Setup Dear ImGui style
        ImGui::StyleColorsDark();

        // Setup scaling
        ImGuiStyle& style = ImGui::GetStyle();
        style.ScaleAllSizes(main_scale);
        style.FontScaleDpi = main_scale;

        // Setup Platform/Renderer backends
        ImGui_ImplWin32_Init(hwnd);
        ImGui_ImplDX11_Init(g_pd3dDevice, g_pd3dDeviceContext);


        MessageDispatch();
        
        //CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)MessageDispatch, NULL, 0, NULL);

        return TRUE;
    }

    auto MessageDispatch() -> void
    {
        bool done = false;
        while (!done)
        {
            MSG msg;
            while (::PeekMessage(&msg, nullptr, 0U, 0U, PM_REMOVE))
            {
                ::TranslateMessage(&msg);
                ::DispatchMessage(&msg);
                if (msg.message == WM_QUIT)
                    done = true;
            }
            if (done)
                break;

            if (g_SwapChainOccluded && g_pSwapChain->Present(0, DXGI_PRESENT_TEST) == DXGI_STATUS_OCCLUDED)
            {
                ::Sleep(10);
                continue;
            }
            g_SwapChainOccluded = false;

            if (g_ResizeWidth != 0 && g_ResizeHeight != 0)
            {
                CleanupRenderTarget();
                g_pSwapChain->ResizeBuffers(0, g_ResizeWidth, g_ResizeHeight, DXGI_FORMAT_UNKNOWN, 0);
                g_ResizeWidth = g_ResizeHeight = 0;
                CreateRenderTarget();
            }

            ImGui_ImplDX11_NewFrame();
            ImGui_ImplWin32_NewFrame();
            ImGui::NewFrame();

            if (DrawAceStyleLoader())
                break;
            

            ImGui::Render();
            const float clear_color_with_alpha[4] = { 0.0f, 0.0f, 0.0f, 1.00f };
            g_pd3dDeviceContext->OMSetRenderTargets(1, &g_mainRenderTargetView, nullptr);
            g_pd3dDeviceContext->ClearRenderTargetView(g_mainRenderTargetView, clear_color_with_alpha);
            ImGui_ImplDX11_RenderDrawData(ImGui::GetDrawData());

            HRESULT hr = g_pSwapChain->Present(1, 0);
            g_SwapChainOccluded = (hr == DXGI_STATUS_OCCLUDED);
        }

        CleanupShieldTexture();
        ImGui_ImplDX11_Shutdown();
        ImGui_ImplWin32_Shutdown();
        ImGui::DestroyContext();

        CleanupDeviceD3D();
        ::DestroyWindow(g_OverlayHwnd);
        ::UnregisterClassW(wc.lpszClassName, wc.hInstance);

        if (g_ComInitialized)
            CoUninitialize();
    }

    // Helper functions
    bool CreateDeviceD3D(HWND hWnd)
    {
        DXGI_SWAP_CHAIN_DESC sd;
        ZeroMemory(&sd, sizeof(sd));
        sd.BufferCount = 2;
        sd.BufferDesc.Width = 0;
        sd.BufferDesc.Height = 0;
        sd.BufferDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
        sd.BufferDesc.RefreshRate.Numerator = 60;
        sd.BufferDesc.RefreshRate.Denominator = 1;
        sd.Flags = DXGI_SWAP_CHAIN_FLAG_ALLOW_MODE_SWITCH;
        sd.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
        sd.OutputWindow = hWnd;
        sd.SampleDesc.Count = 1;
        sd.SampleDesc.Quality = 0;
        sd.Windowed = TRUE;
        sd.SwapEffect = DXGI_SWAP_EFFECT_DISCARD;

        UINT createDeviceFlags = 0;
        D3D_FEATURE_LEVEL featureLevel;
        const D3D_FEATURE_LEVEL featureLevelArray[2] = { D3D_FEATURE_LEVEL_11_0, D3D_FEATURE_LEVEL_10_0, };
        HRESULT res = D3D11CreateDeviceAndSwapChain(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, createDeviceFlags, featureLevelArray, 2, D3D11_SDK_VERSION, &sd, &g_pSwapChain, &g_pd3dDevice, &featureLevel, &g_pd3dDeviceContext);
        if (res == DXGI_ERROR_UNSUPPORTED)
            res = D3D11CreateDeviceAndSwapChain(nullptr, D3D_DRIVER_TYPE_WARP, nullptr, createDeviceFlags, featureLevelArray, 2, D3D11_SDK_VERSION, &sd, &g_pSwapChain, &g_pd3dDevice, &featureLevel, &g_pd3dDeviceContext);
        if (res != S_OK)
            return false;

        CreateRenderTarget();
        return true;
    }

    void CleanupDeviceD3D()
    {
        CleanupRenderTarget();
        if (g_pSwapChain) { g_pSwapChain->Release(); g_pSwapChain = nullptr; }
        if (g_pd3dDeviceContext) { g_pd3dDeviceContext->Release(); g_pd3dDeviceContext = nullptr; }
        if (g_pd3dDevice) { g_pd3dDevice->Release(); g_pd3dDevice = nullptr; }
    }

    void CreateRenderTarget()
    {
        ID3D11Texture2D* pBackBuffer;
        g_pSwapChain->GetBuffer(0, IID_PPV_ARGS(&pBackBuffer));
        g_pd3dDevice->CreateRenderTargetView(pBackBuffer, nullptr, &g_mainRenderTargetView);
        pBackBuffer->Release();
    }

    void CleanupRenderTarget()
    {
        if (g_mainRenderTargetView) { g_mainRenderTargetView->Release(); g_mainRenderTargetView = nullptr; }
    }

    // Win32 message handler
    LRESULT WINAPI WndProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam)
    {
        if (ImGui_ImplWin32_WndProcHandler(hWnd, msg, wParam, lParam))
            return true;

        switch (msg)
        {
        case WM_SIZE:
            if (wParam == SIZE_MINIMIZED)
                return 0;
            g_ResizeWidth = (UINT)LOWORD(lParam);
            g_ResizeHeight = (UINT)HIWORD(lParam);
            return 0;
        case WM_SYSCOMMAND:
            if ((wParam & 0xfff0) == SC_KEYMENU)
                return 0;
            break;
        case WM_DESTROY:
            ::PostQuitMessage(0);
            return 0;
        }
        return ::DefWindowProcW(hWnd, msg, wParam, lParam);
    }
}