/*
 * libdatachannel WebRTC 推流器 (C++) — 应答模式为主, 主动 offer 兜底
 *
 * 协议实测 (signal.js / client.js):
 *   拉流端进房间 -> 主动发 offer (recvonly video+audio + "init" DataChannel)
 *   本设备 (推流端) handle_offer: setRemoteDescription ->
 *     解析 offer 的 H264 PT/fmtp -> 同参数加 sendonly track ->
 *     setLocalDescription 生成 answer -> 回传
 *
 * H.264 协商关键: libdatachannel 生成 answer 时同 mid 有本地 track 就直接
 * 用本地 PT/fmtp, 不做 codec 匹配 —— 必须手动把本地 track 的 PT/fmtp 设置成
 * 与对方 offer 完全一致, 否则浏览器拒收 answer。
 */
#include "webrtc_pusher.h"

#include <rtc/rtc.hpp>
#include <rtc/candidate.hpp>
#include <rtc/common.hpp>
#include <rtc/h264rtppacketizer.hpp>
#include <rtc/rtppacketizationconfig.hpp>
#include <nlohmann/json.hpp>

#include <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <mutex>
#include <thread>
#include <vector>

namespace webrtcpush {

using json = nlohmann::json;

/* ---------- SDP 解析: 从 offer 里选 H264 PT + fmtp + video mid ---------- */

struct H264Choice {
    std::string mid;     // video m-line 的 mid
    int         pt  = -1;
    std::string fmtp;    // a=fmtp:<pt> 后面的完整参数串
    bool        ok  = false;
};

static std::string trim(const std::string &s){
    size_t b = s.find_first_not_of(" \t\r\n");
    if(b == std::string::npos) return "";
    size_t e = s.find_last_not_of(" \t\r\n");
    return s.substr(b, e - b + 1);
}

struct AnnexBNalu {
    size_t offset = 0;
    size_t size = 0;
    int type = -1;
};

static std::vector<AnnexBNalu> split_annexb(const uint8_t *data, size_t len){
    std::vector<AnnexBNalu> out;
    if(!data || len == 0) return out;

    auto start_code = [&](size_t pos, size_t *prefix) -> bool {
        if(pos + 3 <= len && data[pos] == 0 && data[pos + 1] == 0){
            if(data[pos + 2] == 1){ *prefix = 3; return true; }
            if(pos + 4 <= len && data[pos + 2] == 0 && data[pos + 3] == 1){
                *prefix = 4;
                return true;
            }
        }
        return false;
    };

    size_t pos = 0;
    while(pos < len){
        size_t prefix = 0;
        while(pos < len && !start_code(pos, &prefix)) ++pos;
        if(pos >= len) break;
        const size_t begin = pos;
        const size_t header = pos + prefix;
        pos = header + 1;
        size_t next = pos;
        size_t ignored = 0;
        while(next < len && !start_code(next, &ignored)) ++next;
        if(header < len)
            out.push_back({begin, next - begin, data[header] & 0x1f});
        pos = next;
    }

    if(out.empty()) out.push_back({0, len, data[0] & 0x1f});
    return out;
}

/* 按 a=rtpmap / a=fmtp 行挑浏览器 offer 里我们也能送的 H264:
 *   优先 packetization-mode=1 且 high profile (640c.., 与 MPP 默认 high 匹配)
 *   其次 packetization-mode=1 + 42e01f (constrained baseline)
 *   再次任意 packetization-mode=1
 *   最后任意 H264 */
static bool profile_matches(const std::string &fmtp, const std::string &profile){
    std::string low = fmtp;
    std::transform(low.begin(), low.end(), low.begin(), ::tolower);
    const std::string key = "profile-level-id=";
    size_t pos = low.find(key);
    if(pos == std::string::npos) return profile == "baseline";
    pos += key.size();
    if(pos + 2 > low.size()) return false;
    const std::string profile_idc = low.substr(pos, 2);
    if(profile == "high") return profile_idc == "64";
    if(profile == "main") return profile_idc == "4d";
    return profile_idc == "42";
}

static std::string profile_fmtp(const std::string &profile){
    const char *profile_level_id = profile == "high" ? "640c1f"
                                 : profile == "main" ? "4d001f"
                                                     : "42e01f";
    return std::string("level-asymmetry-allowed=1;packetization-mode=1;") +
           "profile-level-id=" + profile_level_id;
}

static H264Choice pick_h264_from_offer(const std::string &sdp,
                                       const std::string &profile){
    H264Choice out;
    std::vector<std::string> lines;
    {
        size_t pos = 0;
        while(pos <= sdp.size()){
            size_t nl = sdp.find('\n', pos);
            if(nl == std::string::npos){ lines.push_back(sdp.substr(pos)); break; }
            lines.push_back(sdp.substr(pos, nl - pos));
            pos = nl + 1;
        }
    }

    /* 1. 找 m=video 段的 mid */
    bool in_video = false;
    for(const auto &raw : lines){
        std::string l = trim(raw);
        if(l.rfind("m=video", 0) == 0){ in_video = true; continue; }
        if(l.rfind("m=", 0) == 0 && l.rfind("m=video", 0) != 0){ if(in_video) break; continue; }
        if(!in_video) continue;
        if(l.rfind("a=mid:", 0) == 0){ out.mid = l.substr(6); break; }
    }

    /* 2. 收集 video 段的 H264 rtpmap 和 fmtp */
    std::vector<int> h264_pts;
    std::vector<std::pair<int,std::string>> fmts;
    in_video = false;
    for(const auto &raw : lines){
        std::string l = trim(raw);
        if(l.rfind("m=video", 0) == 0){ in_video = true; continue; }
        if(l.rfind("m=", 0) == 0 && l.rfind("m=video", 0) != 0){ in_video = false; continue; }
        if(!in_video) continue;

        if(l.rfind("a=rtpmap:", 0) == 0){
            /* a=rtpmap:<pt> H264/90000 */
            std::string rest = l.substr(9);
            size_t sp = rest.find(' ');
            if(sp == std::string::npos) continue;
            int pt = atoi(rest.substr(0, sp).c_str());
            std::string fmt = trim(rest.substr(sp + 1));
            if(fmt.rfind("H264/", 0) == 0 || fmt.rfind("h264/", 0) == 0){
                h264_pts.push_back(pt);
            }
        }else if(l.rfind("a=fmtp:", 0) == 0){
            std::string rest = l.substr(7);
            size_t sp = rest.find(' ');
            if(sp == std::string::npos) continue;
            int pt = atoi(rest.substr(0, sp).c_str());
            fmts.emplace_back(pt, trim(rest.substr(sp + 1)));
        }
    }

    /* 3. 打分选最优 */
    struct Cand { int pt; std::string fmtp; int score; };
    std::vector<Cand> cands;
    for(int pt : h264_pts){
        std::string fmtp;
        for(auto &f : fmts) if(f.first == pt){ fmtp = f.second; break; }
        if(!profile_matches(fmtp, profile)) continue;
        int score = 0;
        std::string low = fmtp;
        std::transform(low.begin(), low.end(), low.begin(), ::tolower);
        if(low.find("packetization-mode=1") != std::string::npos) score += 2;
        if(low.find("level-asymmetry-allowed=1") != std::string::npos) score += 1;
        cands.push_back({pt, fmtp, score});
    }
    if(cands.empty()) return out;

    const Cand *best = &cands[0];
    for(const auto &c : cands) if(c.score > best->score) best = &c;

    out.pt   = best->pt;
    out.fmtp = best->fmtp;
    out.ok   = true;
    return out;
}

/* ---------- Impl ---------- */

struct WebrtcPusher::Impl {
    PusherConfig cfg;

    std::shared_ptr<rtc::PeerConnection> pc;
    std::shared_ptr<rtc::Track>          track;
    std::shared_ptr<rtc::RtpPacketizationConfig> rtp_cfg;
    std::shared_ptr<rtc::H264RtpPacketizer>      packetizer;
    std::shared_ptr<rtc::DataChannel>    remote_dc;  /* 拉流端建的 "init" DC (应答模式) */
    std::shared_ptr<rtc::DataChannel>    local_dc;   /* 兜底模式自己建的 DC */
    int cur_pt = 96;                                  /* 本会话实际使用的 PT */

    LocalDescriptionCallback desc_cb;
    LocalCandidateCallback   cand_cb;
    ConnectedCallback        conn_cb;
    BitrateAdjustCallback    bitrate_cb;
    std::mutex cb_mutex;
    std::mutex pc_mutex;      /* 保护 pc/track 的整体重建 */

    std::atomic<bool> connected{false};
    std::atomic<bool> pending{false};     /* 兜底 offer 已发, 等 answer */
    std::atomic<bool> stopped{false};
    std::atomic<int>  cur_bitrate_kbps{0};
    std::atomic<uint32_t> frame_idx{0};
    std::atomic<uint64_t> sent_packets{0};
    std::vector<std::byte> cached_sps;
    std::vector<std::byte> cached_pps;

    std::thread monitor_th;
    std::mutex monitor_mutex;
    std::condition_variable monitor_cv;

    explicit Impl(PusherConfig c) : cfg(std::move(c)) {}

    bool create_status_channel_locked(){
        if(local_dc) return true;
        try {
            rtc::DataChannelInit dc_cfg;
            dc_cfg.reliability.maxRetransmits = 3;
            local_dc = pc->createDataChannel("myDataChannel", dc_cfg);
            return local_dc != nullptr;
        }catch(const std::exception &e){
            std::cerr << "[rtc] createDataChannel warning: " << e.what() << "\n";
            return false;
        }
    }

    void ensure_status_thread(){
        if(monitor_th.joinable()) return;
        monitor_th = std::thread([this](){
            while(!stopped){
                std::unique_lock<std::mutex> wait_lk(monitor_mutex);
                if(monitor_cv.wait_for(
                       wait_lk,
                       std::chrono::milliseconds(cfg.bitrate_adjust_period_ms),
                       [this](){ return stopped.load(); })) break;
                wait_lk.unlock();
                if(!connected) continue;

                std::shared_ptr<rtc::DataChannel> dc;
                {
                    std::lock_guard<std::mutex> lk(pc_mutex);
                    dc = local_dc;
                }
                if(!dc) continue;

                char value[32];
                std::snprintf(value, sizeof(value), "%.2f",
                              cur_bitrate_kbps.load() / 1000.0);
                try { dc->send(std::string(value)); }catch(...){}
            }
        });
    }

    void teardown_locked(){
        connected = false;
        pending   = false;
        try {
            if(remote_dc) remote_dc->close();
            if(local_dc)  local_dc->close();
            if(track)     track->close();
            if(pc)        pc->close();
        }catch(...){}
        remote_dc.reset(); local_dc.reset();
        track.reset(); packetizer.reset(); rtp_cfg.reset(); pc.reset();
    }

    /* 建新 PeerConnection 并接好回调 */
    bool create_pc(){
        rtc::Configuration rcfg;
        for(const auto &url : cfg.ice_servers){
            if(url.rfind("turn:", 0) == 0){
                rtc::IceServer s(url);
                s.username   = cfg.turn_username;
                s.password   = cfg.turn_credential;
                rcfg.iceServers.push_back(s);
            }else{
                rcfg.iceServers.emplace_back(url);
            }
        }
        /* libjuice does not support TURN while UDP mux mode is enabled.
         * This device has one PeerConnection, so use the normal per-session
         * socket and keep host, STUN and TURN candidates available. */
        rcfg.enableIceUdpMux = false;

        pc = std::make_shared<rtc::PeerConnection>(rcfg);

        pc->onLocalDescription([this](rtc::Description d){
            bool is_offer = (d.type() == rtc::Description::Type::Offer);
            LocalDescriptionCallback cb;
            {
                std::lock_guard<std::mutex> lk(cb_mutex);
                cb = desc_cb;
            }
            if(cb){
                std::cerr << "[rtc] local " << (is_offer ? "offer" : "answer")
                          << " ready (" << std::string(d).size() << " bytes, pt="
                          << cur_pt << ")\n";
                cb(std::string(d), is_offer);
            }
        });

        pc->onLocalCandidate([this](rtc::Candidate c){
            /* 与 signal.js 完全一致: {memberId, candidate:{...}} 由信令层包 memberId */
            std::cerr << "[rtc] local ICE candidate mid=" << std::string(c.mid())
                      << " value=" << std::string(c.candidate()) << "\n";
            json j;
            j["candidate"]     = std::string(c.candidate());
            j["sdpMid"]        = std::string(c.mid());
            j["sdpMLineIndex"] = 0;
            LocalCandidateCallback cb;
            {
                std::lock_guard<std::mutex> lk(cb_mutex);
                cb = cand_cb;
            }
            if(cb) cb(j.dump());
        });

        pc->onGatheringStateChange([this](rtc::PeerConnection::GatheringState s){
            std::cerr << "[rtc] gathering-state=" << static_cast<int>(s) << "\n";
            if(s == rtc::PeerConnection::GatheringState::Complete && pending){
                if(!pc) return;
                auto description = pc->localDescription();
                if(!description) return;

                LocalDescriptionCallback cb;
                {
                    std::lock_guard<std::mutex> lk(cb_mutex);
                    cb = desc_cb;
                }
                if(cb){
                    const std::string complete_sdp = std::string(*description);
                    std::cerr << "[rtc] complete local offer ready ("
                              << complete_sdp.size() << " bytes, candidates inline)\n";
                    cb(complete_sdp, true);
                }
            }
        });

        pc->onIceStateChange([](rtc::PeerConnection::IceState s){
            std::cerr << "[rtc] ice-state=" << static_cast<int>(s) << "\n";
        });

        pc->onStateChange([this](rtc::PeerConnection::State s){
            std::cerr << "[rtc] state=" << static_cast<int>(s) << std::endl;
            bool now = (s == rtc::PeerConnection::State::Connected);
            bool was = connected.exchange(now);
            if((now && !was) || s == rtc::PeerConnection::State::Failed){
                ConnectedCallback cb;
                {
                    std::lock_guard<std::mutex> lk(cb_mutex);
                    cb = conn_cb;
                }
                if(cb) cb(now);
                if(now) frame_idx = 0;
            }
            if(s == rtc::PeerConnection::State::Failed ||
               s == rtc::PeerConnection::State::Closed){
                connected = false;
            }
        });

        /* 应答模式: 拉流端在 offer 里建的 "init" DataChannel 到达时拿到引用,
         * 用于 bufferedAmount 拥塞监控 */
        pc->onDataChannel([this](std::shared_ptr<rtc::DataChannel> dc){
            std::cerr << "[rtc] remote dataChannel '" << dc->label() << "'\n";
            remote_dc = dc;
            dc->onMessage(std::function<void(rtc::binary)>{},
                          [](std::string m){ (void)m; });
        });

        return true;
    }

    /* 按指定 PT/fmtp 挂 H264 sendonly track (mid 用对方 offer 的 video mid) */
    bool attach_track_locked(const std::string &mid, int pt, const std::string &fmtp){
        cur_pt = pt;
        rtc::Description::Video media(mid.empty() ? "video" : mid,
                                       rtc::Description::Direction::SendOnly);
        if(fmtp.empty()){
            media.addH264Codec(pt);   /* libdatachannel 默认 42e01f/pm=1 */
        }else{
            media.addH264Codec(pt, fmtp);
        }
        media.addSSRC(cfg.ssrc, cfg.stream_id);

        track = pc->addTrack(media);
        if(!track){
            std::cerr << "[rtc] addTrack failed\n";
            return false;
        }

        rtp_cfg = std::make_shared<rtc::RtpPacketizationConfig>(
            cfg.ssrc, cfg.stream_id, pt, cfg.video_clock);
        packetizer = std::make_shared<rtc::H264RtpPacketizer>(
            rtc::H264RtpPacketizer::Separator::StartSequence, rtp_cfg, cfg.mtu);
        track->setMediaHandler(packetizer);
        return true;
    }
};

/* ---------- 公开接口 ---------- */

WebrtcPusher::WebrtcPusher(const PusherConfig &cfg)
    : impl_(std::make_unique<Impl>(cfg)) {
    impl_->cur_bitrate_kbps = cfg.init_bitrate_kbps;
}

WebrtcPusher::~WebrtcPusher(){ stop(); }

void WebrtcPusher::set_local_description_callback(LocalDescriptionCallback cb){
    std::lock_guard<std::mutex> lk(impl_->cb_mutex);
    impl_->desc_cb = std::move(cb);
}
void WebrtcPusher::set_local_candidate_callback(LocalCandidateCallback cb){
    std::lock_guard<std::mutex> lk(impl_->cb_mutex);
    impl_->cand_cb = std::move(cb);
}
void WebrtcPusher::set_connected_callback(ConnectedCallback cb){
    std::lock_guard<std::mutex> lk(impl_->cb_mutex);
    impl_->conn_cb = std::move(cb);
}
void WebrtcPusher::set_bitrate_adjust_callback(BitrateAdjustCallback cb){
    std::lock_guard<std::mutex> lk(impl_->cb_mutex);
    impl_->bitrate_cb = std::move(cb);
}

bool WebrtcPusher::handle_offer(const std::string &offer_sdp){
    std::lock_guard<std::mutex> lk(impl_->pc_mutex);
    if(impl_->stopped) return false;

    H264Choice pick = pick_h264_from_offer(offer_sdp, impl_->cfg.h264_profile);
    if(!pick.ok){
        std::cerr << "[rtc] offer has no compatible H264 "
                  << impl_->cfg.h264_profile << " profile\n";
        return false;
    }
    std::cerr << "[rtc] offer: video mid=" << pick.mid
              << " -> H264 pt=" << pick.pt << " fmtp=" << pick.fmtp << "\n";

    /* 对端重连/刷新会再发 offer: 直接拆掉旧会话重建 */
    impl_->teardown_locked();
    if(!impl_->create_pc()) return false;

    try {
        impl_->pc->setRemoteDescription(
            rtc::Description(offer_sdp, rtc::Description::Type::Offer));
    }catch(const std::exception &e){
        std::cerr << "[rtc] setRemoteDescription(offer) error: " << e.what() << "\n";
        impl_->teardown_locked();
        return false;
    }

    if(!impl_->attach_track_locked(pick.mid, pick.pt, pick.fmtp)){
        impl_->teardown_locked();
        return false;
    }
    impl_->create_status_channel_locked();
    impl_->ensure_status_thread();

    try {
        /* 远端是 offer -> 这里自动生成 answer, 经 onLocalDescription 回调送出 */
        impl_->pc->setLocalDescription();
    }catch(const std::exception &e){
        std::cerr << "[rtc] setLocalDescription error: " << e.what() << "\n";
        impl_->teardown_locked();
        return false;
    }
    return true;
}

bool WebrtcPusher::start(){
    std::lock_guard<std::mutex> lk(impl_->pc_mutex);
    if(impl_->stopped || impl_->pc) return false;

    if(!impl_->create_pc()) return false;
    if(!impl_->attach_track_locked("video", impl_->cfg.payload_type,
                                   profile_fmtp(impl_->cfg.h264_profile)))
        { impl_->teardown_locked(); return false; }

    impl_->create_status_channel_locked();
    impl_->ensure_status_thread();

    try {
        impl_->pending = true;
        impl_->pc->setLocalDescription(rtc::Description::Type::Offer);
    }catch(const std::exception &e){
        impl_->pending = false;
        std::cerr << "[rtc] setLocalDescription(offer) error: " << e.what() << "\n";
        impl_->teardown_locked();
        return false;
    }
    return true;
}

bool WebrtcPusher::set_remote_answer(const std::string &answer_sdp){
    std::lock_guard<std::mutex> lk(impl_->pc_mutex);
    if(!impl_->pc){ return false; }
    try{
        impl_->pc->setRemoteDescription(
            rtc::Description(answer_sdp, rtc::Description::Type::Answer));
        impl_->pending = false;
        return true;
    }catch(const std::exception &e){
        std::cerr << "[rtc] setRemoteAnswer error: " << e.what() << std::endl;
        return false;
    }
}

bool WebrtcPusher::add_remote_candidate(const std::string &candidate_json){
    std::shared_ptr<rtc::PeerConnection> p;
    {
        /* 不锁 pc_mutex 全程: addRemoteCandidate 可能触发 ICE 内部回调 */
        std::lock_guard<std::mutex> lk(impl_->pc_mutex);
        p = impl_->pc;
    }
    if(!p) return false;
    try{
        auto j = json::parse(candidate_json);
        std::string cand_str = j.value("candidate", "");
        std::string mid      = j.value("sdpMid", "");
        if(cand_str.empty()) return false;
        p->addRemoteCandidate(rtc::Candidate(cand_str, mid));
        return true;
    }catch(const std::exception &e){
        std::cerr << "[rtc] addRemoteCandidate error: " << e.what() << std::endl;
        return false;
    }
}

bool WebrtcPusher::is_connected() const{ return impl_->connected; }
bool WebrtcPusher::has_session() const{
    std::lock_guard<std::mutex> lk(impl_->pc_mutex);
    return impl_->pc != nullptr;
}
bool WebrtcPusher::session_pending() const{ return impl_->pending; }

void WebrtcPusher::reset(){
    std::lock_guard<std::mutex> lk(impl_->pc_mutex);
    if(impl_->pc) std::cerr << "[rtc] reset session\n";
    impl_->teardown_locked();
}

void WebrtcPusher::feed_nalu(const uint8_t *data, size_t len, bool is_keyframe){
    if(impl_->stopped || !data || len == 0) return;
    const auto nalus = split_annexb(data, len);
    bool has_sps = false;
    bool has_pps = false;
    bool has_idr = false;
    for(const auto &nalu : nalus){
        has_sps = has_sps || nalu.type == 7;
        has_pps = has_pps || nalu.type == 8;
        has_idr = has_idr || nalu.type == 5;
    }

    std::shared_ptr<rtc::Track> t;
    std::shared_ptr<rtc::RtpPacketizationConfig> rtp;
    uint32_t frame_idx = 0;
    std::vector<std::byte> buf;
    {
        std::lock_guard<std::mutex> lk(impl_->pc_mutex);
        for(const auto &nalu : nalus){
            if(nalu.offset + nalu.size > len) continue;
            auto begin = reinterpret_cast<const std::byte*>(data + nalu.offset);
            if(nalu.type == 7)
                impl_->cached_sps.assign(begin, begin + nalu.size);
            else if(nalu.type == 8)
                impl_->cached_pps.assign(begin, begin + nalu.size);
        }

        t = impl_->track;
        rtp = impl_->rtp_cfg;
        if(!t || !rtp || !impl_->connected) return;
        frame_idx = impl_->frame_idx.fetch_add(1);

        const bool key = is_keyframe || has_idr;
        if(key && (!has_sps || !has_pps) &&
           !impl_->cached_sps.empty() && !impl_->cached_pps.empty()){
            buf.reserve(impl_->cached_sps.size() + impl_->cached_pps.size() + len);
            buf.insert(buf.end(), impl_->cached_sps.begin(), impl_->cached_sps.end());
            buf.insert(buf.end(), impl_->cached_pps.begin(), impl_->cached_pps.end());
        }else{
            buf.reserve(len);
        }
        auto payload = reinterpret_cast<const std::byte*>(data);
        buf.insert(buf.end(), payload, payload + len);
    }

    /* RTP 时间戳: 90kHz, 按实际帧率步进 */
    uint32_t step = (impl_->cfg.fps > 0)
                    ? (uint32_t)(impl_->cfg.video_clock / impl_->cfg.fps) : 3600;
    rtp->timestamp = frame_idx * step;
    try {
        t->send(buf);
        const uint64_t sent = impl_->sent_packets.fetch_add(1) + 1;
        if(sent <= 3 || ((is_keyframe || has_idr) && sent <= 120) || sent % 240 == 0){
            std::cerr << "[rtc] sent H264 packet=" << sent
                      << " bytes=" << buf.size()
                      << " nal=" << (nalus.empty() ? -1 : nalus.front().type)
                      << " key=" << ((is_keyframe || has_idr) ? "yes" : "no")
                      << " sps=" << (impl_->cached_sps.empty() ? "no" : "yes")
                      << " pps=" << (impl_->cached_pps.empty() ? "no" : "yes")
                      << "\n";
        }
    }catch(const std::exception &e){
        std::cerr << "[rtc] send error: " << e.what() << std::endl;
    }
}

void WebrtcPusher::stop(){
    if(impl_->stopped.exchange(true)) return;
    impl_->monitor_cv.notify_all();
    if(impl_->monitor_th.joinable()) impl_->monitor_th.join();
    std::lock_guard<std::mutex> lk(impl_->pc_mutex);
    impl_->teardown_locked();
}

} // namespace webrtcpush
