#ifndef UNICODE
#define UNICODE
#endif
#ifndef _UNICODE
#define _UNICODE
#endif

#include <windows.h>
#include <vector>
#include <complex>
#include <thread>
#include <mutex>
#include <atomic>
#include <chrono>

// Window size
const int WIDTH = 800;
const int HEIGHT = 600;

// Mandelbrot view
double xmin = -2.0, xmax = 1.0;
double ymin = -1.5, ymax = 1.5;
int max_iter = 500;

// Pixel buffer
std::vector<unsigned char> pixels(WIDTH* HEIGHT * 3);
HBITMAP hBitmap = nullptr;
std::mutex renderMutex;

// Track ongoing rendering
std::atomic<bool> rendering(false);

// Low-res rendering factor
int previewScale = 4;

// Forward declarations
void generateMandelbrot(int width, int height, std::vector<unsigned char>& buffer);

// Create or update the bitmap
void updateBitmap(HDC hdc)
{
    std::lock_guard<std::mutex> lock(renderMutex);

    BITMAPINFO bmi = {};
    bmi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bmi.bmiHeader.biWidth = WIDTH;
    bmi.bmiHeader.biHeight = -HEIGHT; // top-down
    bmi.bmiHeader.biPlanes = 1;
    bmi.bmiHeader.biBitCount = 24;
    bmi.bmiHeader.biCompression = BI_RGB;

    if (!hBitmap)
    {
        void* pBits;
        hBitmap = CreateDIBSection(hdc, &bmi, DIB_RGB_COLORS, &pBits, nullptr, 0);
    }

    SetDIBits(hdc, hBitmap, 0, HEIGHT, pixels.data(), &bmi, DIB_RGB_COLORS);
}

// Multi-threaded block render
void renderBlock(int yStart, int yEnd, int width, int height, std::vector<unsigned char>& buffer)
{
    for (int py = yStart; py < yEnd; ++py)
    {
        double y0 = ymin + (ymax - ymin) * py / height;
        for (int px = 0; px < width; ++px)
        {
            double x0 = xmin + (xmax - xmin) * px / width;
            std::complex<double> c(x0, y0), z(0, 0);
            int iter = 0;
            while (std::abs(z) <= 2.0 && iter < max_iter) z = z * z + c, iter++;

            int idx = (py * width + px) * 3;
            buffer[idx + 0] = iter % 256;
            buffer[idx + 1] = (iter * 2) % 256;
            buffer[idx + 2] = (iter * 3) % 256;
        }
    }
}

// Multi-threaded Mandelbrot
void generateMandelbrot(int width, int height, std::vector<unsigned char>& buffer)
{
    int numThreads = std::thread::hardware_concurrency();
    std::vector<std::thread> threads;
    int blockHeight = height / numThreads;

    for (int i = 0; i < numThreads; ++i)
    {
        int yStart = i * blockHeight;
        int yEnd = (i == numThreads - 1) ? height : yStart + blockHeight;
        threads.emplace_back(renderBlock, yStart, yEnd, width, height, std::ref(buffer));
    }

    for (auto& t : threads) t.join();
}

// Start rendering in a background thread
void startRender(HWND hwnd)
{
    if (rendering) return; // avoid multiple renders
    rendering = true;

    std::thread([hwnd]()
        {
            // 1) Low-res preview
            int lowW = WIDTH / previewScale, lowH = HEIGHT / previewScale;
            std::vector<unsigned char> preview(lowW * lowH * 3);
            generateMandelbrot(lowW, lowH, preview);

            // Upscale preview to full size
            {
                std::lock_guard<std::mutex> lock(renderMutex);
                for (int y = 0;y < HEIGHT;y++)
                    for (int x = 0;x < WIDTH;x++)
                    {
                        int px = x / previewScale;
                        int py = y / previewScale;
                        int idx = (y * WIDTH + x) * 3;
                        int idxSmall = (py * lowW + px) * 3;
                        pixels[idx + 0] = preview[idxSmall + 0];
                        pixels[idx + 1] = preview[idxSmall + 1];
                        pixels[idx + 2] = preview[idxSmall + 2];
                    }
            }
            InvalidateRect(hwnd, nullptr, TRUE);

            // 2) Full-res render
            generateMandelbrot(WIDTH, HEIGHT, pixels);
            InvalidateRect(hwnd, nullptr, TRUE);

            rendering = false;
        }).detach();
}

// Dragging
int lastX = 0, lastY = 0;
bool dragging = false;

// Window procedure
LRESULT CALLBACK WindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
    switch (uMsg)
    {
    case WM_DESTROY:
        PostQuitMessage(0);
        return 0;

    case WM_PAINT:
    {
        PAINTSTRUCT ps;
        HDC hdc = BeginPaint(hwnd, &ps);
        updateBitmap(hdc);

        HDC memDC = CreateCompatibleDC(hdc);
        HBITMAP oldBmp = (HBITMAP)SelectObject(memDC, hBitmap);
        BitBlt(hdc, 0, 0, WIDTH, HEIGHT, memDC, 0, 0, SRCCOPY);
        SelectObject(memDC, oldBmp);
        DeleteDC(memDC);

        EndPaint(hwnd, &ps);
    }
    return 0;

    case WM_MOUSEWHEEL:
    {
        short delta = GET_WHEEL_DELTA_WPARAM(wParam);
        double zoomFactor = (delta > 0) ? 0.8 : 1.25;

        double centerX = (xmin + xmax) / 2;
        double centerY = (ymin + ymax) / 2;
        double width = (xmax - xmin) * zoomFactor;
        double height = (ymax - ymin) * zoomFactor;

        xmin = centerX - width / 2;
        xmax = centerX + width / 2;
        ymin = centerY - height / 2;
        ymax = centerY + height / 2;

        startRender(hwnd);
    }
    return 0;

    case WM_LBUTTONDOWN:
        dragging = true;
        lastX = LOWORD(lParam);
        lastY = HIWORD(lParam);
        return 0;

    case WM_LBUTTONUP:
        dragging = false;
        return 0;

    case WM_MOUSEMOVE:
        if (dragging)
        {
            int x = LOWORD(lParam);
            int y = HIWORD(lParam);
            double dx = (x - lastX) * (xmax - xmin) / WIDTH;
            double dy = (y - lastY) * (ymax - ymin) / HEIGHT;

            xmin -= dx; xmax -= dx;
            ymin += dy; ymax += dy;

            lastX = x;
            lastY = y;

            startRender(hwnd);
        }
        return 0;
    }

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

// Entry point
int WINAPI wWinMain(HINSTANCE hInstance, HINSTANCE, PWSTR, int nCmdShow)
{
    const wchar_t CLASS_NAME[] = L"MandelbrotWindow";

    WNDCLASS wc = {};
    wc.lpfnWndProc = WindowProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = CLASS_NAME;
    RegisterClass(&wc);

    HWND hwnd = CreateWindowEx(
        0, CLASS_NAME, L"Mandelbrot Explorer", WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT, WIDTH, HEIGHT,
        nullptr, nullptr, hInstance, nullptr
    );

    ShowWindow(hwnd, nCmdShow);

    // Initial render
    startRender(hwnd);

    MSG msg = {};
    while (GetMessage(&msg, nullptr, 0, 0))
    {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    return 0;
}