#include <algorithm>
#include <cstring>
#include <string>
#include <vector>
#include <mxl/flow.h>
#include <mxl/time.h>
#include "NodeBase.hpp"
#include "FlowDef.hpp"
#include "V210.hpp"

// Round down to nearest V210-aligned pixel count (multiple of 6).
static int v210_align(int pixels) { return (pixels / 6) * 6; }

class PiPNode : public dmf::NodeBase {
    void run() override {
        if (!config().contains("background_flow_id")) { log("no background connected"); return; }
        if (!config().contains("inset_flow_id"))      { log("no inset connected");      return; }
        if (!config().contains("video_flow_id"))      { log("no output connected");      return; }

        const auto bg_id    = config().at("background_flow_id").at("id").get<std::string>();
        const auto inset_id = config().at("inset_flow_id").at("id").get<std::string>();
        const auto out_id   = config().at("video_flow_id").at("id").get<std::string>();

        // Position and size of the inset in the output frame.
        // x and width are snapped to 6-pixel V210 boundaries.
        const int pip_x = v210_align(config().value("x", 0));
        const int pip_y = config().value("y", 0);
        const int pip_w = v210_align(config().value("width",  480));
        const int pip_h = config().value("height", 270);

        // --- Wait for both input flows ---
        for (const auto* fid : {&bg_id, &inset_id}) {
            log("waiting for flow %s...", fid->c_str());
            bool active = false;
            while (!active && dmf::g_running.load(std::memory_order_relaxed)) {
                mxlIsFlowActive(instance(), fid->c_str(), &active);
                if (!active) mxlSleepForNs(100'000'000);
            }
            if (!dmf::g_running) return;
        }

        // --- Create readers ---
        mxlFlowReader     bg_reader{}, inset_reader{};
        mxlFlowConfigInfo bg_cfg{}, inset_cfg{};

        if (mxlCreateFlowReader(instance(), bg_id.c_str(), "", &bg_reader) != MXL_STATUS_OK) {
            log("background mxlCreateFlowReader failed"); return;
        }
        if (mxlCreateFlowReader(instance(), inset_id.c_str(), "", &inset_reader) != MXL_STATUS_OK) {
            log("inset mxlCreateFlowReader failed");
            mxlReleaseFlowReader(instance(), bg_reader);
            return;
        }
        mxlFlowReaderGetConfigInfo(bg_reader,    &bg_cfg);
        mxlFlowReaderGetConfigInfo(inset_reader, &inset_cfg);

        const uint32_t bg_stride    = bg_cfg.discrete.sliceSizes[0];
        const uint32_t inset_stride = inset_cfg.discrete.sliceSizes[0];

        // --- Read formats from flow_def.json ---
        const auto bg_fi    = dmf::read_video_flow_info(domain(), bg_id);
        const auto inset_fi = dmf::read_video_flow_info(domain(), inset_id);

        const int bg_w    = bg_fi.width;
        const int bg_h    = bg_fi.height;
        const int fps_num = bg_fi.fps_num;
        const int fps_den = bg_fi.fps_den;
        const int inset_w = inset_fi.width;
        const int inset_h = inset_fi.height;

        log("background: %dx%d @ %d/%d fps  stride=%u",
            bg_w, bg_h, fps_num, fps_den, bg_stride);
        log("inset src:  %dx%d  stride=%u", inset_w, inset_h, inset_stride);
        log("pip region: %dx%d at (%d,%d)", pip_w, pip_h, pip_x, pip_y);

        // Clamp pip region to background bounds
        const int clamped_w = v210_align(std::min(pip_w, bg_w - pip_x));
        const int clamped_h = std::min(pip_h, bg_h - pip_y);
        if (clamped_w <= 0 || clamped_h <= 0) {
            log("pip region is outside background bounds — exiting");
            mxlReleaseFlowReader(instance(), bg_reader);
            mxlReleaseFlowReader(instance(), inset_reader);
            return;
        }

        // --- Create output writer (same format as background) ---
        mxlFlowWriter     out_writer{};
        mxlFlowConfigInfo out_cfg{};
        bool created = false;
        mxlStatus vst = mxlCreateFlowWriter(
            instance(),
            dmf::make_video_flow_def(out_id, node_id(), bg_w, bg_h, fps_num, fps_den).c_str(),
            "", &out_writer, &out_cfg, &created);
        if (vst != MXL_STATUS_OK) {
            log("mxlCreateFlowWriter failed (%s)", dmf::mxl_status_str(vst));
            mxlReleaseFlowReader(instance(), bg_reader);
            mxlReleaseFlowReader(instance(), inset_reader);
            return;
        }
        const uint32_t out_stride = out_cfg.discrete.sliceSizes[0];
        log("output: stride=%u  grain=%u B  ring=%u grains",
            out_stride, out_stride * static_cast<uint32_t>(bg_h), out_cfg.discrete.grainCount);

        // --- Pre-allocate bilinear scaling workspace (reused every frame) ---
        std::vector<uint16_t> Y0(inset_w), Y1(inset_w);
        std::vector<uint16_t> Cb0(inset_w / 2), Cb1(inset_w / 2);
        std::vector<uint16_t> Cr0(inset_w / 2), Cr1(inset_w / 2);

        // --- Clock: driven by background (inset follows best-effort) ---
        const mxlRational rate = {fps_num, fps_den};
        uint64_t index = mxlGetCurrentIndex(&rate);
        log("start index=%llu", index);

        uint64_t frame_count = 0, stall_count = 0;
        bool fatal = false;

        while (dmf::g_running.load(std::memory_order_relaxed)) {
            // Background is the master clock.
            mxlGrainInfo bg_grain{};
            uint8_t* bg_buf = nullptr;
            const mxlStatus bg_st = mxlFlowReaderGetGrain(
                bg_reader, index, 80'000'000, &bg_grain, &bg_buf);

            if (bg_st == MXL_STATUS_OK && bg_buf) {
                // Try inset at same index (short timeout — hardware often lands one frame behind).
                // On any miss, fall back to headIndex: the ring buffer is the cache.
                mxlGrainInfo inset_grain{};
                uint8_t* inset_buf = nullptr;
                mxlStatus in_st = mxlFlowReaderGetGrain(
                    inset_reader, index, 8'000'000, &inset_grain, &inset_buf);

                if (in_st != MXL_STATUS_OK) {
                    mxlFlowRuntimeInfo ri{};
                    mxlFlowReaderGetRuntimeInfo(inset_reader, &ri);
                    in_st = mxlFlowReaderGetGrain(
                        inset_reader, ri.headIndex, 8'000'000, &inset_grain, &inset_buf);
                    stall_count++;
                }

                mxlGrainInfo out_grain{};
                uint8_t* out_buf = nullptr;
                const mxlStatus wst = mxlFlowWriterOpenGrain(
                    out_writer, index, &out_grain, &out_buf);
                if (wst == MXL_STATUS_OK) {
                    std::memcpy(out_buf, bg_buf,
                                static_cast<size_t>(bg_stride) * static_cast<size_t>(bg_h));
                    if (in_st == MXL_STATUS_OK && inset_buf) {
                        dmf::v210::scale_and_overlay(
                            inset_buf,  inset_stride, inset_w, inset_h,
                            out_buf,    out_stride,
                            pip_x, pip_y, clamped_w, clamped_h,
                            Y0, Y1, Cb0, Cb1, Cr0, Cr1);
                    }
                    out_grain.flags = bg_grain.flags & MXL_GRAIN_FLAG_INVALID;
                    out_grain.validSlices = out_grain.totalSlices;
                    mxlFlowWriterCommitGrain(out_writer, &out_grain);
                    frame_count++;
                    if (frame_count % 25 == 0)
                        log("heartbeat frames=%llu  stalls=%llu  index=%llu",
                            frame_count, stall_count, index);
                } else {
                    log("writer OpenGrain failed (%s) at index=%llu",
                        dmf::mxl_status_str(wst), index);
                }
                index++;

            } else if (bg_st == MXL_ERR_OUT_OF_RANGE_TOO_EARLY) {
                // Background stalled longer than 80 ms — skip this index.
                stall_count++;
                log("bg stall (TOO_EARLY) at index=%llu  frames=%llu", index, frame_count);
                index++;

            } else if (bg_st == MXL_ERR_OUT_OF_RANGE_TOO_LATE) {
                stall_count++;
                mxlFlowRuntimeInfo ri{};
                mxlFlowReaderGetRuntimeInfo(bg_reader, &ri);
                log("bg TOO_LATE at index=%llu → jumping to %llu  frames=%llu",
                    index, ri.headIndex, frame_count);
                index = ri.headIndex;

            } else {
                log("bg fatal (%s) at index=%llu", dmf::mxl_status_str(bg_st), index);
                fatal = true;
                break;
            }
        }

        log("stopped: %s  frames=%llu  stalls=%llu  index=%llu",
            fatal ? "fatal error" : "shutdown signal",
            frame_count, stall_count, index);

        mxlReleaseFlowReader(instance(), bg_reader);
        mxlReleaseFlowReader(instance(), inset_reader);
        mxlReleaseFlowWriter(instance(), out_writer);
    }
};

int main() {
    PiPNode node;
    return node.execute();
}