#include "waterfall.h"
#include <cmath>
#include <algorithm>

Waterfall::Waterfall(WGPUDevice device, float x, float y, float width, float height, const std::string& title)
    : device_(device), vertexBuffer_(nullptr), x_(x), y_(y), width_(width), height_(height),
      title_(title), showGrid_(true) {}

Waterfall::~Waterfall() {
    if (vertexBuffer_) wgpuBufferRelease(vertexBuffer_);
}

bool Waterfall::initialize() {
    // Create vertex buffer (large enough for grid, border, and waterfall lines)
    WGPUBufferDescriptor bufferDesc = {};
    bufferDesc.nextInChain = nullptr;
    bufferDesc.label = {nullptr, WGPU_STRLEN};
    bufferDesc.size = sizeof(Vertex) * POINTS_PER_LINE * 100;
    bufferDesc.usage = WGPUBufferUsage_Vertex | WGPUBufferUsage_CopyDst;
    bufferDesc.mappedAtCreation = false;

    vertexBuffer_ = wgpuDeviceCreateBuffer(device_, &bufferDesc);
    return vertexBuffer_ != nullptr;
}

void Waterfall::update(float time) {
    // Add new line every 10 frames
    if (static_cast<int>(time * 60.0f) % 10 == 0 && lines_.size() < MAX_LINES) {
        std::vector<Vertex> line;
        float phase = time;
        float freq = 2.0f + std::sin(time * 0.5f) * 1.0f;

        for (int i = 0; i < POINTS_PER_LINE; ++i) {
            float x = (static_cast<float>(i) / POINTS_PER_LINE) * 2.0f - 1.0f;
            float y = std::sin(static_cast<float>(i) * 0.1f * freq + phase) * 0.3f;

            float hue = std::fmod(time * 0.1f, 1.0f);
            Vertex v{
                {x, y},
                {
                    std::abs(std::sin(hue * 6.0f)),
                    std::abs(std::sin((hue + 0.33f) * 6.0f)),
                    std::abs(std::sin((hue + 0.66f) * 6.0f))
                }
            };
            line.push_back(v);
        }
        lines_.push_back(line);
    }

    // Scroll lines down
    for (auto& line : lines_) {
        for (auto& vertex : line) {
            vertex.position[1] -= 0.01f;
        }
    }

    // Remove lines that scrolled off screen
    lines_.erase(
        std::remove_if(lines_.begin(), lines_.end(),
            [](const std::vector<Vertex>& line) {
                return !line.empty() && line[0].position[1] < -1.1f;
            }),
        lines_.end()
    );
}

void Waterfall::render(WGPURenderPassEncoder pass,
                       WGPURenderPipeline linePipeline,
                       WGPURenderPipeline lineListPipeline,
                       int windowWidth, int windowHeight) {
    // Set viewport
    wgpuRenderPassEncoderSetViewport(pass,
        x_ * windowWidth,
        y_ * windowHeight,
        width_ * windowWidth,
        height_ * windowHeight,
        0.0f, 1.0f
    );

    // Collect all vertices
    std::vector<Vertex> allVertices;

    // Border
    auto borderVertices = generateBorder();
    size_t borderOffset = allVertices.size();
    allVertices.insert(allVertices.end(), borderVertices.begin(), borderVertices.end());

    // Grid
    size_t gridOffset = allVertices.size();
    size_t gridCount = 0;
    if (showGrid_) {
        auto gridVertices = generateGridLines();
        gridCount = gridVertices.size();
        allVertices.insert(allVertices.end(), gridVertices.begin(), gridVertices.end());
    }

    // Waterfall lines
    size_t linesOffset = allVertices.size();
    for (const auto& line : lines_) {
        allVertices.insert(allVertices.end(), line.begin(), line.end());
    }

    // Upload vertices
    if (!allVertices.empty()) {
        WGPUQueue queue = wgpuDeviceGetQueue(device_);
        wgpuQueueWriteBuffer(queue, vertexBuffer_, 0, allVertices.data(),
                            allVertices.size() * sizeof(Vertex));
    }

    // Draw border
    wgpuRenderPassEncoderSetPipeline(pass, lineListPipeline);
    wgpuRenderPassEncoderSetVertexBuffer(pass, 0, vertexBuffer_, 0, allVertices.size() * sizeof(Vertex));
    wgpuRenderPassEncoderDraw(pass, borderVertices.size(), 1, borderOffset, 0);

    // Draw grid
    if (showGrid_ && gridCount > 0) {
        wgpuRenderPassEncoderSetPipeline(pass, lineListPipeline);
        wgpuRenderPassEncoderDraw(pass, gridCount, 1, gridOffset, 0);
    }

    // Draw waterfall lines
    if (!lines_.empty()) {
        wgpuRenderPassEncoderSetPipeline(pass, linePipeline);
        for (size_t i = 0; i < lines_.size(); ++i) {
            uint32_t start = linesOffset + i * POINTS_PER_LINE;
            wgpuRenderPassEncoderDraw(pass, POINTS_PER_LINE, 1, start, 0);
        }
    }
}

void Waterfall::toggleGrid() {
    showGrid_ = !showGrid_;
}

std::vector<Vertex> Waterfall::generateGridLines() {
    std::vector<Vertex> vertices;
    float gridColor[3] = {0.3f, 0.7f, 0.9f};

    // Vertical lines
    for (int i = 0; i <= 10; ++i) {
        float x = -1.0f + (static_cast<float>(i) / 10.0f) * 2.0f;
        vertices.push_back({{x, -1.0f}, {gridColor[0], gridColor[1], gridColor[2]}});
        vertices.push_back({{x, 1.0f}, {gridColor[0], gridColor[1], gridColor[2]}});
    }

    // Horizontal lines
    for (int i = 0; i <= 10; ++i) {
        float y = -1.0f + (static_cast<float>(i) / 10.0f) * 2.0f;
        vertices.push_back({{-1.0f, y}, {gridColor[0], gridColor[1], gridColor[2]}});
        vertices.push_back({{1.0f, y}, {gridColor[0], gridColor[1], gridColor[2]}});
    }

    return vertices;
}

std::vector<Vertex> Waterfall::generateBorder() {
    float borderColor[3] = {0.6f, 0.7f, 0.7f};

    return {
        {{-1.0f, 1.0f}, {borderColor[0], borderColor[1], borderColor[2]}},
        {{1.0f, 1.0f}, {borderColor[0], borderColor[1], borderColor[2]}},
        {{1.0f, 1.0f}, {borderColor[0], borderColor[1], borderColor[2]}},
        {{1.0f, -1.0f}, {borderColor[0], borderColor[1], borderColor[2]}},
        {{1.0f, -1.0f}, {borderColor[0], borderColor[1], borderColor[2]}},
        {{-1.0f, -1.0f}, {borderColor[0], borderColor[1], borderColor[2]}},
        {{-1.0f, -1.0f}, {borderColor[0], borderColor[1], borderColor[2]}},
        {{-1.0f, 1.0f}, {borderColor[0], borderColor[1], borderColor[2]}}
    };
}