#include "viz/candlestick.h"
#include "imgui.h"
#include "implot.h"

void candlestick(const char* title, const double* dates, const double* opens,
                 const double* closes, const double* lows, const double* highs,
                 int count, float width_percent, bool tooltip) {
    if (count <= 0) return;

    // Compute half-width of each candle body in data coordinates.
    // Use spacing between consecutive dates when count > 1, else fallback to 0.5.
    double spacing = (count > 1) ? (dates[1] - dates[0]) : 1.0;
    double half_w  = spacing * (double)width_percent * 0.5;

    ImPlot::SetupAxes("Date", "Price", ImPlotAxisFlags_None, ImPlotAxisFlags_AutoFit);
    ImPlot::SetupAxisScale(ImAxis_X1, ImPlotScale_Time);

    // Auto-fit X axis to the data range with a small margin.
    ImPlot::SetupAxisLimits(ImAxis_X1,
                            dates[0]         - spacing,
                            dates[count - 1] + spacing,
                            ImGuiCond_Always);

    if (ImPlot::BeginPlot(title, ImVec2(-1, 0))) {
        ImDrawList* draw = ImPlot::GetPlotDrawList();

        const ImU32 col_bull = IM_COL32(0,  200, 80,  255);  // green — close >= open
        const ImU32 col_bear = IM_COL32(220, 50, 50,  255);  // red   — close < open

        int hovered_idx = -1;

        for (int i = 0; i < count; i++) {
            double x     = dates[i];
            double open  = opens[i];
            double close = closes[i];
            double low   = lows[i];
            double high  = highs[i];

            bool bullish = (close >= open);
            ImU32 col   = bullish ? col_bull : col_bear;

            // Convert data coordinates to screen pixels.
            ImVec2 body_tl = ImPlot::PlotToPixels(x - half_w, bullish ? close : open);
            ImVec2 body_br = ImPlot::PlotToPixels(x + half_w, bullish ? open  : close);
            ImVec2 wick_hi = ImPlot::PlotToPixels(x, high);
            ImVec2 wick_lo = ImPlot::PlotToPixels(x, low);
            float  cx      = (body_tl.x + body_br.x) * 0.5f;

            // Wick (high-low vertical line).
            draw->AddLine(ImVec2(cx, wick_hi.y), ImVec2(cx, wick_lo.y), col, 1.5f);

            // Body rectangle (open-close).
            // Ensure at least 1px height so flat candles are visible.
            if (body_br.y <= body_tl.y + 1.0f) body_br.y = body_tl.y + 1.0f;
            draw->AddRectFilled(body_tl, body_br, col);
            draw->AddRect(body_tl, body_br, col);

            // Track hovered candle for tooltip.
            if (tooltip && ImPlot::IsPlotHovered()) {
                ImVec2 mouse = ImGui::GetMousePos();
                if (mouse.x >= body_tl.x - 4 && mouse.x <= body_br.x + 4 &&
                    mouse.y >= wick_hi.y  - 4 && mouse.y <= wick_lo.y  + 4) {
                    hovered_idx = i;
                }
            }
        }

        // Tooltip for the hovered candle.
        if (tooltip && hovered_idx >= 0) {
            int i = hovered_idx;
            ImGui::BeginTooltip();
            ImGui::Text("O: %.4f", opens[i]);
            ImGui::Text("H: %.4f", highs[i]);
            ImGui::Text("L: %.4f", lows[i]);
            ImGui::Text("C: %.4f", closes[i]);
            ImGui::EndTooltip();
        }

        ImPlot::EndPlot();
    }
}