#include "rtmp_client.h"

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

extern "C" {
#include <librtmp/rtmp.h>
#include <librtmp/log.h>
}

#define TAG "LckRtmpClient"
#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__)

RtmpClient::RtmpClient() {}

RtmpClient::~RtmpClient() {
    Disconnect();
}

bool RtmpClient::Connect(const std::string& rtmpUrl) {
    if (connected) {
        LOGW("Already connected, disconnecting first");
        Disconnect();
    }

    RTMP_LogSetLevel(RTMP_LOGWARNING);

    rtmpContext = RTMP_Alloc();
    if (!rtmpContext) {
        LOGE("RTMP_Alloc failed");
        return false;
    }

    RTMP_Init(rtmpContext);

    // RTMP_SetupURL needs a mutable char*
    std::vector<char> urlBuffer(rtmpUrl.begin(), rtmpUrl.end());
    urlBuffer.push_back('\0');

    if (!RTMP_SetupURL(rtmpContext, urlBuffer.data())) {
        LOGE("RTMP_SetupURL failed");
        RTMP_Free(rtmpContext);
        rtmpContext = nullptr;
        return false;
    }

    RTMP_EnableWrite(rtmpContext);

    if (!RTMP_Connect(rtmpContext, nullptr)) {
        LOGE("RTMP_Connect failed");
        RTMP_Free(rtmpContext);
        rtmpContext = nullptr;
        return false;
    }

    if (!RTMP_ConnectStream(rtmpContext, 0)) {
        LOGE("RTMP_ConnectStream failed");
        RTMP_Close(rtmpContext);
        RTMP_Free(rtmpContext);
        rtmpContext = nullptr;
        return false;
    }

    connected = true;
    LOGI("RTMP connected");
    return true;
}

void RtmpClient::Disconnect() {
    if (rtmpContext) {
        RTMP_Close(rtmpContext);
        RTMP_Free(rtmpContext);
        rtmpContext = nullptr;
        LOGI("RTMP disconnected");
    }
    connected = false;
}

bool RtmpClient::IsConnected() const {
    return connected && rtmpContext && RTMP_IsConnected(rtmpContext);
}

bool RtmpClient::SendRtmpPacket(uint8_t packetType, uint32_t timestampMs, const uint8_t* data, uint32_t size) {
    if (!IsConnected())
        return false;

    RTMPPacket pkt;
    RTMPPacket_Alloc(&pkt, size);
    pkt.m_packetType = packetType;
    pkt.m_nChannel = (packetType == RTMP_PACKET_TYPE_VIDEO) ? 0x06 : 0x07;
    pkt.m_headerType = RTMP_PACKET_SIZE_LARGE;
    pkt.m_nTimeStamp = timestampMs;
    pkt.m_hasAbsTimestamp = 1;
    pkt.m_nInfoField2 = rtmpContext->m_stream_id;
    pkt.m_nBodySize = size;
    memcpy(pkt.m_body, data, size);

    int ret = RTMP_SendPacket(rtmpContext, &pkt, 0);
    RTMPPacket_Free(&pkt);

    if (!ret) {
        LOGW("RTMP_SendPacket failed (type=%d, size=%u)", packetType, size);
        connected = false;
    }

    return ret != 0;
}

bool RtmpClient::SendAvcSequenceHeader(const uint8_t* extraData, uint32_t extraDataSize) {
    // FLV video tag: keyframe(1) + AVC(7) = 0x17, AVC sequence header = 0x00, composition time = 0
    uint32_t bodySize = 5 + extraDataSize;
    std::vector<uint8_t> body(bodySize);
    body[0] = 0x17; // keyframe + AVC
    body[1] = 0x00; // AVC sequence header
    body[2] = 0x00; // composition time
    body[3] = 0x00;
    body[4] = 0x00;
    memcpy(body.data() + 5, extraData, extraDataSize);

    return SendRtmpPacket(RTMP_PACKET_TYPE_VIDEO, 0, body.data(), bodySize);
}

void RtmpClient::BuildAudioSpecificConfig(uint8_t outConfig[2], uint32_t sampleRate, uint32_t numChannels) {
    static const uint32_t sampleRateTable[] = {
        96000, 88200, 64000, 48000, 44100, 32000,
        24000, 22050, 16000, 12000, 11025, 8000, 7350
    };

    uint8_t freqIndex = 4; // default 44100
    for (uint8_t i = 0; i < 13; ++i) {
        if (sampleRateTable[i] == sampleRate) {
            freqIndex = i;
            break;
        }
    }

    uint8_t channelConfig = static_cast<uint8_t>(numChannels < 1 ? 1 : (numChannels > 7 ? 7 : numChannels));

    // Pack: AAAAA FFFF CCCC 000
    outConfig[0] = (2 << 3) | (freqIndex >> 1);
    outConfig[1] = ((freqIndex & 1) << 7) | (channelConfig << 3);
}

bool RtmpClient::SendAacSequenceHeader(uint32_t sampleRate, uint32_t numChannels) {
    // FLV audio tag: AAC(10) + 44100(3) + 16-bit(1) + stereo(1) = 0xAF, AAC sequence header = 0x00
    uint8_t body[4];
    body[0] = 0xAF;
    body[1] = 0x00;
    BuildAudioSpecificConfig(body + 2, sampleRate, numChannels);

    return SendRtmpPacket(RTMP_PACKET_TYPE_AUDIO, 0, body, sizeof(body));
}

bool RtmpClient::SendVideoPacket(const uint8_t* data, uint32_t size, uint32_t timestampMs, bool isKeyframe) {
    uint32_t bodySize = 5 + size;
    std::vector<uint8_t> body(bodySize);
    body[0] = isKeyframe ? 0x17 : 0x27;
    body[1] = 0x01; // AVC NALU
    body[2] = 0x00; // composition time offset
    body[3] = 0x00;
    body[4] = 0x00;
    memcpy(body.data() + 5, data, size);

    return SendRtmpPacket(RTMP_PACKET_TYPE_VIDEO, timestampMs, body.data(), bodySize);
}

bool RtmpClient::SendAudioPacket(const uint8_t* data, uint32_t size, uint32_t timestampMs) {
    uint32_t bodySize = 2 + size;
    std::vector<uint8_t> body(bodySize);
    body[0] = 0xAF;
    body[1] = 0x01; // AAC raw
    memcpy(body.data() + 2, data, size);

    return SendRtmpPacket(RTMP_PACKET_TYPE_AUDIO, timestampMs, body.data(), bodySize);
}
