#include <windows.h>
#include <d3d11.h>
#include <d3dcompiler.h>
#include <stdio.h>

#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "D3DCompiler.lib")

HWND hwnd = nullptr;
ID3D11Device* device = nullptr;
ID3D11DeviceContext* context = nullptr;
IDXGISwapChain* swapChain = nullptr;
ID3D11RenderTargetView* rtv = nullptr;
ID3D11VertexShader* vs = nullptr;
ID3D11PixelShader* ps = nullptr;
ID3D11InputLayout* inputLayout = nullptr;
ID3D11Buffer* vertexBuffer = nullptr;

// Vertex struct
struct Vertex { float x, y; };

// Vertex shader
const char* vsSrc = R"(
struct VS_INPUT { float2 pos : POSITION; };
struct PS_INPUT { float4 pos : SV_POSITION; float2 uv : TEXCOORD; };
PS_INPUT main(VS_INPUT input) {
    PS_INPUT output;
    output.pos = float4(input.pos,0,1);
    output.uv = input.pos*0.5+0.5;
    return output;
}
)";

char psSrc[2048]; // Pixel shader

// Mandelbrot view presets
struct View { float xMin, xMax, yMin, yMax; };
View views[] = {
    {-2.5f, 1.0f, -1.2f, 1.2f},       // Full set
    {-0.748f, -0.743f, 0.1f, 0.105f}, // Seahorse Valley
    {-0.74877f, -0.74872f, 0.06505f, 0.06510f}, // Mini zoom
};
int currentView = 0;

// Forward declarations
LRESULT CALLBACK WndProc(HWND, UINT, WPARAM, LPARAM);
bool InitWindow(HINSTANCE, int);
bool InitD3D();
void CleanUp();
bool CompileShaders();
void Render();

// Main
int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE, LPSTR, int nCmdShow) {
    if (!InitWindow(hInstance, nCmdShow)) return 0;
    if (!InitD3D()) return 0;

    CompileShaders();

    MSG msg = {};
    while (msg.message != WM_QUIT) {
        if (PeekMessage(&msg, nullptr, 0, 0, PM_REMOVE)) {
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        }
        else {
            Render(); // non-blocking rendering
        }
    }

    CleanUp();
    return 0;
}

// --- Window ---
bool InitWindow(HINSTANCE hInstance, int nCmdShow) {
    const char CLASS_NAME[] = "MandelbrotDX11";
    WNDCLASS wc = {};
    wc.lpfnWndProc = WndProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = CLASS_NAME;
    RegisterClass(&wc);

    hwnd = CreateWindowEx(0, CLASS_NAME, "DX11 Mandelbrot", WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT, 800, 600, nullptr, nullptr, hInstance, nullptr);
    if (!hwnd) return false;
    ShowWindow(hwnd, nCmdShow);
    return true;
}

// --- WinProc ---
LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    switch (msg) {
    case WM_DESTROY: PostQuitMessage(0); return 0;
    case WM_KEYDOWN:
        if (wParam == VK_ESCAPE) PostQuitMessage(0);
        if (wParam == VK_LEFT) currentView = (currentView + 2) % 3;
        if (wParam == VK_RIGHT) currentView = (currentView + 1) % 3;
        CompileShaders();
        return 0;
    }
    return DefWindowProc(hwnd, msg, wParam, lParam);
}

// --- D3D11 Init ---
bool InitD3D() {
    DXGI_SWAP_CHAIN_DESC scd = {};
    scd.BufferCount = 1;
    scd.BufferDesc.Width = 800;
    scd.BufferDesc.Height = 600;
    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.Windowed = TRUE;

    D3D_FEATURE_LEVEL featureLevel;
    if (FAILED(D3D11CreateDeviceAndSwapChain(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr,
        0, nullptr, 0, D3D11_SDK_VERSION, &scd,
        &swapChain, &device, &featureLevel, &context))) return false;

    ID3D11Texture2D* backBuffer = nullptr;
    swapChain->GetBuffer(0, __uuidof(ID3D11Texture2D), (void**)&backBuffer);
    device->CreateRenderTargetView(backBuffer, nullptr, &rtv);
    backBuffer->Release();
    context->OMSetRenderTargets(1, &rtv, nullptr);

    D3D11_VIEWPORT vp = {};
    vp.Width = 800; vp.Height = 600; vp.MinDepth = 0; vp.MaxDepth = 1;
    context->RSSetViewports(1, &vp);

    Vertex verts[] = { {-1,-1},{-1,1},{1,-1},{1,1} };
    D3D11_BUFFER_DESC bd = {};
    bd.Usage = D3D11_USAGE_DEFAULT;
    bd.ByteWidth = sizeof(verts);
    bd.BindFlags = D3D11_BIND_VERTEX_BUFFER;
    D3D11_SUBRESOURCE_DATA initData = {};
    initData.pSysMem = verts;
    device->CreateBuffer(&bd, &initData, &vertexBuffer);

    return true;
}

// --- Compile shaders ---
bool CompileShaders() {
    // Pixel shader template with smooth coloring
    snprintf(psSrc, sizeof(psSrc),
        R"(struct PS_INPUT { float4 pos:SV_POSITION; float2 uv: TEXCOORD; };
float mandelbrot(float2 c){
    float2 z=0;
    float iter=0;
    const int MAX_ITER=256;
    for(int i=0;i<MAX_ITER;i++){
        float x=z.x*z.x - z.y*z.y + (c.x*%.8f+%.8f);
        float y=2*z.x*z.y + (c.y*%.8f+%.8f);
        z=float2(x,y);
        if(dot(z,z)>4){iter=i; break;}
    }
    return iter;
}
float4 main(PS_INPUT input):SV_TARGET{
    float iter=mandelbrot(input.uv);
    const float MAX_ITER=256.0;
    // Original color mapping (smooth)
    float t = iter / MAX_ITER;
    float r = clamp(9.0*(1.0-t)*t*t*t,0,1);
    float g = clamp(15.0*(1.0-t)*(1.0-t)*t*t,0,1);
    float b = clamp(8.5*(1.0-t)*(1.0-t)*(1.0-t)*t,0,1);
    return float4(r,g,b,1.0);
})",
(views[currentView].xMax - views[currentView].xMin), views[currentView].xMin,
(views[currentView].yMax - views[currentView].yMin), views[currentView].yMin
);

    ID3DBlob* vsBlob = nullptr, * psBlob = nullptr, * err = nullptr;

    if (FAILED(D3DCompile(vsSrc, strlen(vsSrc), nullptr, nullptr, nullptr, "main", "vs_5_0", 0, 0, &vsBlob, &err))) {
        if (err) { OutputDebugStringA((char*)err->GetBufferPointer()); err->Release(); } return false;
    }
    if (FAILED(D3DCompile(psSrc, strlen(psSrc), nullptr, nullptr, nullptr, "main", "ps_5_0", 0, 0, &psBlob, &err))) {
        if (err) { OutputDebugStringA((char*)err->GetBufferPointer()); err->Release(); } return false;
    }

    if (vs) vs->Release();
    if (ps) ps->Release();
    if (inputLayout) inputLayout->Release();

    device->CreateVertexShader(vsBlob->GetBufferPointer(), vsBlob->GetBufferSize(), nullptr, &vs);
    device->CreatePixelShader(psBlob->GetBufferPointer(), psBlob->GetBufferSize(), nullptr, &ps);

    D3D11_INPUT_ELEMENT_DESC layout[] = { {"POSITION",0,DXGI_FORMAT_R32G32_FLOAT,0,0,D3D11_INPUT_PER_VERTEX_DATA,0} };
    device->CreateInputLayout(layout, 1, vsBlob->GetBufferPointer(), vsBlob->GetBufferSize(), &inputLayout);

    vsBlob->Release(); psBlob->Release();
    return true;
}

// --- Render ---
void Render() {
    float clear[4] = { 0,0,0,1 };
    context->ClearRenderTargetView(rtv, clear);

    UINT stride = sizeof(Vertex), offset = 0;
    context->IASetVertexBuffers(0, 1, &vertexBuffer, &stride, &offset);
    context->IASetInputLayout(inputLayout);
    context->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLESTRIP);

    context->VSSetShader(vs, nullptr, 0);
    context->PSSetShader(ps, nullptr, 0);

    context->Draw(4, 0);
    swapChain->Present(0, 0); // non-blocking present
}

// --- Cleanup ---
void CleanUp() {
    if (vertexBuffer) vertexBuffer->Release();
    if (inputLayout) inputLayout->Release();
    if (vs) vs->Release();
    if (ps) ps->Release();
    if (rtv) rtv->Release();
    if (swapChain) swapChain->Release();
    if (context) context->Release();
    if (device) device->Release();
}