#include <atomic>
#include <chrono>
#include <functional>
#include <iostream>
#include <thread>
#include <vector>

#include "node.h"
#include "utils/thread_affinity.h"

#if defined(RK3588_ENABLE_MPP)
extern "C" {
#include <rockchip/mpp_buffer.h>
#include <rockchip/mpp_frame.h>
#include <rockchip/mpp_packet.h>
#include <rockchip/rk_mpi.h>
}
#include <unistd.h>
#endif

#ifdef RK3588_ENABLE_FFMPEG
extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libavutil/opt.h>
}
#endif

namespace rk3588 {

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

    bool Init(const SimpleJson& config, const NodeContext& ctx) override {
        id_ = config.ValueOr<std::string>("id", "input_rtsp");
        url_ = config.ValueOr<std::string>("url", "");
        fps_ = config.ValueOr<int>("fps", 25);
        width_ = config.ValueOr<int>("width", 1920);
        height_ = config.ValueOr<int>("height", 1080);
        use_ffmpeg_ = config.ValueOr<bool>("use_ffmpeg", false);
        use_mpp_ = config.ValueOr<bool>("use_mpp", true);
        ffmpeg_force_tcp_ = config.ValueOr<bool>("force_tcp", true);
        reconnect_sec_ = config.ValueOr<int>("reconnect_sec", 5);
        reconnect_backoff_max_sec_ = config.ValueOr<int>("reconnect_backoff_max_sec", 30);
        fallback_to_stub_on_fail_ = config.ValueOr<bool>("fallback_to_stub_on_fail", false);
        cpu_affinity_ = ParseCpuAffinity(config);
        if (ctx.output_queues.empty()) {
            std::cerr << "[input_rtsp] no downstream queue configured for node " << id_ << "\n";
            return false;
        }
        out_queues_ = ctx.output_queues;
        return true;
    }

    bool Start() override {
        if (out_queues_.empty()) return false;
        running_.store(true);
#if defined(RK3588_ENABLE_FFMPEG) && defined(RK3588_ENABLE_MPP)
        if (use_mpp_) {
            worker_ = std::thread(&InputRtspNode::LoopFfmpegMpp, this);
        } else if (use_ffmpeg_) {
            worker_ = std::thread(&InputRtspNode::LoopFfmpegCpu, this);
        } else {
            worker_ = std::thread(&InputRtspNode::LoopStub, this);
        }
#elif defined(RK3588_ENABLE_FFMPEG)
        if (use_ffmpeg_) {
            worker_ = std::thread(&InputRtspNode::LoopFfmpegCpu, this);
        } else {
            worker_ = std::thread(&InputRtspNode::LoopStub, this);
        }
#else
        if (use_ffmpeg_ || use_mpp_) {
            std::cerr << "[input_rtsp] requested ffmpeg/mpp but not enabled at build time" << "\n";
        }
        worker_ = std::thread(&InputRtspNode::LoopStub, this);
#endif

        std::cout << "[input_rtsp] start url=" << url_ << " fps=" << fps_;
#if defined(RK3588_ENABLE_MPP)
        if (use_mpp_) std::cout << " (ffmpeg demux + mpp decode)";
        else if (use_ffmpeg_) std::cout << " (ffmpeg cpu decode" << (ffmpeg_force_tcp_ ? ", tcp" : ", udp") << ")";
        else std::cout << " (stub)";
#else
        if (use_ffmpeg_) std::cout << " (ffmpeg cpu decode" << (ffmpeg_force_tcp_ ? ", tcp" : ", udp") << ")";
        else std::cout << " (stub)";
#endif
        std::cout << "\n";
        return true;
    }

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

    void Drain() override {
        running_.store(false);
    }

private:
    void ApplyAffinity() {
        if (cpu_affinity_.empty()) return;
        std::string aerr;
        if (!SetCurrentThreadAffinity(cpu_affinity_, aerr)) {
            std::cerr << "[input_rtsp] SetCurrentThreadAffinity failed: " << aerr << "\n";
        }
    }

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

    void LoopStub() {
        ApplyAffinity();
        using namespace std::chrono;
        auto frame_interval = fps_ > 0 ? milliseconds(1000 / fps_) : milliseconds(40);
        while (running_.load()) {
            auto frame = std::make_shared<Frame>();
            frame->width = width_;
            frame->height = height_;
            frame->format = PixelFormat::NV12;
            frame->plane_count = 2;
            frame->planes[0] = {nullptr, width_, width_ * height_, 0};
            frame->planes[1] = {nullptr, width_, width_ * height_ / 2, width_ * height_};
            frame->frame_id = ++frame_id_;
            frame->pts = duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count();
            PushToDownstream(frame);

            if (frame_id_ % 100 == 0) {
                std::cout << "[input_rtsp] generated frame " << frame_id_ << " queue="
                          << (out_queues_.empty() ? 0 : out_queues_[0]->Size())
                          << " drops=" << (out_queues_.empty() ? 0 : out_queues_[0]->DroppedCount())
                          << "\n";
            }
            std::this_thread::sleep_for(frame_interval);
        }
    }
#if defined(RK3588_ENABLE_FFMPEG)
    void LoopFfmpegCpu() {
        ApplyAffinity();
        using namespace std::chrono;

        int backoff = std::max(1, reconnect_sec_);
        const int backoff_max = std::max(backoff, reconnect_backoff_max_sec_);

        while (running_.load()) {
            AVFormatContext* fmt_ctx = nullptr;
            AVCodecContext* codec_ctx = nullptr;
            AVPacket* pkt = av_packet_alloc();
            AVFrame* frm = av_frame_alloc();
            int video_stream = -1;
            AVRational time_base{1, 1000};

            AVDictionary* opts = nullptr;
            if (ffmpeg_force_tcp_) av_dict_set(&opts, "rtsp_transport", "tcp", 0);
            if (avformat_open_input(&fmt_ctx, url_.c_str(), nullptr, &opts) < 0) {
                std::cerr << "[input_rtsp] avformat_open_input failed: " << url_ << "\n";
                av_dict_free(&opts);
                Cleanup(fmt_ctx, codec_ctx, pkt, frm);
                if (fallback_to_stub_on_fail_) {
                    LoopStub();
                    return;
                }
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }
            av_dict_free(&opts);
            if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
                std::cerr << "[input_rtsp] avformat_find_stream_info failed\n";
                Cleanup(fmt_ctx, codec_ctx, pkt, frm);
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }
            for (unsigned i = 0; i < fmt_ctx->nb_streams; ++i) {
                if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
                    video_stream = static_cast<int>(i);
                    time_base = fmt_ctx->streams[i]->time_base;
                    break;
                }
            }
            if (video_stream < 0) {
                std::cerr << "[input_rtsp] no video stream\n";
                Cleanup(fmt_ctx, codec_ctx, pkt, frm);
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }

            const AVCodec* codec = avcodec_find_decoder(fmt_ctx->streams[video_stream]->codecpar->codec_id);
            if (!codec) {
                std::cerr << "[input_rtsp] decoder not found\n";
                Cleanup(fmt_ctx, codec_ctx, pkt, frm);
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }
            codec_ctx = avcodec_alloc_context3(codec);
            avcodec_parameters_to_context(codec_ctx, fmt_ctx->streams[video_stream]->codecpar);
            if (avcodec_open2(codec_ctx, codec, nullptr) < 0) {
                std::cerr << "[input_rtsp] avcodec_open2 failed\n";
                Cleanup(fmt_ctx, codec_ctx, pkt, frm);
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }

            backoff = std::max(1, reconnect_sec_); // reset after successful open
            int read_fail = 0;
            while (running_.load()) {
                if (av_read_frame(fmt_ctx, pkt) < 0) {
                    if (++read_fail >= 50) {
                        break;
                    }
                    std::this_thread::sleep_for(milliseconds(10));
                    continue;
                }
                read_fail = 0;
                if (pkt->stream_index != video_stream) {
                    av_packet_unref(pkt);
                    continue;
                }
                if (avcodec_send_packet(codec_ctx, pkt) == 0) {
                    while (avcodec_receive_frame(codec_ctx, frm) == 0) {
                        PushFrameFromAVFrame(*frm, time_base);
                        av_frame_unref(frm);
                    }
                }
                av_packet_unref(pkt);
            }

            Cleanup(fmt_ctx, codec_ctx, pkt, frm);
            if (!running_.load()) break;
            std::this_thread::sleep_for(seconds(std::max(1, reconnect_sec_)));
        }
    }

    void PushFrameFromAVFrame(const AVFrame& f, AVRational time_base) {
        PixelFormat fmt = PixelFormat::UNKNOWN;
        int plane_count = 0;
        int width = f.width;
        int height = f.height;

        if (f.format == AV_PIX_FMT_NV12) {
            fmt = PixelFormat::NV12;
            plane_count = 2;
        } else if (f.format == AV_PIX_FMT_YUV420P || f.format == AV_PIX_FMT_YUVJ420P) {
            fmt = PixelFormat::YUV420;
            plane_count = 3;
        }

        int buf_size = av_image_get_buffer_size(static_cast<AVPixelFormat>(f.format), width, height, 1);
        if (buf_size <= 0) return;
        auto buffer = std::make_shared<std::vector<uint8_t>>(buf_size);
        if (av_image_copy_to_buffer(buffer->data(), buffer->size(), f.data, f.linesize,
                                    static_cast<AVPixelFormat>(f.format), width, height, 1) < 0) {
            return;
        }

        auto frame = std::make_shared<Frame>();
        frame->width = width;
        frame->height = height;
        frame->format = fmt;
        frame->data = buffer->data();
        frame->data_size = buffer->size();
        frame->plane_count = plane_count;
        frame->frame_id = ++frame_id_;
        auto best_pts = f.pts; // av_frame_get_best_effort_timestamp is deprecated
        frame->pts = best_pts == AV_NOPTS_VALUE ? 0
                     : static_cast<uint64_t>(av_rescale_q(best_pts, time_base, {1, 1000}));
        frame->data_owner = buffer;

        if (fmt == PixelFormat::NV12 && plane_count == 2) {
            int y_size = width * height;
            frame->planes[0] = {buffer->data(), width, y_size, 0};
            frame->planes[1] = {buffer->data() + y_size, width, y_size / 2, y_size};
        } else if (fmt == PixelFormat::YUV420 && plane_count == 3) {
            int y_size = width * height;
            int uv_stride = width / 2;
            int uv_h = height / 2;
            int u_size = uv_stride * uv_h;
            frame->planes[0] = {buffer->data(), width, y_size, 0};
            frame->planes[1] = {buffer->data() + y_size, uv_stride, u_size, y_size};
            frame->planes[2] = {buffer->data() + y_size + u_size, uv_stride, u_size, y_size + u_size};
        }

        PushToDownstream(frame);
        if (frame_id_ % 100 == 0) {
            std::cout << "[input_rtsp] recv frame " << frame->frame_id
                      << " queue=" << (out_queues_.empty() ? 0 : out_queues_[0]->Size())
                      << " drops=" << (out_queues_.empty() ? 0 : out_queues_[0]->DroppedCount())
                      << "\n";
        }
    }

    void Cleanup(AVFormatContext* fmt, AVCodecContext* dec, AVPacket* pkt, AVFrame* frm) {
        if (dec) avcodec_free_context(&dec);
        if (fmt) avformat_close_input(&fmt);
        if (pkt) av_packet_free(&pkt);
        if (frm) av_frame_free(&frm);
    }
#endif

#if defined(RK3588_ENABLE_FFMPEG) && defined(RK3588_ENABLE_MPP)
    struct MppDecoderWrapper {
        ~MppDecoderWrapper() { Shutdown(); }

        bool Init(MppCodingType type) {
            MPP_RET ret = mpp_create(&ctx, &mpi);
            if (ret != MPP_OK) return false;
            ret = mpp_init(ctx, MPP_CTX_DEC, type);
            if (ret != MPP_OK) return false;
            MppDecCfg cfg = nullptr;
            mpp_dec_cfg_init(&cfg);
            ret = mpi->control(ctx, MPP_DEC_GET_CFG, cfg);
            if (ret != MPP_OK) return false;
            ret = mpp_dec_cfg_set_u32(cfg, "base:split_parse", 1);
            if (ret != MPP_OK) return false;
            ret = mpi->control(ctx, MPP_DEC_SET_CFG, cfg);
            mpp_dec_cfg_deinit(cfg);
            return ret == MPP_OK;
        }

        void Shutdown() {
            if (packet) {
                mpp_packet_deinit(&packet);
                packet = nullptr;
            }
            if (frame_group) {
                mpp_buffer_group_put(frame_group);
                frame_group = nullptr;
            }
            if (ctx) {
                mpp_destroy(ctx);
                ctx = nullptr;
                mpi = nullptr;
            }
        }

        bool Decode(const uint8_t* data, size_t size, bool eos, int64_t pts_ms,
                    const std::function<void(MppFrame)>& on_frame) {
            if (!ctx || !mpi || !data || size == 0) return false;

            if (!packet) {
                if (mpp_packet_init(&packet, nullptr, 0) != MPP_OK) return false;
            }
            void* pkt_data = const_cast<uint8_t*>(data);
            mpp_packet_set_data(packet, pkt_data);
            mpp_packet_set_size(packet, size);
            mpp_packet_set_pos(packet, pkt_data);
            mpp_packet_set_length(packet, size);
            mpp_packet_set_pts(packet, pts_ms);
            if (eos) mpp_packet_set_eos(packet);

            bool pkt_done = false;
            while (true) {
                if (!pkt_done) {
                    if (mpi->decode_put_packet(ctx, packet) == MPP_OK) {
                        pkt_done = true;
                    } else {
                        usleep(2000);
                    }
                }

                int get_frm = 0;
                MPP_RET ret = mpi->decode_get_frame(ctx, &frame);
                if (ret == MPP_ERR_TIMEOUT) {
                    usleep(2000);
                    continue;
                }
                if (ret != MPP_OK) break;

                if (frame) {
                    if (mpp_frame_get_info_change(frame)) {
                        HandleInfoChange();
                    } else {
                        if (on_frame) on_frame(frame);
                    }
                    get_frm = 1;
                    mpp_frame_deinit(&frame);
                    frame = nullptr;
                }

                if (get_frm) continue;
                break;
            }

            mpp_packet_deinit(&packet);
            packet = nullptr;
            return true;
        }

        void HandleInfoChange() {
            if (!ctx || !mpi || !frame) return;
            MppBufferGroup group = frame_group;
            size_t buf_size = mpp_frame_get_buf_size(frame);
            if (!group) {
                if (mpp_buffer_group_get_internal(&group, MPP_BUFFER_TYPE_DRM, NULL) != MPP_OK) return;
                frame_group = group;
            } else {
                mpp_buffer_group_clear(group);
            }
            mpp_buffer_group_limit_config(group, buf_size, 24);
            mpi->control(ctx, MPP_DEC_SET_EXT_BUF_GROUP, group);
            mpi->control(ctx, MPP_DEC_SET_INFO_CHANGE_READY, nullptr);
        }

        MppCtx ctx = nullptr;
        MppApi* mpi = nullptr;
        MppBufferGroup frame_group = nullptr;
        MppPacket packet = nullptr;
        MppFrame frame = nullptr;
    };

    void LoopFfmpegMpp() {
        ApplyAffinity();
        using namespace std::chrono;

        int backoff = std::max(1, reconnect_sec_);
        const int backoff_max = std::max(backoff, reconnect_backoff_max_sec_);

        while (running_.load()) {
            AVFormatContext* fmt_ctx = nullptr;
            AVPacket* pkt = av_packet_alloc();
            int video_stream = -1;
            AVRational time_base{1, 1000};

            AVDictionary* opts = nullptr;
            if (ffmpeg_force_tcp_) av_dict_set(&opts, "rtsp_transport", "tcp", 0);
            if (avformat_open_input(&fmt_ctx, url_.c_str(), nullptr, &opts) < 0) {
                std::cerr << "[input_rtsp] avformat_open_input failed: " << url_ << "\n";
                av_dict_free(&opts);
                Cleanup(fmt_ctx, nullptr, pkt, nullptr);
                if (fallback_to_stub_on_fail_) {
                    LoopStub();
                    return;
                }
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }
            av_dict_free(&opts);
            if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
                std::cerr << "[input_rtsp] avformat_find_stream_info failed\n";
                Cleanup(fmt_ctx, nullptr, pkt, nullptr);
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }
            for (unsigned i = 0; i < fmt_ctx->nb_streams; ++i) {
                if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
                    video_stream = static_cast<int>(i);
                    time_base = fmt_ctx->streams[i]->time_base;
                    break;
                }
            }
            if (video_stream < 0) {
                std::cerr << "[input_rtsp] no video stream\n";
                Cleanup(fmt_ctx, nullptr, pkt, nullptr);
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }

            MppCodingType coding = MPP_VIDEO_CodingAVC;
            auto codec_id = fmt_ctx->streams[video_stream]->codecpar->codec_id;
            if (codec_id == AV_CODEC_ID_H264) coding = MPP_VIDEO_CodingAVC;
            else if (codec_id == AV_CODEC_ID_HEVC) coding = MPP_VIDEO_CodingHEVC;
            else {
                std::cerr << "[input_rtsp] unsupported codec for mpp\n";
                Cleanup(fmt_ctx, nullptr, pkt, nullptr);
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }

            MppDecoderWrapper dec;
            if (!dec.Init(coding)) {
                std::cerr << "[input_rtsp] mpp init failed\n";
                Cleanup(fmt_ctx, nullptr, pkt, nullptr);
                std::this_thread::sleep_for(seconds(backoff));
                backoff = std::min(backoff_max, backoff * 2);
                continue;
            }

            backoff = std::max(1, reconnect_sec_);
            int read_fail = 0;
            while (running_.load()) {
                if (av_read_frame(fmt_ctx, pkt) < 0) {
                    if (++read_fail >= 50) break;
                    std::this_thread::sleep_for(milliseconds(10));
                    continue;
                }
                read_fail = 0;
                if (pkt->stream_index != video_stream) {
                    av_packet_unref(pkt);
                    continue;
                }
                int64_t pts_ms = pkt->pts == AV_NOPTS_VALUE ? 0
                                  : av_rescale_q(pkt->pts, time_base, {1, 1000});
                dec.Decode(pkt->data, pkt->size, false, pts_ms,
                           [&](MppFrame frm) { PushFrameFromMpp(frm); });
                av_packet_unref(pkt);
            }

            Cleanup(fmt_ctx, nullptr, pkt, nullptr);
            if (!running_.load()) break;
            std::this_thread::sleep_for(seconds(std::max(1, reconnect_sec_)));
        }
    }

    void PushFrameFromMpp(MppFrame frm) {
        MppBuffer buf = mpp_frame_get_buffer(frm);
        if (!buf) return;
        mpp_buffer_inc_ref(buf);
        auto owner = std::shared_ptr<void>(buf, [](void* p) { mpp_buffer_put(reinterpret_cast<MppBuffer>(p)); });

        int width = mpp_frame_get_width(frm);
        int height = mpp_frame_get_height(frm);
        int hor_stride = mpp_frame_get_hor_stride(frm);
        int ver_stride = mpp_frame_get_ver_stride(frm);
        MppFrameFormat mpp_fmt = mpp_frame_get_fmt(frm);

        PixelFormat fmt = PixelFormat::UNKNOWN;
        int plane_count = 0;
        if (mpp_fmt == MPP_FMT_YUV420SP
#ifdef MPP_FMT_YUV420SP_10BIT
            || mpp_fmt == MPP_FMT_YUV420SP_10BIT
#endif
        ) {
            fmt = PixelFormat::NV12;
            plane_count = 2;
        } else if (mpp_fmt == MPP_FMT_YUV420P
#ifdef MPP_FMT_YUV420P10
                   || mpp_fmt == MPP_FMT_YUV420P10
#endif
        ) {
            fmt = PixelFormat::YUV420;
            plane_count = 3;
        }

        if (plane_count == 0) {
            return;
        }

        auto frame = std::make_shared<Frame>();
        frame->width = width;
        frame->height = height;
        frame->format = fmt;
        frame->stride = hor_stride;
        frame->plane_count = plane_count;
        frame->dma_fd = mpp_buffer_get_fd(buf);
        frame->data = static_cast<uint8_t*>(mpp_buffer_get_ptr(buf));
        frame->data_size = mpp_buffer_get_size(buf);
        frame->data_owner = owner;
        frame->frame_id = ++frame_id_;
        frame->pts = mpp_frame_get_pts(frm);

        if (fmt == PixelFormat::NV12 && plane_count == 2) {
            int y_size = hor_stride * ver_stride;
            frame->planes[0] = {frame->data, hor_stride, y_size, 0};
            frame->planes[1] = {frame->data + y_size, hor_stride, y_size / 2, y_size};
        } else if (fmt == PixelFormat::YUV420 && plane_count == 3) {
            int y_size = hor_stride * ver_stride;
            int uv_stride = hor_stride / 2;
            int uv_h = ver_stride / 2;
            int u_size = uv_stride * uv_h;
            frame->planes[0] = {frame->data, hor_stride, y_size, 0};
            frame->planes[1] = {frame->data + y_size, uv_stride, u_size, y_size};
            frame->planes[2] = {frame->data + y_size + u_size, uv_stride, u_size, y_size + u_size};
        }

        PushToDownstream(frame);

        if (frame_id_ % 100 == 0) {
            std::cout << "[input_rtsp] mpp frame " << frame->frame_id
                      << " queue=" << (out_queues_.empty() ? 0 : out_queues_[0]->Size())
                      << " drops=" << (out_queues_.empty() ? 0 : out_queues_[0]->DroppedCount())
                      << "\n";
        }
    }
#endif

    std::string id_;
    std::string url_;
    int fps_ = 25;
    int width_ = 1920;
    int height_ = 1080;
    std::atomic<bool> running_{false};
    std::vector<std::shared_ptr<SpscQueue<FramePtr>>> out_queues_;
    std::thread worker_;
    uint64_t frame_id_ = 0;
    bool use_ffmpeg_ = false;
    bool use_mpp_ = true;
    bool ffmpeg_force_tcp_ = true;

    int reconnect_sec_ = 5;
    int reconnect_backoff_max_sec_ = 30;
    bool fallback_to_stub_on_fail_ = false;
    std::vector<int> cpu_affinity_;
};

REGISTER_NODE(InputRtspNode, "input_rtsp");

}  // namespace rk3588