#include <windows.h>
#include <d3d11.h>
#include <d3dcompiler.h>
#include <wrl/client.h>
#include <algorithm>
#include <tchar.h>
#include <cmath>
#include <sstream>
#include <iostream>
#include <iomanip>
#include <mutex>
#include <unordered_map>
#include "mandelbrot_app.h"

#ifndef GET_X_LPARAM
#define GET_X_LPARAM(lParam) ((int)(short)LOWORD(lParam))
#endif

#ifndef GET_Y_LPARAM
#define GET_Y_LPARAM(lParam) ((int)(short)HIWORD(lParam))
#endif

const char* computeShaderSource = R"(
cbuffer CB : register(b0) {
    float xmin, xmax, ymin, ymax;
    int width, height, maxIter;
    float time;
};

RWTexture2D<float4> Output : register(u0);

float3 palette(float t, int scheme) {
    if (scheme == 0) {
        float3 a = float3(0.5, 0.5, 0.5);
        float3 b = float3(0.5, 0.5, 0.5);
        float3 c = float3(1.0, 1.0, 1.0);
        float3 d = float3(0.0, 0.33, 0.67);
        return a + b * cos(6.28318 * (c * t + d));
    } else if (scheme == 1) {
        float3 a = float3(0.5, 0.5, 0.0);
        float3 b = float3(0.5, 0.5, 0.0);
        float3 c = float3(1.0, 0.7, 0.4);
        float3 d = float3(0.0, 0.15, 0.20);
        return a + b * cos(6.28318 * (c * t + d));
    } else if (scheme == 2) {
        float3 a = float3(0.2, 0.5, 0.8);
        float3 b = float3(0.2, 0.4, 0.2);
        float3 c = float3(2.0, 1.0, 1.0);
        float3 d = float3(0.0, 0.25, 0.25);
        return a + b * cos(6.28318 * (c * t + d));
    } else if (scheme == 3) {
        float3 a = float3(0.5, 0.2, 0.8);
        float3 b = float3(0.5, 0.5, 0.5);
        float3 c = float3(1.0, 1.0, 0.5);
        float3 d = float3(0.8, 0.9, 0.3);
        return a + b * cos(6.28318 * (c * t + d));
    } else {
        float v = 0.5 + 0.5 * cos(6.28318 * t);
        return float3(v, v, v);
    }
}

[numthreads(8, 8, 1)]
void main(uint3 DTid : SV_DispatchThreadID) {
    int x = DTid.x;
    int y = DTid.y;
    
    if (x >= width || y >= height) return;
    
    float cx = xmin + (xmax - xmin) * x / (float)width;
    float cy = ymin + (ymax - ymin) * y / (float)height;
    
    float zx = 0.0;
    float zy = 0.0;
    float zx2 = 0.0;
    float zy2 = 0.0;
    int iter = 0;
    
    while (iter < maxIter && (zx2 + zy2) < 4.0) {
        zy = 2.0 * zx * zy + cy;
        zx = zx2 - zy2 + cx;
        zx2 = zx * zx;
        zy2 = zy * zy;
        iter++;
    }
    
    float4 color;
    if (iter == maxIter) {
        color = float4(0.0, 0.0, 0.0, 1.0);
    } else {
        float log_zn = log(zx2 + zy2) * 0.5;
        float nu = log2(log_zn);
        float smooth_iter = iter + 1.0 - nu;
        
        int scheme = maxIter >> 16;
        int actualMaxIter = maxIter & 0xFFFF;
        
        float t = smooth_iter / 50.0 + time * 0.02;
        float3 rgb = palette(t, scheme);
        
        float brightness = 0.5 + 0.5 * sin(smooth_iter * 0.1);
        rgb *= brightness;
        
        color = float4(rgb, 1.0);
    }
    
    Output[uint2(x, y)] = color;
}
)";

const char* vertexShaderSource = R"(
struct VS_OUTPUT {
    float4 pos : SV_POSITION;
    float2 tex : TEXCOORD0;
};

VS_OUTPUT main(uint id : SV_VertexID) {
    VS_OUTPUT output;
    output.tex = float2((id << 1) & 2, id & 2);
    output.pos = float4(output.tex * float2(2, -2) + float2(-1, 1), 0, 1);
    return output;
}
)";

const char* pixelShaderSource = R"(
Texture2D tex : register(t0);
SamplerState samp : register(s0);

float4 main(float4 pos : SV_POSITION, float2 texCoord : TEXCOORD0) : SV_TARGET {
    return tex.Sample(samp, texCoord);
}
)";

MandelbrotApp::MandelbrotApp(HINSTANCE hInstance) : hInstance_(hInstance) {
    QueryPerformanceFrequency(&perfFreq_);
    QueryPerformanceCounter(&lastFrameTime_);

    CreateMainWindow();
    if (hwnd_) {
        InitD3D();
        CreateComputeShader();
        UpdateTitle();
    }
}

MandelbrotApp::~MandelbrotApp() {
    if (context_) context_->ClearState();
}

void MandelbrotApp::CreateMainWindow() {
    WNDCLASSEX wc = {};
    wc.cbSize = sizeof(WNDCLASSEX);
    wc.style = CS_HREDRAW | CS_VREDRAW;
    wc.lpfnWndProc = WndProc;
    wc.hInstance = hInstance_;
    wc.hCursor = LoadCursor(nullptr, IDC_ARROW);
    wc.hbrBackground = (HBRUSH)(COLOR_WINDOW + 1);
    wc.lpszClassName = _T("MandelbrotViewerGPU");

    if (!RegisterClassEx(&wc)) {
        MessageBox(nullptr, L"Failed to register window class", L"Error", MB_OK);
        return;
    }

    RECT rc = { 0, 0, 1280, 720 };
    AdjustWindowRect(&rc, WS_OVERLAPPEDWINDOW, FALSE);

    hwnd_ = CreateWindowEx(
        0,
        _T("MandelbrotViewerGPU"),
        _T("GPU Mandelbrot Viewer"),
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT,
        rc.right - rc.left, rc.bottom - rc.top,
        nullptr, nullptr, hInstance_, this
    );

    if (!hwnd_) {
        MessageBox(nullptr, L"Failed to create window", L"Error", MB_OK);
        return;
    }

    ShowWindow(hwnd_, SW_SHOWDEFAULT);
    UpdateWindow(hwnd_);
}

void MandelbrotApp::InitD3D() {
    RECT rc;
    GetClientRect(hwnd_, &rc);

    imageWidth_ = rc.right - rc.left;
    imageHeight_ = rc.bottom - rc.top;
    if (imageWidth_ < 1) imageWidth_ = 1;
    if (imageHeight_ < 1) imageHeight_ = 1;

    DXGI_SWAP_CHAIN_DESC scd = {};
    scd.BufferCount = 2;
    scd.BufferDesc.Width = imageWidth_;
    scd.BufferDesc.Height = imageHeight_;
    scd.BufferDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    scd.BufferDesc.RefreshRate.Numerator = 60;
    scd.BufferDesc.RefreshRate.Denominator = 1;
    scd.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
    scd.OutputWindow = hwnd_;
    scd.SampleDesc.Count = 1;
    scd.SampleDesc.Quality = 0;
    scd.Windowed = TRUE;
    scd.SwapEffect = DXGI_SWAP_EFFECT_DISCARD;

    D3D_FEATURE_LEVEL featureLevels[] = {
        D3D_FEATURE_LEVEL_11_0,
        D3D_FEATURE_LEVEL_10_1,
        D3D_FEATURE_LEVEL_10_0
    };
    UINT numFeatureLevels = ARRAYSIZE(featureLevels);

    D3D_FEATURE_LEVEL featureLevel;
    UINT flags = 0;

    HRESULT hr = D3D11CreateDeviceAndSwapChain(
        nullptr,
        D3D_DRIVER_TYPE_HARDWARE,
        nullptr,
        flags,
        featureLevels,
        numFeatureLevels,
        D3D11_SDK_VERSION,
        &scd,
        &swapChain_,
        &device_,
        &featureLevel,
        &context_
    );

    if (FAILED(hr)) {
        MessageBox(nullptr, L"Failed to create D3D11 device and swap chain", L"Error", MB_OK);
        return;
    }

    ComPtr<ID3D11Texture2D> backBuffer;
    hr = swapChain_->GetBuffer(0, IID_PPV_ARGS(&backBuffer));
    if (FAILED(hr)) {
        MessageBox(nullptr, L"Failed to get swap chain back buffer", L"Error", MB_OK);
        return;
    }

    hr = device_->CreateRenderTargetView(backBuffer.Get(), nullptr, &renderTargetView_);
    if (FAILED(hr)) {
        MessageBox(nullptr, L"Failed to create render target view", L"Error", MB_OK);
        return;
    }

    D3D11_SAMPLER_DESC sampDesc = {};
    sampDesc.Filter = D3D11_FILTER_MIN_MAG_MIP_LINEAR;
    sampDesc.AddressU = D3D11_TEXTURE_ADDRESS_CLAMP;
    sampDesc.AddressV = D3D11_TEXTURE_ADDRESS_CLAMP;
    sampDesc.AddressW = D3D11_TEXTURE_ADDRESS_CLAMP;
    sampDesc.ComparisonFunc = D3D11_COMPARISON_NEVER;
    sampDesc.MinLOD = 0;
    sampDesc.MaxLOD = D3D11_FLOAT32_MAX;

    hr = device_->CreateSamplerState(&sampDesc, &samplerState_);
    if (FAILED(hr)) {
        MessageBox(nullptr, L"Failed to create sampler state", L"Error", MB_OK);
        return;
    }

    // Set the viewport
    D3D11_VIEWPORT vp;
    vp.Width = (float)imageWidth_;
    vp.Height = (float)imageHeight_;
    vp.MinDepth = 0.0f;
    vp.MaxDepth = 1.0f;
    vp.TopLeftX = 0;
    vp.TopLeftY = 0;
    context_->RSSetViewports(1, &vp);

    // Set render target
    context_->OMSetRenderTargets(1, renderTargetView_.GetAddressOf(), nullptr);

    ResizeBuffers(imageWidth_, imageHeight_);
}

void MandelbrotApp::CreateComputeShader() {
    if (!device_) return;

    ComPtr<ID3DBlob> csBlob, errBlob;

    HRESULT hr = D3DCompile(
        computeShaderSource,
        strlen(computeShaderSource),
        nullptr,
        nullptr,
        nullptr,
        "main",
        "cs_5_0",
        D3DCOMPILE_ENABLE_STRICTNESS,
        0,
        &csBlob,
        &errBlob
    );

    if (FAILED(hr)) {
        if (errBlob) {
            OutputDebugStringA((char*)errBlob->GetBufferPointer());
        }
        MessageBox(nullptr, L"Failed to compile compute shader", L"Error", MB_OK);
        return;
    }

    hr = device_->CreateComputeShader(csBlob->GetBufferPointer(), csBlob->GetBufferSize(), nullptr, &computeShader_);
    if (FAILED(hr)) {
        MessageBox(nullptr, L"Failed to create compute shader", L"Error", MB_OK);
        return;
    }

    // Vertex shader
    ComPtr<ID3DBlob> vsBlob;
    hr = D3DCompile(
        vertexShaderSource,
        strlen(vertexShaderSource),
        nullptr,
        nullptr,
        nullptr,
        "main",
        "vs_5_0",
        D3DCOMPILE_ENABLE_STRICTNESS,
        0,
        &vsBlob,
        &errBlob
    );
    if (SUCCEEDED(hr)) {
        device_->CreateVertexShader(vsBlob->GetBufferPointer(), vsBlob->GetBufferSize(), nullptr, &vertexShader_);
    }
    else {
        MessageBox(nullptr, L"Failed to compile vertex shader", L"Error", MB_OK);
    }

    // Pixel shader
    ComPtr<ID3DBlob> psBlob;
    hr = D3DCompile(
        pixelShaderSource,
        strlen(pixelShaderSource),
        nullptr,
        nullptr,
        nullptr,
        "main",
        "ps_5_0",
        D3DCOMPILE_ENABLE_STRICTNESS,
        0,
        &psBlob,
        &errBlob
    );
    if (SUCCEEDED(hr)) {
        device_->CreatePixelShader(psBlob->GetBufferPointer(), psBlob->GetBufferSize(), nullptr, &pixelShader_);
    }
    else {
        MessageBox(nullptr, L"Failed to compile pixel shader", L"Error", MB_OK);
    }

    // Constant buffer - FIXED: Ensure proper alignment
    D3D11_BUFFER_DESC cbDesc = {};
    cbDesc.ByteWidth = (sizeof(ConstantBuffer) + 15) & ~15; // Align to 16 bytes
    cbDesc.Usage = D3D11_USAGE_DYNAMIC;
    cbDesc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    cbDesc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    cbDesc.MiscFlags = 0;
    cbDesc.StructureByteStride = 0;

    hr = device_->CreateBuffer(&cbDesc, nullptr, &constantBuffer_);
    if (FAILED(hr)) {
        MessageBox(nullptr, L"Failed to create constant buffer", L"Error", MB_OK);
    }
}

void MandelbrotApp::ResizeBuffers(int width, int height) {
    if (width <= 0 || height <= 0 || !device_) return;

    outputTexture_.Reset();
    outputUAV_.Reset();
    outputSRV_.Reset();

    // Create output texture for compute shader
    D3D11_TEXTURE2D_DESC texDesc = {};
    texDesc.Width = width;
    texDesc.Height = height;
    texDesc.MipLevels = 1;
    texDesc.ArraySize = 1;
    texDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    texDesc.SampleDesc.Count = 1;
    texDesc.SampleDesc.Quality = 0;
    texDesc.Usage = D3D11_USAGE_DEFAULT;
    texDesc.BindFlags = D3D11_BIND_UNORDERED_ACCESS | D3D11_BIND_SHADER_RESOURCE;
    texDesc.CPUAccessFlags = 0;
    texDesc.MiscFlags = 0;

    HRESULT hr = device_->CreateTexture2D(&texDesc, nullptr, &outputTexture_);
    if (FAILED(hr)) {
        MessageBox(nullptr, L"Failed to create output texture", L"Error", MB_OK);
        return;
    }

    // Create UAV for compute shader
    D3D11_UNORDERED_ACCESS_VIEW_DESC uavDesc = {};
    uavDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    uavDesc.ViewDimension = D3D11_UAV_DIMENSION_TEXTURE2D;
    uavDesc.Texture2D.MipSlice = 0;

    hr = device_->CreateUnorderedAccessView(outputTexture_.Get(), &uavDesc, &outputUAV_);
    if (FAILED(hr)) {
        MessageBox(nullptr, L"Failed to create UAV", L"Error", MB_OK);
        return;
    }

    // Create SRV for pixel shader
    D3D11_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
    srvDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    srvDesc.ViewDimension = D3D11_SRV_DIMENSION_TEXTURE2D;
    srvDesc.Texture2D.MostDetailedMip = 0;
    srvDesc.Texture2D.MipLevels = 1;

    hr = device_->CreateShaderResourceView(outputTexture_.Get(), &srvDesc, &outputSRV_);
    if (FAILED(hr)) {
        MessageBox(nullptr, L"Failed to create SRV", L"Error", MB_OK);
        return;
    }

    imageWidth_ = width;
    imageHeight_ = height;
}

void MandelbrotApp::UpdateTitle() {
    double zoom = 4.0 / (xmax_ - xmin_);
    std::wostringstream title;
    title << L"Mandelbrot Viewer | Zoom: " << std::fixed << std::setprecision(2) << zoom
        << L"x | Iter: " << maxIter_
        << L" | Color: " << colorScheme_ + 1
        << L" | Anim: " << (animationEnabled_ ? L"ON" : L"OFF")
        << L" | [H for Help]";
    SetWindowText(hwnd_, title.str().c_str());
}

void MandelbrotApp::SaveBookmark() {
    for (int i = 0; i < 10; i++) {
        if (!bookmarks_[i].saved || i == 0) {
            std::lock_guard<std::mutex> lock(viewMutex_);
            bookmarks_[i].xmin = xmin_;
            bookmarks_[i].xmax = xmax_;
            bookmarks_[i].ymin = ymin_;
            bookmarks_[i].ymax = ymax_;
            bookmarks_[i].maxIter = maxIter_;
            bookmarks_[i].saved = true;

            std::wostringstream msg;
            msg << L"Bookmark saved to slot " << i;
            MessageBox(hwnd_, msg.str().c_str(), L"Bookmark", MB_OK | MB_ICONINFORMATION);
            break;
        }
    }
}

void MandelbrotApp::LoadBookmark(int slot) {
    if (slot >= 0 && slot < 10 && bookmarks_[slot].saved) {
        std::lock_guard<std::mutex> lock(viewMutex_);
        xmin_ = bookmarks_[slot].xmin;
        xmax_ = bookmarks_[slot].xmax;
        ymin_ = bookmarks_[slot].ymin;
        ymax_ = bookmarks_[slot].ymax;
        maxIter_ = bookmarks_[slot].maxIter;
        UpdateTitle();
    }
}

void MandelbrotApp::ResetView() {
    std::lock_guard<std::mutex> lock(viewMutex_);
    xmin_ = -2.5; xmax_ = 1.5;
    ymin_ = -1.5; ymax_ = 1.5;
    maxIter_ = 256;
    UpdateTitle();
}

void MandelbrotApp::ToggleAnimation() {
    animationEnabled_ = !animationEnabled_;
    if (!animationEnabled_) {
        animTime_ = 0.0f;
    }
    UpdateTitle();
}

void MandelbrotApp::AdjustIterations(int delta) {
    std::lock_guard<std::mutex> lock(viewMutex_);
    maxIter_ = std::max(64, std::min(8192, maxIter_ + delta));
    UpdateTitle();
}

void MandelbrotApp::CycleColorScheme() {
    colorScheme_ = (colorScheme_ + 1) % 5;
    UpdateTitle();
}

void MandelbrotApp::Zoom(double factor, int centerX, int centerY) {
    std::lock_guard<std::mutex> lock(viewMutex_);

    double centerReal = xmin_ + (xmax_ - xmin_) * centerX / imageWidth_;
    double centerImag = ymin_ + (ymax_ - ymin_) * centerY / imageHeight_;

    double width = (xmax_ - xmin_) * factor;
    double height = (ymax_ - ymin_) * factor;

    xmin_ = centerReal - width * 0.5;
    xmax_ = centerReal + width * 0.5;
    ymin_ = centerImag - height * 0.5;
    ymax_ = centerImag + height * 0.5;
}

void MandelbrotApp::Pan(int dx, int dy) {
    std::lock_guard<std::mutex> lock(viewMutex_);

    double deltaX = (xmax_ - xmin_) * dx / imageWidth_;
    double deltaY = (ymax_ - ymin_) * dy / imageHeight_;

    xmin_ -= deltaX;
    xmax_ -= deltaX;
    ymin_ += deltaY;
    ymax_ += deltaY;
}

void MandelbrotApp::RenderMandelbrot() {
    if (!device_ || !context_ || !swapChain_ || !computeShader_ ||
        !constantBuffer_ || !outputUAV_ || !renderTargetView_) {
        return;
    }

    ConstantBuffer cb;
    {
        std::lock_guard<std::mutex> lock(viewMutex_);

        double zoom = 4.0 / (xmax_ - xmin_);
        int adaptiveMaxIter = maxIter_;
        if (zoom > 10) adaptiveMaxIter = std::max(maxIter_, 512);
        if (zoom > 100) adaptiveMaxIter = std::max(maxIter_, 1024);
        if (zoom > 1000) adaptiveMaxIter = std::max(maxIter_, 2048);

        cb.xmin = (float)xmin_;
        cb.xmax = (float)xmax_;
        cb.ymin = (float)ymin_;
        cb.ymax = (float)ymax_;
        cb.width = imageWidth_;
        cb.height = imageHeight_;
        cb.maxIter = (colorScheme_ << 16) | (adaptiveMaxIter & 0xFFFF);
        cb.time = animTime_;
    }

    // Update constant buffer
    D3D11_MAPPED_SUBRESOURCE mapped;
    HRESULT hr = context_->Map(constantBuffer_.Get(), 0, D3D11_MAP_WRITE_DISCARD, 0, &mapped);
    if (SUCCEEDED(hr)) {
        memcpy(mapped.pData, &cb, sizeof(ConstantBuffer));
        context_->Unmap(constantBuffer_.Get(), 0);
    }

    // Run compute shader
    context_->CSSetShader(computeShader_.Get(), nullptr, 0);
    context_->CSSetConstantBuffers(0, 1, constantBuffer_.GetAddressOf());
    context_->CSSetUnorderedAccessViews(0, 1, outputUAV_.GetAddressOf(), nullptr);

    UINT dispatchX = (imageWidth_ + 7) / 8;
    UINT dispatchY = (imageHeight_ + 7) / 8;
    context_->Dispatch(dispatchX, dispatchY, 1);

    // Clear compute shader bindings
    ID3D11UnorderedAccessView* nullUAV[] = { nullptr };
    context_->CSSetUnorderedAccessViews(0, 1, nullUAV, nullptr);
    context_->CSSetShader(nullptr, nullptr, 0);

    // Clear render target
    float clearColor[4] = { 0.0f, 0.0f, 0.0f, 1.0f };
    context_->ClearRenderTargetView(renderTargetView_.Get(), clearColor);

    // Set viewport
    D3D11_VIEWPORT vp = {};
    vp.Width = (float)imageWidth_;
    vp.Height = (float)imageHeight_;
    vp.MinDepth = 0.0f;
    vp.MaxDepth = 1.0f;
    context_->RSSetViewports(1, &vp);

    // Render to screen
    context_->OMSetRenderTargets(1, renderTargetView_.GetAddressOf(), nullptr);
    context_->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
    context_->VSSetShader(vertexShader_.Get(), nullptr, 0);
    context_->PSSetShader(pixelShader_.Get(), nullptr, 0);
    context_->PSSetShaderResources(0, 1, outputSRV_.GetAddressOf());
    context_->PSSetSamplers(0, 1, samplerState_.GetAddressOf());
    context_->Draw(3, 0);

    swapChain_->Present(1, 0);

    if (animationEnabled_) {
        animTime_ += 0.016f;
    }
}

LRESULT CALLBACK MandelbrotApp::WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    MandelbrotApp* pApp = nullptr;

    if (msg == WM_NCCREATE) {
        CREATESTRUCT* pCreate = reinterpret_cast<CREATESTRUCT*>(lParam);
        pApp = reinterpret_cast<MandelbrotApp*>(pCreate->lpCreateParams);
        SetWindowLongPtr(hwnd, GWLP_USERDATA, reinterpret_cast<LONG_PTR>(pApp));
    }
    else {
        pApp = reinterpret_cast<MandelbrotApp*>(GetWindowLongPtr(hwnd, GWLP_USERDATA));
    }

    if (pApp) {
        return pApp->HandleMessage(hwnd, msg, wParam, lParam);
    }

    return DefWindowProc(hwnd, msg, wParam, lParam);
}

LRESULT MandelbrotApp::HandleMessage(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    switch (msg) {
    case WM_SIZE: {
        if (swapChain_ && wParam != SIZE_MINIMIZED) {
            RECT rc;
            GetClientRect(hwnd, &rc);
            int width = rc.right - rc.left;
            int height = rc.bottom - rc.top;

            if (width > 0 && height > 0) {
                context_->OMSetRenderTargets(0, nullptr, nullptr);
                renderTargetView_.Reset();

                HRESULT hr = swapChain_->ResizeBuffers(0, width, height, DXGI_FORMAT_UNKNOWN, 0);
                if (SUCCEEDED(hr)) {
                    ComPtr<ID3D11Texture2D> backBuffer;
                    hr = swapChain_->GetBuffer(0, IID_PPV_ARGS(&backBuffer));
                    if (SUCCEEDED(hr)) {
                        device_->CreateRenderTargetView(backBuffer.Get(), nullptr, &renderTargetView_);
                        ResizeBuffers(width, height);

                        if (context_ && renderTargetView_) {
                            context_->OMSetRenderTargets(1, renderTargetView_.GetAddressOf(), nullptr);

                            // Reset viewport
                            D3D11_VIEWPORT vp;
                            vp.Width = (float)width;
                            vp.Height = (float)height;
                            vp.MinDepth = 0.0f;
                            vp.MaxDepth = 1.0f;
                            vp.TopLeftX = 0;
                            vp.TopLeftY = 0;
                            context_->RSSetViewports(1, &vp);
                        }
                    }
                }
            }
        }
        return 0;
    }

    case WM_DESTROY:
        PostQuitMessage(0);
        return 0;

    case WM_MOUSEWHEEL: {
        int delta = GET_WHEEL_DELTA_WPARAM(wParam);
        POINT pt = { GET_X_LPARAM(lParam), GET_Y_LPARAM(lParam) };
        ScreenToClient(hwnd, &pt);

        bool shiftPressed = (GetKeyState(VK_SHIFT) & 0x8000) != 0;
        double zoomFactor = shiftPressed ?
            (delta > 0 ? 0.5 : 2.0) :
            (delta > 0 ? 0.8 : 1.25);

        Zoom(zoomFactor, pt.x, pt.y);
        UpdateTitle();
        return 0;
    }

    case WM_RBUTTONDOWN: {
        SetCapture(hwnd);
        lastMouseX_ = GET_X_LPARAM(lParam);
        lastMouseY_ = GET_Y_LPARAM(lParam);
        rightDragging_ = true;
        return 0;
    }

    case WM_RBUTTONUP: {
        if (rightDragging_) {
            int currentX = GET_X_LPARAM(lParam);
            int currentY = GET_Y_LPARAM(lParam);

            if (abs(currentX - lastMouseX_) < 5 && abs(currentY - lastMouseY_) < 5) {
                ResetView();
            }
        }
        rightDragging_ = false;
        ReleaseCapture();
        return 0;
    }

    case WM_LBUTTONDOWN: {
        SetCapture(hwnd);
        lastMouseX_ = GET_X_LPARAM(lParam);
        lastMouseY_ = GET_Y_LPARAM(lParam);
        dragging_ = true;
        SetCursor(LoadCursor(nullptr, IDC_SIZEALL));
        return 0;
    }

    case WM_LBUTTONUP: {
        dragging_ = false;
        ReleaseCapture();
        SetCursor(LoadCursor(nullptr, IDC_ARROW));
        return 0;
    }

    case WM_MBUTTONDOWN: {
        int x = GET_X_LPARAM(lParam);
        int y = GET_Y_LPARAM(lParam);
        Zoom(1.0, x, y);
        UpdateTitle();
        return 0;
    }

    case WM_MOUSEMOVE: {
        if (dragging_) {
            int currentX = GET_X_LPARAM(lParam);
            int currentY = GET_Y_LPARAM(lParam);
            int dx = currentX - lastMouseX_;
            int dy = currentY - lastMouseY_;

            if (dx != 0 || dy != 0) {
                Pan(dx, dy);
                lastMouseX_ = currentX;
                lastMouseY_ = currentY;
            }
        }
        else if (rightDragging_) {
            int currentY = GET_Y_LPARAM(lParam);
            int dy = currentY - lastMouseY_;

            if (abs(dy) > 2) {
                AdjustIterations(-dy * 5);
                lastMouseY_ = currentY;
            }
        }
        return 0;
    }

    case WM_KEYDOWN: {
        keysPressed_[wParam] = true;

        switch (wParam) {
        case 'H': {
            const wchar_t* helpText =
                L"MANDELBROT VIEWER CONTROLS\n\n"
                L"Mouse:\n"
                L"  Left Drag - Pan view\n"
                L"  Wheel - Zoom in/out\n"
                L"  Shift+Wheel - Fast zoom\n"
                L"  Middle Click - Center on point\n"
                L"  Right Drag - Adjust iterations\n"
                L"  Right Click - Reset view\n\n"
                L"Keyboard:\n"
                L"  Arrow Keys - Pan view\n"
                L"  Shift+Arrows - Fast pan\n"
                L"  +/= / - - Zoom in/out\n"
                L"  [ / ] - Adjust iterations\n"
                L"  C - Cycle color schemes (5 total)\n"
                L"  A - Toggle animation\n"
                L"  R - Reset view\n"
                L"  S - Save bookmark\n"
                L"  0-9 - Load bookmark\n"
                L"  F11 - Toggle fullscreen\n"
                L"  ESC - Exit fullscreen\n"
                L"  H - This help";
            MessageBox(hwnd, helpText, L"Help", MB_OK | MB_ICONINFORMATION);
            return 0;
        }

        case 'R':
            ResetView();
            return 0;

        case 'C':
            CycleColorScheme();
            return 0;

        case 'A':
            ToggleAnimation();
            return 0;

        case 'S':
            SaveBookmark();
            return 0;

        case '0': case '1': case '2': case '3': case '4':
        case '5': case '6': case '7': case '8': case '9':
            LoadBookmark(wParam - '0');
            return 0;

        case VK_OEM_PLUS:
        case VK_ADD:
            Zoom(0.8, imageWidth_ / 2, imageHeight_ / 2);
            UpdateTitle();
            return 0;

        case VK_OEM_MINUS:
        case VK_SUBTRACT:
            Zoom(1.25, imageWidth_ / 2, imageHeight_ / 2);
            UpdateTitle();
            return 0;

        case 'D': // [
            AdjustIterations(-64);
            return 0;

        case 'I': // ]
            AdjustIterations(64);
            return 0;

        case VK_F11: {
            static bool fullscreen = false;
            static RECT savedRect;
            static DWORD savedStyle;

            if (!fullscreen) {
                GetWindowRect(hwnd, &savedRect);
                savedStyle = GetWindowLong(hwnd, GWL_STYLE);

                SetWindowLong(hwnd, GWL_STYLE, savedStyle & ~(WS_CAPTION | WS_THICKFRAME));

                MONITORINFO mi = { sizeof(mi) };
                GetMonitorInfo(MonitorFromWindow(hwnd, MONITOR_DEFAULTTOPRIMARY), &mi);
                SetWindowPos(hwnd, HWND_TOP,
                    mi.rcMonitor.left, mi.rcMonitor.top,
                    mi.rcMonitor.right - mi.rcMonitor.left,
                    mi.rcMonitor.bottom - mi.rcMonitor.top,
                    SWP_FRAMECHANGED);
                fullscreen = true;
            }
            else {
                SetWindowLong(hwnd, GWL_STYLE, savedStyle);
                SetWindowPos(hwnd, nullptr,
                    savedRect.left, savedRect.top,
                    savedRect.right - savedRect.left,
                    savedRect.bottom - savedRect.top,
                    SWP_FRAMECHANGED);
                fullscreen = false;
            }
            return 0;
        }

        case VK_ESCAPE: {
            DWORD currentStyle = GetWindowLong(hwnd, GWL_STYLE);
            if (!(currentStyle & WS_CAPTION)) {
                SetWindowLong(hwnd, GWL_STYLE, WS_OVERLAPPEDWINDOW);
                SetWindowPos(hwnd, nullptr, 100, 100, 1280, 720, SWP_FRAMECHANGED);
            }
            return 0;
        }

        case VK_LEFT:
        case VK_RIGHT:
        case VK_UP:
        case VK_DOWN: {
            bool shift = (GetKeyState(VK_SHIFT) & 0x8000) != 0;
            double speed = shift ? 50 : 10;

            int dx = 0, dy = 0;
            if (wParam == VK_LEFT) dx = (int)speed;
            if (wParam == VK_RIGHT) dx = -(int)speed;
            if (wParam == VK_UP) dy = (int)speed;
            if (wParam == VK_DOWN) dy = -(int)speed;

            Pan(dx, dy);
            return 0;
        }
        }
        return 0;
    }

    case WM_KEYUP:
        keysPressed_[wParam] = false;
        return 0;

    default:
        return DefWindowProc(hwnd, msg, wParam, lParam);
    }
}

void MandelbrotApp::Run() {
    MSG msg = {};
    LARGE_INTEGER lastTime;
    QueryPerformanceCounter(&lastTime);

    double targetFrameTime = 1.0 / 120.0; // 60 FPS
    LARGE_INTEGER freq;
    QueryPerformanceFrequency(&freq);

    while (true) {
        if (PeekMessage(&msg, nullptr, 0, 0, PM_REMOVE)) {
            if (msg.message == WM_QUIT)
                break;
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        }
        else {
            RenderMandelbrot();

            // Limit frame rate
            LARGE_INTEGER currentTime;
            QueryPerformanceCounter(&currentTime);
            double elapsed = double(currentTime.QuadPart - lastTime.QuadPart) / freq.QuadPart;

            if (elapsed < targetFrameTime) {
                DWORD sleepMs = DWORD((targetFrameTime - elapsed) * 1000);
                if (sleepMs > 0) Sleep(sleepMs);
            }
            QueryPerformanceCounter(&lastTime);
        }
    }
}