#include <atomic>
#include <chrono>
#include <cstring>
#include <iostream>
#include <memory>
#include <thread>
#include <vector>

#include "node.h"

#if defined(RK3588_ENABLE_RGA)
#include "im2d.hpp"
#include "im2d_buffer.h"
#include "im2d_type.h"
#endif

#if defined(RK3588_ENABLE_FFMPEG)
extern "C" {
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
}
#endif

namespace rk3588 {

namespace {

inline int Align16(int v) { return (v + 15) & ~15; }

#if defined(RK3588_ENABLE_RGA)
int ToRgaFormat(PixelFormat fmt) {
    switch (fmt) {
        case PixelFormat::NV12: return RK_FORMAT_YCbCr_420_SP;
        case PixelFormat::YUV420: return RK_FORMAT_YCbCr_420_P;
        case PixelFormat::RGB: return RK_FORMAT_RGB_888;
        case PixelFormat::BGR: return RK_FORMAT_BGR_888;
        default: return RK_FORMAT_UNKNOWN;
    }
}
#endif

PixelFormat ParseFormat(const std::string& s) {
    if (s == "nv12" || s == "NV12") return PixelFormat::NV12;
    if (s == "yuv420" || s == "YUV420") return PixelFormat::YUV420;
    if (s == "rgb" || s == "RGB") return PixelFormat::RGB;
    if (s == "bgr" || s == "BGR") return PixelFormat::BGR;
    return PixelFormat::UNKNOWN;
}

size_t CalcImageSize(int w, int h, PixelFormat fmt) {
    switch (fmt) {
        case PixelFormat::NV12:
        case PixelFormat::YUV420:
            return static_cast<size_t>(w) * h * 3 / 2;
        case PixelFormat::RGB:
        case PixelFormat::BGR:
            return static_cast<size_t>(w) * h * 3;
        default:
            return 0;
    }
}

}  // namespace

class PreprocessNode : public INode {
public:
    std::string Id() const override { return id_; }
    std::string Type() const override { return "preprocess"; }

    bool Init(const SimpleJson& config, const NodeContext& ctx) override {
        id_ = config.ValueOr<std::string>("id", "preprocess");
        dst_w_ = config.ValueOr<int>("dst_w", 640);
        dst_h_ = config.ValueOr<int>("dst_h", 640);
        keep_ratio_ = config.ValueOr<bool>("keep_ratio", false);
        std::string fmt_str = config.ValueOr<std::string>("dst_format", "");
        if (!fmt_str.empty()) {
            dst_fmt_ = ParseFormat(fmt_str);
        }
        use_rga_ = config.ValueOr<bool>("use_rga", true);

        input_queue_ = ctx.input_queue;
        if (!input_queue_) {
            std::cerr << "[preprocess] no input queue for node " << id_ << "\n";
            return false;
        }
        if (ctx.output_queues.empty()) {
            std::cerr << "[preprocess] no output queue for node " << id_ << "\n";
            return false;
        }
        output_queues_ = ctx.output_queues;

#if !defined(RK3588_ENABLE_RGA)
        use_rga_ = false;
#endif
#if !defined(RK3588_ENABLE_FFMPEG)
        if (!use_rga_) {
            std::cerr << "[preprocess] neither RGA nor FFmpeg enabled\n";
            return false;
        }
#endif
        return true;
    }

    bool Start() override {
        if (!input_queue_) return false;
        running_.store(true);

#if defined(RK3588_ENABLE_RGA)
        if (use_rga_) {
            worker_ = std::thread(&PreprocessNode::LoopRga, this);
        } else {
            worker_ = std::thread(&PreprocessNode::LoopSwscale, this);
        }
#elif defined(RK3588_ENABLE_FFMPEG)
        worker_ = std::thread(&PreprocessNode::LoopSwscale, this);
#else
        worker_ = std::thread(&PreprocessNode::LoopPassthrough, this);
#endif
        std::cout << "[preprocess] start dst=" << dst_w_ << "x" << dst_h_
                  << (use_rga_ ? " (rga)" : " (swscale)") << "\n";
        return true;
    }

    void Stop() override {
        running_.store(false);
        if (input_queue_) input_queue_->Stop();
        for (auto& q : output_queues_) q->Stop();
        if (worker_.joinable()) worker_.join();
    }

private:
    void PushToDownstream(FramePtr frame) {
        for (auto& q : output_queues_) {
            q->Push(frame);
        }
    }

    void LoopPassthrough() {
        using namespace std::chrono;
        FramePtr frame;
        while (running_.load()) {
            if (!input_queue_->Pop(frame, milliseconds(200))) continue;
            if (!frame) continue;
            PushToDownstream(frame);
            ++processed_;
            if (processed_ % 100 == 0) {
                std::cout << "[preprocess] passthrough frame " << frame->frame_id << "\n";
            }
        }
    }

#if defined(RK3588_ENABLE_RGA)
    void LoopRga() {
        using namespace std::chrono;
        FramePtr frame;
        while (running_.load()) {
            if (!input_queue_->Pop(frame, milliseconds(200))) continue;
            if (!frame) continue;

            PixelFormat out_fmt = (dst_fmt_ != PixelFormat::UNKNOWN) ? dst_fmt_ : frame->format;
            int out_w = dst_w_;
            int out_h = dst_h_;

            if (keep_ratio_ && frame->width > 0 && frame->height > 0) {
                float scale = std::min(static_cast<float>(dst_w_) / frame->width,
                                       static_cast<float>(dst_h_) / frame->height);
                out_w = static_cast<int>(frame->width * scale);
                out_h = static_cast<int>(frame->height * scale);
                out_w = (out_w + 1) & ~1;
                out_h = (out_h + 1) & ~1;
            }

            size_t out_size = CalcImageSize(out_w, out_h, out_fmt);
            if (out_size == 0) {
                PushToDownstream(frame);
                continue;
            }

            auto buffer = std::make_shared<std::vector<uint8_t>>(out_size);
            int src_fmt_rga = ToRgaFormat(frame->format);
            int dst_fmt_rga = ToRgaFormat(out_fmt);

            if (src_fmt_rga == RK_FORMAT_UNKNOWN || dst_fmt_rga == RK_FORMAT_UNKNOWN) {
                std::cerr << "[preprocess] unsupported format for RGA\n";
                PushToDownstream(frame);
                continue;
            }

            int src_stride = frame->planes[0].stride > 0 ? frame->planes[0].stride
                             : (frame->stride > 0 ? frame->stride : frame->width);

            rga_buffer_t src_buf{};
            rga_buffer_t dst_buf{};

            if (frame->dma_fd >= 0) {
                src_buf = wrapbuffer_fd_t(frame->dma_fd, frame->width, frame->height,
                                          src_stride, Align16(frame->height), src_fmt_rga);
            } else if (frame->data) {
                src_buf = wrapbuffer_virtualaddr_t(frame->data, frame->width, frame->height,
                                                   src_stride, Align16(frame->height), src_fmt_rga);
            } else {
                PushToDownstream(frame);
                continue;
            }

            dst_buf = wrapbuffer_virtualaddr_t(buffer->data(), out_w, out_h,
                                               out_w, out_h, dst_fmt_rga);

            IM_STATUS status = IM_STATUS_SUCCESS;
            bool need_cvt = (src_fmt_rga != dst_fmt_rga);
            bool need_resize = (frame->width != out_w || frame->height != out_h);

            if (need_resize && need_cvt) {
                // Resize + color convert in one call using improcess
                status = imresize(src_buf, dst_buf);
                if (status == IM_STATUS_SUCCESS) {
                    // imresize doesn't do color conversion; need separate call
                    // For simplicity, do resize first then cvtcolor
                    rga_buffer_t tmp = dst_buf;
                    status = imcvtcolor(tmp, dst_buf, src_fmt_rga, dst_fmt_rga, IM_COLOR_SPACE_DEFAULT);
                }
            } else if (need_resize) {
                status = imresize(src_buf, dst_buf);
            } else if (need_cvt) {
                status = imcvtcolor(src_buf, dst_buf, src_fmt_rga, dst_fmt_rga, IM_COLOR_SPACE_DEFAULT);
            } else {
                status = imcopy(src_buf, dst_buf);
            }

            if (status != IM_STATUS_SUCCESS) {
                std::cerr << "[preprocess] RGA failed: " << imStrError(status) << "\n";
                PushToDownstream(frame);
                continue;
            }

            auto out_frame = std::make_shared<Frame>();
            out_frame->width = out_w;
            out_frame->height = out_h;
            out_frame->format = out_fmt;
            out_frame->stride = out_w;
            out_frame->data = buffer->data();
            out_frame->data_size = buffer->size();
            out_frame->data_owner = buffer;
            out_frame->pts = frame->pts;
            out_frame->frame_id = frame->frame_id;
            out_frame->det = frame->det;
            out_frame->user_meta = frame->user_meta;

            SetupPlanes(*out_frame, out_fmt);
            PushToDownstream(out_frame);
            ++processed_;

            if (processed_ % 100 == 0) {
                std::cout << "[preprocess] rga frame " << out_frame->frame_id
                          << " " << frame->width << "x" << frame->height
                          << " -> " << out_w << "x" << out_h << "\n";
            }
        }
    }
#endif

#if defined(RK3588_ENABLE_FFMPEG)
    void LoopSwscale() {
        using namespace std::chrono;
        FramePtr frame;
        SwsContext* sws_ctx = nullptr;
        int last_src_w = 0, last_src_h = 0;
        AVPixelFormat last_src_fmt = AV_PIX_FMT_NONE;
        AVPixelFormat last_dst_fmt = AV_PIX_FMT_NONE;

        while (running_.load()) {
            if (!input_queue_->Pop(frame, milliseconds(200))) continue;
            if (!frame) continue;

            PixelFormat out_fmt = (dst_fmt_ != PixelFormat::UNKNOWN) ? dst_fmt_ : frame->format;
            int out_w = dst_w_;
            int out_h = dst_h_;

            if (keep_ratio_ && frame->width > 0 && frame->height > 0) {
                float scale = std::min(static_cast<float>(dst_w_) / frame->width,
                                       static_cast<float>(dst_h_) / frame->height);
                out_w = static_cast<int>(frame->width * scale);
                out_h = static_cast<int>(frame->height * scale);
                out_w = (out_w + 1) & ~1;
                out_h = (out_h + 1) & ~1;
            }

            AVPixelFormat src_av_fmt = ToAvFormat(frame->format);
            AVPixelFormat dst_av_fmt = ToAvFormat(out_fmt);

            if (src_av_fmt == AV_PIX_FMT_NONE || dst_av_fmt == AV_PIX_FMT_NONE) {
                PushToDownstream(frame);
                continue;
            }

            if (!sws_ctx || frame->width != last_src_w || frame->height != last_src_h ||
                src_av_fmt != last_src_fmt || dst_av_fmt != last_dst_fmt) {
                if (sws_ctx) sws_freeContext(sws_ctx);
                sws_ctx = sws_getContext(frame->width, frame->height, src_av_fmt,
                                         out_w, out_h, dst_av_fmt,
                                         SWS_BILINEAR, nullptr, nullptr, nullptr);
                last_src_w = frame->width;
                last_src_h = frame->height;
                last_src_fmt = src_av_fmt;
                last_dst_fmt = dst_av_fmt;
            }

            if (!sws_ctx) {
                PushToDownstream(frame);
                continue;
            }

            size_t out_size = CalcImageSize(out_w, out_h, out_fmt);
            auto buffer = std::make_shared<std::vector<uint8_t>>(out_size);

            uint8_t* src_data[4] = {nullptr};
            int src_linesize[4] = {0};
            uint8_t* dst_data[4] = {nullptr};
            int dst_linesize[4] = {0};

            SetupAvPlanes(frame.get(), src_data, src_linesize);
            av_image_fill_arrays(dst_data, dst_linesize, buffer->data(),
                                 dst_av_fmt, out_w, out_h, 1);

            sws_scale(sws_ctx, src_data, src_linesize, 0, frame->height,
                      dst_data, dst_linesize);

            auto out_frame = std::make_shared<Frame>();
            out_frame->width = out_w;
            out_frame->height = out_h;
            out_frame->format = out_fmt;
            out_frame->stride = out_w;
            out_frame->data = buffer->data();
            out_frame->data_size = buffer->size();
            out_frame->data_owner = buffer;
            out_frame->pts = frame->pts;
            out_frame->frame_id = frame->frame_id;
            out_frame->det = frame->det;
            out_frame->user_meta = frame->user_meta;

            SetupPlanes(*out_frame, out_fmt);
            PushToDownstream(out_frame);
            ++processed_;

            if (processed_ % 100 == 0) {
                std::cout << "[preprocess] swscale frame " << out_frame->frame_id
                          << " " << frame->width << "x" << frame->height
                          << " -> " << out_w << "x" << out_h << "\n";
            }
        }

        if (sws_ctx) sws_freeContext(sws_ctx);
    }

    static AVPixelFormat ToAvFormat(PixelFormat fmt) {
        switch (fmt) {
            case PixelFormat::NV12: return AV_PIX_FMT_NV12;
            case PixelFormat::YUV420: return AV_PIX_FMT_YUV420P;
            case PixelFormat::RGB: return AV_PIX_FMT_RGB24;
            case PixelFormat::BGR: return AV_PIX_FMT_BGR24;
            default: return AV_PIX_FMT_NONE;
        }
    }

    static void SetupAvPlanes(const Frame* f, uint8_t* data[4], int linesize[4]) {
        if (!f->data) return;
        if (f->format == PixelFormat::NV12) {
            data[0] = f->planes[0].data ? f->planes[0].data : f->data;
            data[1] = f->planes[1].data ? f->planes[1].data : (f->data + f->width * f->height);
            linesize[0] = f->planes[0].stride > 0 ? f->planes[0].stride : f->width;
            linesize[1] = f->planes[1].stride > 0 ? f->planes[1].stride : f->width;
        } else if (f->format == PixelFormat::YUV420) {
            data[0] = f->planes[0].data ? f->planes[0].data : f->data;
            int y_size = f->width * f->height;
            int uv_size = y_size / 4;
            data[1] = f->planes[1].data ? f->planes[1].data : (f->data + y_size);
            data[2] = f->planes[2].data ? f->planes[2].data : (f->data + y_size + uv_size);
            linesize[0] = f->planes[0].stride > 0 ? f->planes[0].stride : f->width;
            linesize[1] = f->planes[1].stride > 0 ? f->planes[1].stride : f->width / 2;
            linesize[2] = f->planes[2].stride > 0 ? f->planes[2].stride : f->width / 2;
        } else {
            data[0] = f->data;
            linesize[0] = f->stride > 0 ? f->stride : f->width * 3;
        }
    }
#endif

    void SetupPlanes(Frame& f, PixelFormat fmt) {
        if (fmt == PixelFormat::NV12) {
            f.plane_count = 2;
            int y_size = f.width * f.height;
            f.planes[0] = {f.data, f.width, y_size, 0};
            f.planes[1] = {f.data + y_size, f.width, y_size / 2, y_size};
        } else if (fmt == PixelFormat::YUV420) {
            f.plane_count = 3;
            int y_size = f.width * f.height;
            int uv_size = y_size / 4;
            f.planes[0] = {f.data, f.width, y_size, 0};
            f.planes[1] = {f.data + y_size, f.width / 2, uv_size, y_size};
            f.planes[2] = {f.data + y_size + uv_size, f.width / 2, uv_size, y_size + uv_size};
        } else {
            f.plane_count = 1;
            f.planes[0] = {f.data, f.width * 3, static_cast<int>(f.data_size), 0};
        }
    }

    std::string id_;
    int dst_w_ = 640;
    int dst_h_ = 640;
    bool keep_ratio_ = false;
    PixelFormat dst_fmt_ = PixelFormat::UNKNOWN;
    bool use_rga_ = true;

    std::atomic<bool> running_{false};
    std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
    std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;
    std::thread worker_;
    uint64_t processed_ = 0;
};

REGISTER_NODE(PreprocessNode, "preprocess");

}  // namespace rk3588