#include "renderer.h"
#include "waterfall.h"
#include <fstream>
#include <sstream>
#include <iostream>

Renderer::Renderer(WGPUDevice device, WGPUSurface surface, int width, int height)
    : device_(device), surface_(surface), width_(width), height_(height), time_(0.0f),
      surfaceFormat_(WGPUTextureFormat_BGRA8Unorm), linePipeline_(nullptr), lineListPipeline_(nullptr) {}

Renderer::~Renderer() {
    if (linePipeline_) wgpuRenderPipelineRelease(linePipeline_);
    if (lineListPipeline_) wgpuRenderPipelineRelease(lineListPipeline_);
}

bool Renderer::initialize() {
    configureSurface();
    createPipelines();

    // Create two waterfall views side-by-side
    waterfalls_.push_back(std::make_unique<Waterfall>(
        device_, 0.0f, 0.1f, 0.5f, 0.9f, "Frequency vs Time"));
    waterfalls_.push_back(std::make_unique<Waterfall>(
        device_, 0.5f, 0.1f, 0.5f, 0.9f, "Position vs Time"));

    for (auto& waterfall : waterfalls_) {
        if (!waterfall->initialize()) {
            return false;
        }
    }

    return true;
}

void Renderer::configureSurface() {
    WGPUSurfaceConfiguration config = {};
    config.nextInChain = nullptr;
    config.device = device_;
    config.format = surfaceFormat_;
    config.usage = WGPUTextureUsage_RenderAttachment;
    config.width = width_;
    config.height = height_;
    config.presentMode = WGPUPresentMode_Fifo;
    config.alphaMode = WGPUCompositeAlphaMode_Auto;
    config.viewFormatCount = 0;
    config.viewFormats = nullptr;

    wgpuSurfaceConfigure(surface_, &config);
}

void Renderer::createPipelines() {
    // Load shader
    std::ifstream shaderFile("shaders/waterfall.wgsl");
    std::stringstream buffer;
    buffer << shaderFile.rdbuf();
    std::string shaderCode = buffer.str();

    WGPUShaderSourceWGSL wgslSource = {};
    wgslSource.chain.sType = WGPUSType_ShaderSourceWGSL;
    wgslSource.chain.next = nullptr;
    wgslSource.code = {shaderCode.c_str(), shaderCode.length()};

    WGPUShaderModuleDescriptor shaderDesc = {};
    shaderDesc.nextInChain = &wgslSource.chain;
    shaderDesc.label = {nullptr, WGPU_STRLEN};

    WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device_, &shaderDesc);

    // Vertex buffer layout
    WGPUVertexAttribute attributes[2] = {};
    attributes[0].format = WGPUVertexFormat_Float32x2;
    attributes[0].offset = 0;
    attributes[0].shaderLocation = 0;

    attributes[1].format = WGPUVertexFormat_Float32x3;
    attributes[1].offset = 2 * sizeof(float);
    attributes[1].shaderLocation = 1;

    WGPUVertexBufferLayout vertexBufferLayout = {};
    vertexBufferLayout.arrayStride = 5 * sizeof(float);
    vertexBufferLayout.stepMode = WGPUVertexStepMode_Vertex;
    vertexBufferLayout.attributeCount = 2;
    vertexBufferLayout.attributes = attributes;

    // Pipeline layout
    WGPUPipelineLayoutDescriptor layoutDesc = {};
    layoutDesc.nextInChain = nullptr;
    layoutDesc.label = {nullptr, WGPU_STRLEN};
    layoutDesc.bindGroupLayoutCount = 0;
    layoutDesc.bindGroupLayouts = nullptr;

    WGPUPipelineLayout pipelineLayout = wgpuDeviceCreatePipelineLayout(device_, &layoutDesc);

    // Color target
    WGPUBlendState blend = {};
    blend.color.operation = WGPUBlendOperation_Add;
    blend.color.srcFactor = WGPUBlendFactor_One;
    blend.color.dstFactor = WGPUBlendFactor_Zero;
    blend.alpha.operation = WGPUBlendOperation_Add;
    blend.alpha.srcFactor = WGPUBlendFactor_One;
    blend.alpha.dstFactor = WGPUBlendFactor_Zero;

    WGPUColorTargetState colorTarget = {};
    colorTarget.nextInChain = nullptr;
    colorTarget.format = surfaceFormat_;
    colorTarget.blend = &blend;
    colorTarget.writeMask = WGPUColorWriteMask_All;

    WGPUFragmentState fragmentState = {};
    fragmentState.nextInChain = nullptr;
    fragmentState.module = shader;
    fragmentState.entryPoint = {"fs_main", 7};
    fragmentState.constantCount = 0;
    fragmentState.constants = nullptr;
    fragmentState.targetCount = 1;
    fragmentState.targets = &colorTarget;

    // Line strip pipeline
    WGPURenderPipelineDescriptor pipelineDesc = {};
    pipelineDesc.nextInChain = nullptr;
    pipelineDesc.label = {nullptr, WGPU_STRLEN};
    pipelineDesc.layout = pipelineLayout;

    pipelineDesc.vertex.nextInChain = nullptr;
    pipelineDesc.vertex.module = shader;
    pipelineDesc.vertex.entryPoint = {"vs_main", 7};
    pipelineDesc.vertex.constantCount = 0;
    pipelineDesc.vertex.constants = nullptr;
    pipelineDesc.vertex.bufferCount = 1;
    pipelineDesc.vertex.buffers = &vertexBufferLayout;

    pipelineDesc.primitive.nextInChain = nullptr;
    pipelineDesc.primitive.topology = WGPUPrimitiveTopology_LineStrip;
    pipelineDesc.primitive.stripIndexFormat = WGPUIndexFormat_Undefined;
    pipelineDesc.primitive.frontFace = WGPUFrontFace_CCW;
    pipelineDesc.primitive.cullMode = WGPUCullMode_None;

    pipelineDesc.depthStencil = nullptr;
    pipelineDesc.multisample.nextInChain = nullptr;
    pipelineDesc.multisample.count = 1;
    pipelineDesc.multisample.mask = ~0u;
    pipelineDesc.multisample.alphaToCoverageEnabled = false;

    pipelineDesc.fragment = &fragmentState;

    linePipeline_ = wgpuDeviceCreateRenderPipeline(device_, &pipelineDesc);

    // Line list pipeline for grid
    pipelineDesc.primitive.topology = WGPUPrimitiveTopology_LineList;
    lineListPipeline_ = wgpuDeviceCreateRenderPipeline(device_, &pipelineDesc);

    // Cleanup
    wgpuPipelineLayoutRelease(pipelineLayout);
    wgpuShaderModuleRelease(shader);
}

void Renderer::update() {
    time_ += 0.016f; // ~60fps

    for (auto& waterfall : waterfalls_) {
        waterfall->update(time_);
    }
}

void Renderer::render() {
    WGPUSurfaceTexture surfaceTexture;
    wgpuSurfaceGetCurrentTexture(surface_, &surfaceTexture);

    if (surfaceTexture.status != WGPUSurfaceGetCurrentTextureStatus_SuccessOptimal &&
        surfaceTexture.status != WGPUSurfaceGetCurrentTextureStatus_SuccessSuboptimal) {
        std::cerr << "Failed to get surface texture" << std::endl;
        return;
    }

    WGPUTextureViewDescriptor viewDesc = {};
    viewDesc.nextInChain = nullptr;
    viewDesc.label = {nullptr, WGPU_STRLEN};
    viewDesc.format = surfaceFormat_;
    viewDesc.dimension = WGPUTextureViewDimension_2D;
    viewDesc.baseMipLevel = 0;
    viewDesc.mipLevelCount = 1;
    viewDesc.baseArrayLayer = 0;
    viewDesc.arrayLayerCount = 1;
    viewDesc.aspect = WGPUTextureAspect_All;

    WGPUTextureView textureView = wgpuTextureCreateView(surfaceTexture.texture, &viewDesc);

    WGPUCommandEncoderDescriptor encoderDesc = {};
    encoderDesc.nextInChain = nullptr;
    encoderDesc.label = {nullptr, WGPU_STRLEN};

    WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device_, &encoderDesc);

    WGPURenderPassColorAttachment colorAttachment = {};
    colorAttachment.nextInChain = nullptr;
    colorAttachment.view = textureView;
    colorAttachment.depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;
    colorAttachment.resolveTarget = nullptr;
    colorAttachment.loadOp = WGPULoadOp_Clear;
    colorAttachment.storeOp = WGPUStoreOp_Store;
    colorAttachment.clearValue = {0.1, 0.1, 0.15, 1.0};

    WGPURenderPassDescriptor renderPassDesc = {};
    renderPassDesc.nextInChain = nullptr;
    renderPassDesc.label = {nullptr, WGPU_STRLEN};
    renderPassDesc.colorAttachmentCount = 1;
    renderPassDesc.colorAttachments = &colorAttachment;
    renderPassDesc.depthStencilAttachment = nullptr;
    renderPassDesc.occlusionQuerySet = nullptr;
    renderPassDesc.timestampWrites = nullptr;

    WGPURenderPassEncoder pass = wgpuCommandEncoderBeginRenderPass(encoder, &renderPassDesc);

    for (auto& waterfall : waterfalls_) {
        waterfall->render(pass, linePipeline_, lineListPipeline_, width_, height_);
    }

    wgpuRenderPassEncoderEnd(pass);

    WGPUCommandBufferDescriptor cmdBufferDesc = {};
    cmdBufferDesc.nextInChain = nullptr;
    cmdBufferDesc.label = {nullptr, WGPU_STRLEN};

    WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, &cmdBufferDesc);

    WGPUQueue queue = wgpuDeviceGetQueue(device_);
    wgpuQueueSubmit(queue, 1, &commands);
    wgpuSurfacePresent(surface_);

    // Cleanup
    wgpuCommandBufferRelease(commands);
    wgpuCommandEncoderRelease(encoder);
    wgpuRenderPassEncoderRelease(pass);
    wgpuTextureViewRelease(textureView);
}

void Renderer::resize(int width, int height) {
    if (width > 0 && height > 0) {
        width_ = width;
        height_ = height;
        configureSurface();
    }
}

void Renderer::toggleGrid() {
    for (auto& waterfall : waterfalls_) {
        waterfall->toggleGrid();
    }
}