#include "rtmp_sink.h"

#include <android/log.h>
#include <cstring>
#include <algorithm>

#define TAG "LckRtmpSink"
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, TAG, __VA_ARGS__)
#define LOGW(...) __android_log_print(ANDROID_LOG_WARN, TAG, __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, TAG, __VA_ARGS__)

RtmpSink::RtmpSink() {}

RtmpSink::~RtmpSink() {
    if (isOpen) {
        Close();
    }
}

void RtmpSink::SetRtmpUrl(const std::string& url) {
    rtmpUrl = url;
}

bool RtmpSink::Open(uint32_t width, uint32_t height, uint32_t framerate,
                     uint32_t sampleRate, uint32_t numChannels) {
    if (rtmpUrl.empty()) {
        LOGE("RTMP URL not set");
        return false;
    }

    storedSampleRate = sampleRate;
    storedNumChannels = numChannels;

    if (!rtmpClient.Connect(rtmpUrl)) {
        LOGE("Failed to connect RTMP");
        return false;
    }

    isOpen = true;
    videoHeaderSent = false;
    audioHeaderSent = false;
    LOGI("RTMP sink opened: %dx%d@%dfps, %dHz %dch",
         width, height, framerate, sampleRate, numChannels);
    return true;
}

void RtmpSink::OnVideoFormatReady(const uint8_t* extraData, uint32_t extraDataSize) {
    if (!isOpen) return;

    if (extraData && extraDataSize > 0) {
        // Check if already AVCC format (starts with version byte 0x01)
        if (extraDataSize > 4 && extraData[0] == 0x01) {
            if (rtmpClient.SendAvcSequenceHeader(extraData, extraDataSize)) {
                videoHeaderSent = true;
                LOGI("Sent AVC sequence header (AVCC, %u bytes)", extraDataSize);
            }
        } else {
            // Annex-B format - extract and convert
            TryExtractAndSendSequenceHeader(extraData, extraDataSize);
        }
    }
}

void RtmpSink::OnAudioFormatReady(uint32_t sampleRate, uint32_t numChannels) {
    if (!isOpen) return;

    storedSampleRate = sampleRate;
    storedNumChannels = numChannels;

    if (rtmpClient.SendAacSequenceHeader(sampleRate, numChannels)) {
        audioHeaderSent = true;
        LOGI("Sent AAC sequence header (%dHz, %dch)", sampleRate, numChannels);
    }
}

void RtmpSink::SendVideoPacket(const uint8_t* data, uint32_t size,
                                int64_t timestampMs, bool isKeyframe) {
    if (!isOpen || !rtmpClient.IsConnected()) return;

    // If we haven't sent the video sequence header yet and this is a keyframe,
    // try to extract SPS/PPS from it
    if (!videoHeaderSent && isKeyframe) {
        TryExtractAndSendSequenceHeader(data, size);
    }

    if (!videoHeaderSent) return;

    // Send AAC sequence header on first video packet if not sent yet
    if (!audioHeaderSent) {
        if (rtmpClient.SendAacSequenceHeader(storedSampleRate, storedNumChannels)) {
            audioHeaderSent = true;
            LOGI("Sent AAC sequence header (deferred, %dHz, %dch)",
                 storedSampleRate, storedNumChannels);
        }
    }

    uint32_t ts = static_cast<uint32_t>(std::max<int64_t>(timestampMs, 0));

    // Convert Annex-B to AVCC for RTMP/FLV
    std::vector<uint8_t> avccData = ConvertAnnexBToAvcc(data, size);
    if (!avccData.empty()) {
        rtmpClient.SendVideoPacket(avccData.data(), static_cast<uint32_t>(avccData.size()),
                                   ts, isKeyframe);
    }
}

void RtmpSink::SendAudioPacket(const uint8_t* data, uint32_t size, int64_t timestampMs) {
    if (!isOpen || !rtmpClient.IsConnected()) return;
    if (!audioHeaderSent || !videoHeaderSent) return;

    uint32_t ts = static_cast<uint32_t>(std::max<int64_t>(timestampMs, 0));
    rtmpClient.SendAudioPacket(data, size, ts);
}

void RtmpSink::Close() {
    if (isOpen) {
        rtmpClient.Disconnect();
        isOpen = false;
        videoHeaderSent = false;
        audioHeaderSent = false;
        LOGI("RTMP sink closed");
    }
}

bool RtmpSink::IsOpen() const {
    return isOpen;
}

bool RtmpSink::TryExtractAndSendSequenceHeader(const uint8_t* data, uint32_t size) {
    // Parse Annex-B bitstream to find SPS and PPS NALUs
    const uint8_t* sps = nullptr;
    uint32_t spsSize = 0;
    const uint8_t* pps = nullptr;
    uint32_t ppsSize = 0;

    const uint8_t* end = data + size;

    auto findStartCode = [](const uint8_t* p, const uint8_t* end) -> const uint8_t* {
        while (p + 3 <= end) {
            if (p[0] == 0 && p[1] == 0) {
                if (p[2] == 1) return p + 3;
                if (p + 3 < end && p[2] == 0 && p[3] == 1) return p + 4;
            }
            p++;
        }
        return nullptr;
    };

    const uint8_t* pos = findStartCode(data, end);
    while (pos && pos < end) {
        uint8_t currentNaluType = pos[0] & 0x1F;
        const uint8_t* currentNaluStart = pos;

        const uint8_t* nextStart = findStartCode(pos, end);
        const uint8_t* naluEnd;
        if (nextStart) {
            naluEnd = nextStart - 3;
            if (naluEnd > data && *(naluEnd - 1) == 0) naluEnd--;
        } else {
            naluEnd = end;
        }

        uint32_t naluSize = static_cast<uint32_t>(naluEnd - currentNaluStart);

        if (currentNaluType == 7 && !sps) { // SPS
            sps = currentNaluStart;
            spsSize = naluSize;
        } else if (currentNaluType == 8 && !pps) { // PPS
            pps = currentNaluStart;
            ppsSize = naluSize;
        }

        if (sps && pps) break;
        pos = nextStart;
    }

    if (sps && spsSize > 0 && pps && ppsSize > 0) {
        std::vector<uint8_t> avcc = BuildAvccFromAnnexB(sps, spsSize, pps, ppsSize);
        if (rtmpClient.SendAvcSequenceHeader(avcc.data(), static_cast<uint32_t>(avcc.size()))) {
            videoHeaderSent = true;
            LOGI("Sent AVC sequence header (extracted SPS=%u PPS=%u)", spsSize, ppsSize);
            return true;
        } else {
            LOGE("SendAvcSequenceHeader failed (SPS=%u PPS=%u)", spsSize, ppsSize);
        }
    }

    return false;
}

std::vector<uint8_t> RtmpSink::ConvertAnnexBToAvcc(const uint8_t* data, uint32_t size) {
    std::vector<uint8_t> result;
    result.reserve(size);

    auto findStartCode = [](const uint8_t* p, const uint8_t* end, int& startCodeLen) -> const uint8_t* {
        while (p + 3 <= end) {
            if (p[0] == 0 && p[1] == 0) {
                if (p + 3 < end && p[2] == 0 && p[3] == 1) {
                    startCodeLen = 4;
                    return p;
                }
                if (p[2] == 1) {
                    startCodeLen = 3;
                    return p;
                }
            }
            p++;
        }
        return nullptr;
    };

    const uint8_t* pos = data;
    const uint8_t* end = data + size;

    int startCodeLen = 0;
    const uint8_t* startCode = findStartCode(pos, end, startCodeLen);
    if (!startCode) {
        // No start codes found - pass through
        result.insert(result.end(), data, data + size);
        return result;
    }

    while (startCode) {
        const uint8_t* naluStart = startCode + startCodeLen;
        if (naluStart >= end) break;

        int nextStartCodeLen = 0;
        const uint8_t* nextStartCode = findStartCode(naluStart, end, nextStartCodeLen);
        uint32_t naluSize = nextStartCode
            ? static_cast<uint32_t>(nextStartCode - naluStart)
            : static_cast<uint32_t>(end - naluStart);

        if (naluSize > 0) {
            uint8_t naluType = naluStart[0] & 0x1F;
            // Skip SPS (7), PPS (8), AUD (9)
            if (naluType != 7 && naluType != 8 && naluType != 9) {
                result.push_back(static_cast<uint8_t>(naluSize >> 24));
                result.push_back(static_cast<uint8_t>(naluSize >> 16));
                result.push_back(static_cast<uint8_t>(naluSize >> 8));
                result.push_back(static_cast<uint8_t>(naluSize & 0xFF));
                result.insert(result.end(), naluStart, naluStart + naluSize);
            }
        }

        startCode = nextStartCode;
        startCodeLen = nextStartCodeLen;
    }

    return result;
}

std::vector<uint8_t> RtmpSink::BuildAvccFromAnnexB(const uint8_t* sps, uint32_t spsSize,
                                                     const uint8_t* pps, uint32_t ppsSize) {
    // AVCDecoderConfigurationRecord
    std::vector<uint8_t> record;
    record.reserve(11 + spsSize + ppsSize);

    record.push_back(0x01); // configurationVersion
    record.push_back(spsSize > 1 ? sps[1] : 0x42); // AVCProfileIndication
    record.push_back(spsSize > 2 ? sps[2] : 0x00); // profile_compatibility
    record.push_back(spsSize > 3 ? sps[3] : 0x1E); // AVCLevelIndication
    record.push_back(0xFF); // lengthSizeMinusOne = 3 (4 bytes)
    record.push_back(0xE1); // numOfSequenceParameterSets = 1

    // SPS length (big-endian)
    record.push_back(static_cast<uint8_t>(spsSize >> 8));
    record.push_back(static_cast<uint8_t>(spsSize & 0xFF));
    record.insert(record.end(), sps, sps + spsSize);

    record.push_back(0x01); // numOfPictureParameterSets = 1

    // PPS length (big-endian)
    record.push_back(static_cast<uint8_t>(ppsSize >> 8));
    record.push_back(static_cast<uint8_t>(ppsSize & 0xFF));
    record.insert(record.end(), pps, pps + ppsSize);

    return record;
}
