#include "webrtc_stream.h"

#include "camera_v4l2.h"
#include "hw_scaler.h"
#include "isp_aiq.h"
#include "mpp_encoder.h"
#include "test_pattern.h"
#include "webrtc_pusher.h"
#include "ws_signaling.h"

#include <algorithm>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <memory>
#include <mutex>
#include <thread>
#include <vector>

namespace webrtcpush {

int run_stream(const StreamConfig &a, volatile std::sig_atomic_t &running){
    /* AIQ must start before RK_MPI_SYS_Init and before opening the ISP node. */
    isp_aiq_t *isp = nullptr;
    if(a.source == "camera" && a.use_isp){
        isp_aiq_cfg_t ic{};
        ic.camera_id = a.camera_id;
        ic.iq_dir = a.iq_dir.c_str();
        isp = isp_aiq_start(&ic);
        if(!isp){
            std::cerr << "[main] ISP/AIQ startup failed; use --no-isp only for diagnostics\n";
            return 6;
        }
    }

    /* ---- 编码器 (目标 720p NV12) ---- */
    mpp_encoder_cfg_t ec{};
    ec.width = a.width; ec.height = a.height;
    ec.fps = a.fps; ec.bitrate_kbps = a.bitrate;
    ec.gop = a.gop; ec.profile = a.profile.c_str();
    mpp_encoder_t *enc = mpp_encoder_create(&ec);
    if(!enc){
        std::cerr << "encoder create failed\n";
        isp_aiq_stop(isp);
        return 2;
    }

    int dst_hs = 0, dst_vs = 0;
    mpp_encoder_get_stride(enc, &dst_hs, &dst_vs);
    const size_t enc_frame_size = (size_t)dst_hs * dst_vs * 3 / 2;

    /* Open the camera before RGA so the scaler uses the negotiated V4L2 stride. */
    camera_v4l2_t *cam = nullptr;
    int cam_hs = a.src_w, cam_vs = a.src_h;
    if(a.source == "camera"){
        camera_cfg_t cc{};
        cc.dev    = a.device.c_str();
        cc.width  = a.src_w;
        cc.height = a.src_h;
        cc.fps    = a.fps;
        cc.format = a.format.c_str();
        cam = camera_v4l2_open(&cc);
        if(!cam){
            std::cerr << "[cam] open failed\n";
            isp_aiq_stop(isp);
            mpp_encoder_destroy(enc);
            return 4;
        }
        camera_v4l2_get_stride(cam, &cam_hs, &cam_vs);
        if(cam_hs < a.src_w) cam_hs = a.src_w;
        if(cam_vs < a.src_h) cam_vs = a.src_h;
    }

    /* ---- 硬件缩放器 (sensor 原图 -> 720p) ---- */
    std::unique_ptr<webrtcpush::HwScaler> scaler;
    webrtcpush::ScalerCfg sc{};
    if(cam && a.use_scaler && (a.src_w != a.width || a.src_h != a.height)){
        sc.src_w = a.src_w; sc.src_h = a.src_h;
        sc.src_fmt = webrtcpush::PixelFmt::NV12;
        sc.src_hor_stride = cam_hs;
        sc.src_ver_stride = cam_vs;
        sc.dst_w = a.width;  sc.dst_h = a.height;
        sc.dst_fmt = webrtcpush::PixelFmt::NV12;
        sc.dst_hor_stride = dst_hs;
        sc.dst_ver_stride = dst_vs;
        scaler = std::make_unique<webrtcpush::HwScaler>(sc);
        if(!scaler->valid()){
            std::cerr << "[scaler] init failed\n";
            camera_v4l2_close(cam);
            isp_aiq_stop(isp);
            mpp_encoder_destroy(enc);
            return 5;
        }
    }

    /* ---- WebRTC 推流器 ---- */
    webrtcpush::PusherConfig pc;
    pc.h264_profile = a.profile;
    pc.fps = a.fps;
    pc.init_bitrate_kbps = a.bitrate;
    pc.min_bitrate_kbps = std::max(100, a.bitrate / 4);
    pc.max_bitrate_kbps = std::max(a.bitrate, a.bitrate * 2);
    webrtcpush::WebrtcPusher pusher(pc);

    /* ---- WS 信令 ---- */
    std::string ws_url = a.signal_url + "?dev=" + a.room;
    webrtcpush::WsSignaling ws(ws_url);

    /*
     * The deployed pages can negotiate in either direction, but in the normal
     * index.html-first flow client.html becomes the offerer. Its full Chrome
     * offer exceeds this signaling server's 8 KiB message limit, so RV1106
     * must use the other supported direction: compact video-only offer from
     * the device, answer from client.html. The server does not notify the first
     * room member when a viewer joins and does not retain old messages.
     */
    std::atomic<int> room_members{0};
    std::atomic<bool> reset_session{false};
    std::atomic<bool> restart_offer{false};
    std::mutex session_action_mutex;
    std::mutex offer_mutex;
    std::string retry_offer_sdp;
    std::mutex candidate_mutex;
    std::vector<std::string> pending_remote_candidates;

    /* 信令消息路由 */
    ws.set_message_callback([&](const webrtcpush::SignalMessage &msg){
        switch(msg.type){
        case webrtcpush::SignalMessage::Type::Id:
            std::cerr << "[ws] ready; preparing compact device video offer\n";
            restart_offer = true;
            break;
        case webrtcpush::SignalMessage::Type::Offer:
        {
            std::lock_guard<std::mutex> action_lk(session_action_mutex);
            reset_session = false;
            restart_offer = false;
            {
                std::lock_guard<std::mutex> offer_lk(offer_mutex);
                retry_offer_sdp.clear();
            }
            std::cerr << "[ws] got browser offer from " << msg.sender_id
                      << " (" << msg.sdp.size() << " bytes)\n";
            if(!pusher.handle_offer(msg.sdp)){
                std::cerr << "[rtc] browser offer rejected\n";
                break;
            }

            std::vector<std::string> candidates;
            {
                std::lock_guard<std::mutex> candidate_lk(candidate_mutex);
                candidates.swap(pending_remote_candidates);
            }
            for(const auto &candidate : candidates){
                if(!pusher.add_remote_candidate(candidate))
                    std::cerr << "[rtc] queued remote ICE candidate rejected\n";
            }
            break;
        }
        case webrtcpush::SignalMessage::Type::Answer:
            std::cerr << "[ws] got browser answer (" << msg.sdp.size() << " bytes)\n";
            if(std::getenv("WEBRTCPUSH_DUMP_SDP"))
                std::cerr << "[rtc] remote answer SDP begin\n" << msg.sdp
                          << "[rtc] remote answer SDP end\n";
            if(!pusher.set_remote_answer(msg.sdp)){
                std::cerr << "[rtc] browser answer rejected\n";
            }else{
                std::lock_guard<std::mutex> offer_lk(offer_mutex);
                retry_offer_sdp.clear();
            }
            break;
        case webrtcpush::SignalMessage::Type::Num: {
            room_members = msg.num;
            std::cerr << "[ws] room members=" << msg.num << "\n";
            break;
        }
        case webrtcpush::SignalMessage::Type::Candidate:
            std::cerr << "[ws] got remote ICE candidate " << msg.candidate_json << "\n";
            if(pusher.has_session()){
                if(!pusher.add_remote_candidate(msg.candidate_json))
                    std::cerr << "[rtc] remote ICE candidate rejected\n";
            }else{
                std::lock_guard<std::mutex> lk(candidate_mutex);
                pending_remote_candidates.push_back(msg.candidate_json);
            }
            break;
        case webrtcpush::SignalMessage::Type::Leave:
            std::cerr << "[ws] peer " << msg.sender_id << " left"
                      << " (room members=" << room_members.load() << ")\n";
            room_members = 1;
            reset_session = true;
            restart_offer = true;
            break;
        case webrtcpush::SignalMessage::Type::Overload:
            std::cerr << "[ws] room overloaded; this device connection was rejected"
                      << " (memberId=" << msg.message << ")\n";
            break;
        case webrtcpush::SignalMessage::Type::Pong:
        default:
            break;  /* 心跳等, 不动作 */
        }
    });

    /* pusher 回调 -> WS 信令 */
    pusher.set_local_description_callback(
        [&](const std::string &sdp, bool is_offer){
            if(is_offer){
                {
                    std::lock_guard<std::mutex> offer_lk(offer_mutex);
                    retry_offer_sdp = sdp;
                }
                /* The first callback precedes ICE gathering. Only the complete
                 * SDP is useful to a future viewer because the signaling server
                 * does not retain separately trickled candidates. */
                if(sdp.find("a=end-of-candidates") == std::string::npos)
                    return;
            }
            ws.send_description(sdp, is_offer);
        });
    pusher.set_local_candidate_callback([&](const std::string &candidate_json){
        /* Device-offer candidates are embedded in the complete SDP. In answer
         * mode candidates still need to be trickled to the browser. */
        if(pusher.session_pending()) return;
        ws.send_candidate(candidate_json);
    });
    pusher.set_connected_callback([&](bool connected){
        if(connected){
            std::cerr << "[main] webrtc connected, requesting IDR\n";
            {
                std::lock_guard<std::mutex> offer_lk(offer_mutex);
                retry_offer_sdp.clear();
            }
            mpp_encoder_request_idr(enc);
        }else{
            std::cerr << "[main] webrtc failed; rebuilding compact device offer\n";
            reset_session = true;
            restart_offer = true;
        }
    });
    pusher.set_bitrate_adjust_callback([&](int new_kbps){
        mpp_encoder_set_bitrate(enc, new_kbps);
    });

    /* 编码器输出 -> pusher.feed_nalu */
    mpp_encoder_set_nalu_callback(enc,
        [](const uint8_t *data, size_t len, uint64_t pts, int is_key, void *ud){
            auto *p = static_cast<webrtcpush::WebrtcPusher*>(ud);
            p->feed_nalu(data, len, is_key != 0);
        }, &pusher);

    /* ---- 启动 WS ---- */
    if(!ws.start()){
        std::cerr << "ws start failed\n";
        if(cam) camera_v4l2_close(cam);
        isp_aiq_stop(isp);
        mpp_encoder_destroy(enc);
        return 3;
    }

    /* Session maintenance runs outside WebSocket callbacks so closing a
     * PeerConnection cannot re-enter the signaling callback stack. The server
     * does not notify the first room member when a viewer joins, therefore the
     * complete compact offer is replayed until a browser answer is received. */
    std::thread session_th([&](){
        auto last_offer_send = std::chrono::steady_clock::now() -
                               std::chrono::seconds(1);
        uint64_t offer_retry_count = 0;
        while(running){
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
            if(!running) break;

            if(reset_session.load() || restart_offer.load()){
                std::lock_guard<std::mutex> action_lk(session_action_mutex);
                const bool should_restart = restart_offer.exchange(false);
                if(reset_session.exchange(false) || should_restart){
                    pusher.reset();
                    {
                        std::lock_guard<std::mutex> candidate_lk(candidate_mutex);
                        pending_remote_candidates.clear();
                    }
                    {
                        std::lock_guard<std::mutex> offer_lk(offer_mutex);
                        retry_offer_sdp.clear();
                    }
                    if(should_restart){
                        std::cerr << "[ws] start compact H264 device offer\n";
                        if(!pusher.start())
                            std::cerr << "[rtc] compact device offer failed\n";
                    }
                }
                last_offer_send = std::chrono::steady_clock::now();
                offer_retry_count = 0;
                continue;
            }

            const auto now = std::chrono::steady_clock::now();
            /* The signaling server does not retain offers. The stock viewer
             * starts its own oversized Chromium offer after joining, so the
             * compact device offer must arrive first. 300 ms leaves enough
             * time for the viewer to answer without queueing rapid repeats. */
            if(pusher.session_pending() &&
               now - last_offer_send >= std::chrono::milliseconds(300)){
                std::string sdp;
                {
                    std::lock_guard<std::mutex> offer_lk(offer_mutex);
                    sdp = retry_offer_sdp;
                }
                if(sdp.find("a=end-of-candidates") != std::string::npos){
                    const bool sent = ws.send_description(sdp, true, false);
                    ++offer_retry_count;
                    if(offer_retry_count == 1 || offer_retry_count % 100 == 0){
                        std::cerr << "[ws] replay compact offer count="
                                  << offer_retry_count
                                  << " sent=" << (sent ? "yes" : "no") << "\n";
                    }
                }
                last_offer_send = now;
            }
        }
    });

    /* ---- 采集/编码循环 ---- */
    /* RGA 路径使用 HwScaler 内部的连续 CMA 缓冲；其他路径使用普通帧缓冲。 */
    std::vector<uint8_t> enc_frame(scaler ? 0 : enc_frame_size);

    std::cerr << "[main] running " << a.width << "x" << a.height << "@" << a.fps
              << " src=" << a.source;
    if(scaler) std::cerr << " (scaled from " << a.src_w << "x" << a.src_h << ")";
    std::cerr << " enc_frame=" << enc_frame_size << "\n";

    auto period = std::chrono::microseconds(1000000 / std::max(1, a.fps));
    uint32_t fidx = 0;
    while(running){
        auto t0 = std::chrono::steady_clock::now();
        uint8_t *encode_data = scaler ? scaler->dst_data() : enc_frame.data();
        size_t encode_size = scaler ? scaler->dst_frame_size() : enc_frame.size();

        if(cam){
            if(scaler){
                /* ISP 直接写 CMA 输入，RGA 缩放到 CMA 编码帧。 */
                if(camera_v4l2_get_frame(cam, scaler->src_data(), scaler->src_frame_size(),
                                         cam_hs, cam_vs, a.src_h) < 0){
                    if(running) std::cerr << "[cam] get_frame fail\n";
                    running = 0;
                    break;
                }
                if(scaler->scale() != 0){
                    std::cerr << "[scaler] scale fail\n";
                    running = 0;
                    break;
                }
            }else{
                /* --no-scaler: ISP 输出直接写到 enc_frame (按编码器 stride 布局) */
                if(camera_v4l2_get_frame(cam, enc_frame.data(), enc_frame.size(),
                                         dst_hs, dst_vs, a.height) < 0){
                    if(running) std::cerr << "[cam] get_frame fail\n";
                    running = 0;
                    break;
                }
            }
        }else{
            /* test 彩条: 直接生成 720p NV12 (无需缩放) */
            test_pattern_nv12(enc_frame.data(), dst_hs, dst_vs, a.height, fidx);
        }

        if(mpp_encoder_encode(enc, encode_data, encode_size,
                              (uint64_t)fidx * 1000000ULL / a.fps) < 0){
            std::cerr << "[enc] encode fail\n";
            running = 0;
            break;
        }
        fidx++;

        std::this_thread::sleep_until(t0 + period);
    }

    running = 0;
    std::cerr << "[main] stopping\n";
    std::cerr << "[main] stopping session maintenance\n";
    if(session_th.joinable()) session_th.join();
    std::cerr << "[main] stopping signaling\n";
    ws.stop();
    std::cerr << "[main] stopping WebRTC\n";
    pusher.stop();
    std::cerr << "[main] closing camera\n";
    if(cam) camera_v4l2_close(cam);
    std::cerr << "[main] stopping ISP/AIQ\n";
    isp_aiq_stop(isp);
    std::cerr << "[main] stopping encoder\n";
    mpp_encoder_destroy(enc);
    std::cerr << "[main] stopped\n";
    return 0;
}

} // namespace webrtcpush
