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

Renderer::Renderer(wgpu::Device device, wgpu::Surface surface, int width, int height)
    : device_(device), surface_(surface), width_(width), height_(height), time_(0.0f),
      surfaceFormat_(wgpu::TextureFormat::BGRA8Unorm) {}

Renderer::~Renderer() = default;

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() {
    wgpu::SurfaceConfiguration config{};
    config.device = device_;
    config.format = surfaceFormat_;
    config.usage = wgpu::TextureUsage::RenderAttachment;
    config.width = width_;
    config.height = height_;
    config.presentMode = wgpu::PresentMode::Fifo;
    config.alphaMode = wgpu::CompositeAlphaMode::Auto;

    surface_.configure(config);
}

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

    wgpu::ShaderSourceWGSL wgslSource{};
    wgslSource.chain.sType = wgpu::SType::ShaderSourceWGSL;
    wgslSource.code.data = shaderCode.c_str();
    wgslSource.code.length = shaderCode.length();

    wgpu::ShaderModuleDescriptor shaderDesc{};
    shaderDesc.nextInChain = &wgslSource.chain;
    wgpu::ShaderModule shader = device_.createShaderModule(shaderDesc);

    // Vertex buffer layout
    wgpu::VertexAttribute attributes[2];
    attributes[0].format = wgpu::VertexFormat::Float32x2;
    attributes[0].offset = 0;
    attributes[0].shaderLocation = 0;

    attributes[1].format = wgpu::VertexFormat::Float32x3;
    attributes[1].offset = 2 * sizeof(float);
    attributes[1].shaderLocation = 1;

    wgpu::VertexBufferLayout vertexBufferLayout{};
    vertexBufferLayout.arrayStride = 5 * sizeof(float);
    vertexBufferLayout.stepMode = wgpu::VertexStepMode::Vertex;
    vertexBufferLayout.attributeCount = 2;
    vertexBufferLayout.attributes = attributes;

    // Pipeline layout
    wgpu::PipelineLayoutDescriptor layoutDesc{};
    wgpu::PipelineLayout pipelineLayout = device_.createPipelineLayout(layoutDesc);

    // Color target
    wgpu::BlendState blend{};
    blend.color.operation = wgpu::BlendOperation::Add;
    blend.color.srcFactor = wgpu::BlendFactor::One;
    blend.color.dstFactor = wgpu::BlendFactor::Zero;
    blend.alpha.operation = wgpu::BlendOperation::Add;
    blend.alpha.srcFactor = wgpu::BlendFactor::One;
    blend.alpha.dstFactor = wgpu::BlendFactor::Zero;

    wgpu::ColorTargetState colorTarget{};
    colorTarget.format = surfaceFormat_;
    colorTarget.blend = &blend;
    colorTarget.writeMask = wgpu::ColorWriteMask::All;

    wgpu::FragmentState fragmentState{};
    fragmentState.module = shader;
    fragmentState.entryPoint.data = "fs_main";
    fragmentState.entryPoint.length = 7;
    fragmentState.targetCount = 1;
    fragmentState.targets = &colorTarget;

    // Line strip pipeline
    wgpu::RenderPipelineDescriptor pipelineDesc{};
    pipelineDesc.layout = pipelineLayout;
    pipelineDesc.vertex.module = shader;
    pipelineDesc.vertex.entryPoint.data = "vs_main";
    pipelineDesc.vertex.entryPoint.length = 7;
    pipelineDesc.vertex.bufferCount = 1;
    pipelineDesc.vertex.buffers = &vertexBufferLayout;
    pipelineDesc.primitive.topology = wgpu::PrimitiveTopology::LineStrip;
    pipelineDesc.fragment = &fragmentState;
    pipelineDesc.multisample.count = 1;
    pipelineDesc.multisample.mask = ~0u;

    linePipeline_ = device_.createRenderPipeline(pipelineDesc);

    // Line list pipeline for grid
    pipelineDesc.primitive.topology = wgpu::PrimitiveTopology::LineList;
    lineListPipeline_ = device_.createRenderPipeline(pipelineDesc);
}

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

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

void Renderer::render() {
    wgpu::SurfaceTexture surfaceTexture;
    surface_.getCurrentTexture(&surfaceTexture);

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

    wgpu::Texture texture(surfaceTexture.texture);
    wgpu::TextureView textureView = texture.createView();

    wgpu::CommandEncoder encoder = device_.createCommandEncoder();

    wgpu::RenderPassColorAttachment colorAttachment{};
    colorAttachment.view = textureView;
    colorAttachment.loadOp = wgpu::LoadOp::Clear;
    colorAttachment.storeOp = wgpu::StoreOp::Store;
    colorAttachment.clearValue = {0.1, 0.1, 0.15, 1.0};

    wgpu::RenderPassDescriptor renderPassDesc{};
    renderPassDesc.colorAttachmentCount = 1;
    renderPassDesc.colorAttachments = &colorAttachment;

    wgpu::RenderPassEncoder pass = encoder.beginRenderPass(renderPassDesc);

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

    pass.end();

    wgpu::CommandBuffer commands = encoder.finish();
    device_.getQueue().submit(1, &commands);
    surface_.present();
}

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();
    }
}