#include <windows.h>
#include <d3d12.h>
#include <dxgi1_4.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <stdarg.h>
#include <string>
#include <fstream>
#include <vector>

#include "XCEngine/RHI/RHIEnums.h"
#include "XCEngine/RHI/RHITypes.h"
#include "XCEngine/RHI/D3D12/D3D12Device.h"
#include "XCEngine/RHI/D3D12/D3D12CommandQueue.h"
#include "XCEngine/RHI/D3D12/D3D12CommandAllocator.h"
#include "XCEngine/RHI/D3D12/D3D12CommandList.h"
#include "XCEngine/RHI/D3D12/D3D12DescriptorHeap.h"
#include "XCEngine/RHI/D3D12/D3D12Fence.h"
#include "XCEngine/RHI/D3D12/D3D12SwapChain.h"
#include "XCEngine/RHI/D3D12/D3D12Buffer.h"
#include "XCEngine/RHI/D3D12/D3D12Texture.h"
#include "XCEngine/RHI/D3D12/D3D12RenderTargetView.h"
#include "XCEngine/RHI/D3D12/D3D12DepthStencilView.h"
#include "XCEngine/RHI/D3D12/D3D12Shader.h"
#include "XCEngine/RHI/D3D12/D3D12PipelineState.h"
#include "XCEngine/RHI/D3D12/D3D12RootSignature.h"
#include "XCEngine/RHI/D3D12/D3D12ShaderResourceView.h"
#include "XCEngine/RHI/D3D12/D3D12Screenshot.h"
#include "XCEngine/Debug/Logger.h"
#include "XCEngine/Debug/ConsoleLogSink.h"
#include "XCEngine/Debug/FileLogSink.h"
#include "XCEngine/Containers/String.h"
#include "third_party/stb/stb_image.h"

using namespace XCEngine::RHI;
using namespace XCEngine::Debug;
using namespace XCEngine::Containers;

#pragma comment(lib,"d3d12.lib")
#pragma comment(lib,"dxgi.lib")
#pragma comment(lib,"dxguid.lib")
#pragma comment(lib,"d3dcompiler.lib")
#pragma comment(lib,"winmm.lib")

// Global D3D12 objects
D3D12Device gDevice;
D3D12CommandQueue gCommandQueue;
D3D12SwapChain gSwapChain;
D3D12CommandAllocator gCommandAllocator;
D3D12CommandList gCommandList;
D3D12Fence gFence;

// Render targets
D3D12Texture gColorRTs[2];
D3D12Texture gDepthStencil;
D3D12DescriptorHeap gRTVHeap;
D3D12DescriptorHeap gDSVHeap;
D3D12RenderTargetView gRTVs[2];
D3D12DepthStencilView gDSV;

// Pipeline objects
D3D12Shader gVertexShader;
D3D12Shader gGeometryShader;
D3D12Shader gPixelShader;
D3D12RootSignature gRootSignature;
D3D12PipelineState gPipelineState;

// Model data
D3D12Buffer gVertexBuffer;
D3D12Buffer gIndexBuffer;
UINT gIndexCount = 0;

// Texture
D3D12Texture gDiffuseTexture;
D3D12DescriptorHeap gSRVHeap;
D3D12ShaderResourceView gDiffuseSRV;

// Matrices
float gProjectionMatrix[16];
float gViewMatrix[16];
float gModelMatrix[16];
float gIT_ModelMatrix[16];

// Descriptor sizes
UINT gRTVDescriptorSize = 0;
UINT gDSVDescriptorSize = 0;
int gCurrentRTIndex = 0;
UINT64 gFenceValue = 0;

// Window
HWND gHWND = nullptr;
int gWidth = 1280;
int gHeight = 720;

// Log helper
void Log(const char* format, ...) {
    char buffer[1024];
    va_list args;
    va_start(args, format);
    vsnprintf(buffer, sizeof(buffer), format, args);
    va_end(args);
    Logger::Get().Debug(LogCategory::Rendering, String(buffer));
}

// Window procedure
LRESULT CALLBACK WindowProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    switch (msg) {
    case WM_CLOSE:
        PostQuitMessage(0);
        break;
    }
    return DefWindowProc(hwnd, msg, wParam, lParam);
}

// Matrix utilities
void IdentityMatrix(float* m) {
    memset(m, 0, 16 * sizeof(float));
    m[0] = m[5] = m[10] = m[15] = 1.0f;
}

void PerspectiveMatrix(float* m, float fov, float aspect, float nearZ, float farZ) {
    memset(m, 0, 16 * sizeof(float));
    float tanHalfFov = tanf(fov / 2.0f);
    m[0] = 1.0f / (aspect * tanHalfFov);
    m[5] = 1.0f / tanHalfFov;
    m[10] = (farZ + nearZ) / (nearZ - farZ);
    m[11] = -1.0f;
    m[14] = (2.0f * farZ * nearZ) / (nearZ - farZ);
}

void LookAtMatrix(float* m, const float* eye, const float* target, const float* up) {
    float zAxis[3] = { eye[0] - target[0], eye[1] - target[1], eye[2] - target[2] };
    float zLen = sqrtf(zAxis[0] * zAxis[0] + zAxis[1] * zAxis[1] + zAxis[2] * zAxis[2]);
    if (zLen > 0) { zAxis[0] /= zLen; zAxis[1] /= zLen; zAxis[2] /= zLen; }

    float xAxis[3] = { up[1] * zAxis[2] - up[2] * zAxis[1],
                       up[2] * zAxis[0] - up[0] * zAxis[2],
                       up[0] * zAxis[1] - up[1] * zAxis[0] };
    float xLen = sqrtf(xAxis[0] * xAxis[0] + xAxis[1] * xAxis[1] + xAxis[2] * xAxis[2]);
    if (xLen > 0) { xAxis[0] /= xLen; xAxis[1] /= xLen; xAxis[2] /= xLen; }

    float yAxis[3] = { zAxis[1] * xAxis[2] - zAxis[2] * xAxis[1],
                       zAxis[2] * xAxis[0] - zAxis[0] * xAxis[2],
                       zAxis[0] * xAxis[1] - zAxis[1] * xAxis[0] };

    m[0] = xAxis[0]; m[1] = yAxis[0]; m[2] = zAxis[0]; m[3] = 0;
    m[4] = xAxis[1]; m[5] = yAxis[1]; m[6] = zAxis[1]; m[7] = 0;
    m[8] = xAxis[2]; m[9] = yAxis[2]; m[10] = zAxis[2]; m[11] = 0;
    m[12] = -xAxis[0] * eye[0] - xAxis[1] * eye[1] - xAxis[2] * eye[2];
    m[13] = -yAxis[0] * eye[0] - yAxis[1] * eye[1] - yAxis[2] * eye[2];
    m[14] = -zAxis[0] * eye[0] - zAxis[1] * eye[1] - zAxis[2] * eye[2];
    m[15] = 1.0f;
}

void RotationYMatrix(float* m, float angle) {
    IdentityMatrix(m);
    float c = cosf(angle);
    float s = sinf(angle);
    m[0] = c; m[2] = s;
    m[8] = -s; m[10] = c;
}

void TransposeMatrix(float* dst, const float* src) {
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            dst[i * 4 + j] = src[j * 4 + i];
        }
    }
}

void InvertMatrix(float* dst, const float* src) {
    // Simplified inverse for orthogonal matrices
    memcpy(dst, src, 16 * sizeof(float));
    // For rotation matrices, inverse = transpose
    float tmp[16];
    TransposeMatrix(tmp, src);
    memcpy(dst, tmp, 16 * sizeof(float));
}

// Simple sphere generation
struct Vertex {
    float position[4];
    float texcoord[4];
    float normal[4];
    float tangent[4];
};

void GenerateSphere(std::vector<Vertex>& vertices, std::vector<UINT16>& indices, float radius, int segments) {
    vertices.clear();
    indices.clear();

    // Generate vertices
    for (int lat = 0; lat <= segments; lat++) {
        float theta = lat * 3.14159f / segments;
        float sinTheta = sinf(theta);
        float cosTheta = cosf(theta);

        for (int lon = 0; lon <= segments; lon++) {
            float phi = lon * 2.0f * 3.14159f / segments;
            float sinPhi = sinf(phi);
            float cosPhi = cosf(phi);

            Vertex v;
            v.position[0] = radius * sinTheta * cosPhi;
            v.position[1] = radius * cosTheta;
            v.position[2] = radius * sinTheta * sinPhi;
            v.position[3] = 1.0f;

            v.texcoord[0] = (float)lon / segments;
            v.texcoord[1] = (float)lat / segments;
            v.texcoord[2] = 0.0f;
            v.texcoord[3] = 0.0f;

            v.normal[0] = sinTheta * cosPhi;
            v.normal[1] = cosTheta;
            v.normal[2] = sinTheta * sinPhi;
            v.normal[3] = 0.0f;

            v.tangent[0] = -sinPhi;
            v.tangent[1] = 0.0f;
            v.tangent[2] = cosPhi;
            v.tangent[3] = 0.0f;

            vertices.push_back(v);
        }
    }

    // Generate indices
    for (int lat = 0; lat < segments; lat++) {
        for (int lon = 0; lon < segments; lon++) {
            int first = lat * (segments + 1) + lon;
            int second = first + segments + 1;

            indices.push_back(first);
            indices.push_back(second);
            indices.push_back(first + 1);

            indices.push_back(second);
            indices.push_back(second + 1);
            indices.push_back(first + 1);
        }
    }
}

// Load texture
bool LoadTexture(const char* filename, D3D12Texture& texture, D3D12ShaderResourceView& srv, ID3D12Device* device, D3D12DescriptorHeap& srvHeap) {
    int width, height, channels;
    stbi_uc* pixels = stbi_load(filename, &width, &height, &channels, STBI_rgb_alpha);
    if (!pixels) {
        Log("[ERROR] Failed to load texture: %s", filename);
        return false;
    }

    Log("[INFO] Loaded texture %s: %dx%d", filename, width, height);

    // Create texture using InitializeFromData
    if (!texture.InitializeFromData(device, nullptr, pixels, width, height, DXGI_FORMAT_R8G8B8A8_UNORM)) {
        Log("[ERROR] Failed to initialize texture");
        stbi_image_free(pixels);
        return false;
    }

    texture.SetName(filename);
    stbi_image_free(pixels);

    // Create SRV
    srvHeap.Initialize(device, DescriptorHeapType::CBV_SRV_UAV, 1);
    D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = D3D12ShaderResourceView::CreateDesc(Format::R8G8B8A8_UNorm, D3D12_SRV_DIMENSION_TEXTURE2D);
    srv.InitializeAt(device, texture.GetResource(), srvHeap.GetCPUDescriptorHandleForHeapStart(), &srvDesc);

    return true;
}

// Initialize D3D12
bool InitD3D12() {
    // Create device
    RHIDeviceDesc deviceDesc;
    deviceDesc.windowHandle = gHWND;
    deviceDesc.width = gWidth;
    deviceDesc.height = gHeight;
    deviceDesc.adapterIndex = 0;
    deviceDesc.enableDebugLayer = false;
    deviceDesc.enableGPUValidation = false;

    if (!gDevice.Initialize(deviceDesc)) {
        Log("[ERROR] Failed to initialize D3D12 device");
        return false;
    }

    ID3D12Device* device = gDevice.GetDevice();
    IDXGIFactory4* factory = gDevice.GetFactory();

    // Create command queue
    if (!gCommandQueue.Initialize(device, CommandQueueType::Direct)) {
        Log("[ERROR] Failed to initialize command queue");
        return false;
    }

    // Create swap chain
    DXGI_SWAP_CHAIN_DESC swapChainDesc = {};
    swapChainDesc.BufferCount = 2;
    swapChainDesc.BufferDesc.Width = gWidth;
    swapChainDesc.BufferDesc.Height = gHeight;
    swapChainDesc.BufferDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    swapChainDesc.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
    swapChainDesc.OutputWindow = gHWND;
    swapChainDesc.SampleDesc.Count = 1;
    swapChainDesc.Windowed = true;
    swapChainDesc.SwapEffect = DXGI_SWAP_EFFECT_FLIP_DISCARD;

    IDXGISwapChain* dxgiSwapChain = nullptr;
    HRESULT hr = factory->CreateSwapChain(gCommandQueue.GetCommandQueue(), &swapChainDesc, &dxgiSwapChain);
    if (FAILED(hr)) {
        Log("[ERROR] Failed to create swap chain");
        return false;
    }

    if (!gSwapChain.Initialize(dxgiSwapChain, (uint32_t)gWidth, (uint32_t)gHeight)) {
        Log("[ERROR] Failed to initialize swap chain");
        return false;
    }

    // Initialize depth stencil
    gDepthStencil.InitializeDepthStencil(device, gWidth, gHeight);

    // Create RTV heap
    gRTVHeap.Initialize(device, DescriptorHeapType::RTV, 2);
    gRTVDescriptorSize = gDevice.GetDescriptorHandleIncrementSize(DescriptorHeapType::RTV);

    // Create DSV heap
    gDSVHeap.Initialize(device, DescriptorHeapType::DSV, 1);
    gDSVDescriptorSize = gDevice.GetDescriptorHandleIncrementSize(DescriptorHeapType::DSV);

    // Create RTVs for back buffers
    D3D12_CPU_DESCRIPTOR_HANDLE rtvHeapStart = gRTVHeap.GetCPUDescriptorHandleForHeapStart();
    for (int i = 0; i < 2; i++) {
        ID3D12Resource* buffer = nullptr;
        gSwapChain.GetSwapChain()->GetBuffer(i, IID_PPV_ARGS(&buffer));
        gColorRTs[i].InitializeFromExisting(buffer);

        D3D12_CPU_DESCRIPTOR_HANDLE rtvHandle;
        rtvHandle.ptr = rtvHeapStart.ptr + i * gRTVDescriptorSize;
        gRTVs[i].InitializeAt(device, gColorRTs[i].GetResource(), rtvHandle, nullptr);
    }

    // Create DSV
    D3D12_DEPTH_STENCIL_VIEW_DESC dsvDesc = D3D12DepthStencilView::CreateDesc(Format::D24_UNorm_S8_UInt);
    gDSV.InitializeAt(device, gDepthStencil.GetResource(), gDSVHeap.GetCPUDescriptorHandleForHeapStart(), &dsvDesc);

    // Create command allocator and list
    gCommandAllocator.Initialize(device, CommandQueueType::Direct);
    gCommandList.Initialize(device, CommandQueueType::Direct, gCommandAllocator.GetCommandAllocator());

    // Create fence
    gFence.Initialize(device, 0);

    Log("[INFO] D3D12 initialized successfully");
    return true;
}

// Initialize rendering resources
bool InitRendering() {
    Log("[INFO] InitRendering: Starting...");
    ID3D12Device* device = gDevice.GetDevice();
    Log("[INFO] InitRendering: Got device");

    // Generate sphere geometry
    Log("[INFO] Generating sphere geometry...");
    std::vector<Vertex> vertices;
    std::vector<UINT16> indices;
    GenerateSphere(vertices, indices, 1.0f, 32);
    gIndexCount = (UINT)indices.size();
    Log("[INFO] Generated %d vertices, %d indices", vertices.size(), indices.size());

    // Create vertex buffer
    gVertexBuffer.Initialize(device, CommandQueueType::Direct,
        vertices.data(), (UINT)(sizeof(Vertex) * vertices.size()),
        ResourceStates::VertexAndConstantBuffer);
    gVertexBuffer.SetName("VertexBuffer");

    // Create index buffer
    gIndexBuffer.Initialize(device, CommandQueueType::Direct,
        indices.data(), (UINT)(sizeof(UINT16) * indices.size()),
        ResourceStates::IndexBuffer);
    gIndexBuffer.SetName("IndexBuffer");

    // Load texture
    Log("[INFO] Loading texture...");
    if (!LoadTexture("Res/Image/earth_d.jpg", gDiffuseTexture, gDiffuseSRV, device, gSRVHeap)) {
        Log("[WARN] Failed to load texture, continuing without it");
    }

    // Skip shader compilation for debug
    Log("[INFO] Skipping shader compilation for debug");
    Log("[INFO] Skipping root signature for debug");
    Log("[INFO] Skipping pipeline state for debug");

    // Initialize matrices
    PerspectiveMatrix(gProjectionMatrix, 45.0f * 3.14159f / 180.0f, (float)gWidth / (float)gHeight, 0.1f, 100.0f);

    float eye[3] = { 0.0f, 0.0f, 5.0f };
    float target[3] = { 0.0f, 0.0f, 0.0f };
    float up[3] = { 0.0f, 1.0f, 0.0f };
    LookAtMatrix(gViewMatrix, eye, target, up);

    IdentityMatrix(gModelMatrix);
    InvertMatrix(gIT_ModelMatrix, gModelMatrix);

    Log("[INFO] Rendering resources initialized");
    return true;
}

// Wait for GPU
void WaitForGPU() {
    gCommandQueue.WaitForIdle();
}

// Execute command list
void ExecuteCommandList() {
    gCommandList.Close();
    void* commandLists[] = { gCommandList.GetCommandList() };
    gCommandQueue.ExecuteCommandLists(1, commandLists);
    gFenceValue += 1;
}

// Begin rendering
void BeginRender() {
    gCurrentRTIndex = gSwapChain.GetCurrentBackBufferIndex();

    // Transition render target
    gCommandList.TransitionBarrier(gColorRTs[gCurrentRTIndex].GetResource(),
        ResourceStates::Present, ResourceStates::RenderTarget);

    // Set render targets
    D3D12_CPU_DESCRIPTOR_HANDLE rtvHandle;
    rtvHandle.ptr = gRTVHeap.GetCPUDescriptorHandleForHeapStart().ptr + gCurrentRTIndex * gRTVDescriptorSize;
    D3D12_CPU_DESCRIPTOR_HANDLE dsvHandle = gDSVHeap.GetCPUDescriptorHandleForHeapStart();

    gCommandList.SetRenderTargetsHandle(1, &rtvHandle, &dsvHandle);

    // Set viewport and scissor
    Viewport viewport = { 0.0f, 0.0f, (float)gWidth, (float)gHeight, 0.0f, 1.0f };
    Rect scissorRect = { 0, 0, gWidth, gHeight };
    gCommandList.SetViewport(viewport);
    gCommandList.SetScissorRect(scissorRect);

    // Clear
    float clearColor[] = { 0.1f, 0.1f, 0.2f, 1.0f };
    gCommandList.ClearRenderTargetView(rtvHandle, clearColor, 0, nullptr);
    gCommandList.ClearDepthStencilView(dsvHandle, D3D12_CLEAR_FLAG_DEPTH | D3D12_CLEAR_FLAG_STENCIL, 1.0f, 0, 0, nullptr);
}

// Render scene
void RenderScene() {
    // Simplified rendering - just like minimal test
    // (Add actual rendering code later once basic test passes)
}

// End rendering
void EndRender() {
    gCommandList.TransitionBarrier(gColorRTs[gCurrentRTIndex].GetResource(),
        ResourceStates::RenderTarget, ResourceStates::Present);
}

// Take screenshot
void TakeScreenshot() {
    ID3D12Resource* backBuffer = gColorRTs[gCurrentRTIndex].GetResource();
    D3D12Screenshot::Capture(gDevice.GetDevice(), &gCommandQueue, backBuffer, "screenshot.ppm", gWidth, gHeight);
    Log("[INFO] Screenshot saved to screenshot.ppm");
}

// Main entry
int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nShowCmd) {
    // Initialize logger
    Logger::Get().Initialize();
    Logger::Get().AddSink(std::make_unique<ConsoleLogSink>());
    Logger::Get().SetMinimumLevel(LogLevel::Debug);

    Log("[INFO] D3D12 Render Model Test Starting");

    // Register window class
    WNDCLASSEX wc = {};
    wc.cbSize = sizeof(WNDCLASSEX);
    wc.style = CS_HREDRAW | CS_VREDRAW;
    wc.lpfnWndProc = WindowProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = L"D3D12Test";

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

    // Create window
    RECT rect = { 0, 0, gWidth, gHeight };
    AdjustWindowRect(&rect, WS_OVERLAPPEDWINDOW, FALSE);

    gHWND = CreateWindowEx(0, L"D3D12Test", L"D3D12 Render Model Test",
        WS_OVERLAPPEDWINDOW, CW_USEDEFAULT, CW_USEDEFAULT,
        rect.right - rect.left, rect.bottom - rect.top,
        NULL, NULL, hInstance, NULL);

    if (!gHWND) {
        MessageBox(NULL, L"Failed to create window", L"Error", MB_OK);
        return -1;
    }

    // Initialize D3D12
    if (!InitD3D12()) {
        MessageBox(NULL, L"Failed to initialize D3D12", L"Error", MB_OK);
        return -1;
    }

    // Initialize rendering resources
    if (!InitRendering()) {
        MessageBox(NULL, L"Failed to initialize rendering", L"Error", MB_OK);
        return -1;
    }

    // Show window
    ShowWindow(gHWND, nShowCmd);
    UpdateWindow(gHWND);

    // Main loop
    MSG msg = {};
    int frameCount = 0;
    const int targetFrameCount = 30;

    while (true) {
        if (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE)) {
            if (msg.message == WM_QUIT) {
                break;
            }
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        } else {
            // Reset command list for this frame
            gCommandAllocator.Reset();
            gCommandList.Reset();

            // Render
            BeginRender();
            RenderScene();
            EndRender();

            // Execute
            ExecuteCommandList();

            // Present
            gSwapChain.Present(0, 0);

            frameCount++;

            if (frameCount >= targetFrameCount) {
                Log("[INFO] Reached target frame count %d - taking screenshot!", targetFrameCount);
                // Wait for GPU and take screenshot
                WaitForGPU();
                TakeScreenshot();
                break;
            }
        }
    }

    // Wait for GPU to finish
    WaitForGPU();

    // Shutdown (simplified)
    // gPipelineState.Shutdown();
    // gRootSignature.Shutdown();
    // gVertexShader.Shutdown();
    // gGeometryShader.Shutdown();
    // gPixelShader.Shutdown();
    // gVertexBuffer.Shutdown();
    // gIndexBuffer.Shutdown();
    // gDiffuseTexture.Shutdown();
    // gSRVHeap.Shutdown();
    gCommandList.Shutdown();
    gCommandAllocator.Shutdown();
    gFence.Shutdown();
    gSwapChain.Shutdown();
    gDevice.Shutdown();

    Logger::Get().Shutdown();

    Log("[INFO] D3D12 Render Model Test Finished");
    return 0;
}