/*
 * Copyright (c) 2026 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "AudioVividFlvMuxerManager.h"
#include "AudioConfig.h"
#include <multimedia/player_framework/native_avbuffer.h>
#include <multimedia/player_framework/native_avcodec_base.h>
#include <multimedia/player_framework/native_audio_vivid.h>
#include <cstring>
#include <chrono>

namespace {
using AudioConfig::SAMPLE_RATE;
using AudioConfig::CHANNEL_COUNT;
using AudioConfig::CHANNEL_LAYOUT;
using AudioConfig::BITRATE;
using AudioConfig::FRAME_SIZE;
using AudioConfig::BYTES_PER_INPUT_FRAME;
constexpr int32_t PCM_BIT_DEPTH = 24; // 24-bit PCM采样位深
constexpr int32_t MAX_METADATA_SIZE = 10 * 1024 * 1024; // 元数据大小上限10MB
constexpr int32_t US_PER_SECOND = 1000000; // 每秒对应的微秒数，用于PTS计算
constexpr int32_t MAX_PUSH_RETRY_COUNT = 100; // 推入编码器缓冲区最大重试次数
constexpr int32_t POLL_INTERVAL_MS = 10; // 编码器轮询等待间隔(毫秒)
} // namespace

AudioVividFlvMuxerManager::AudioVividFlvMuxerManager() {}

AudioVividFlvMuxerManager::~AudioVividFlvMuxerManager()
{
    Release();
}

int32_t AudioVividFlvMuxerManager::Initialize(int32_t fd)
{
    int32_t ret = pcmReader_.Open("/data/storage/el2/base/haps/entry/files/resfile/2ch_2obj_48k_24bit.pcm",
        SAMPLE_RATE, CHANNEL_COUNT, PCM_BIT_DEPTH);
    if (ret != 0) {
        AVCODEC_SAMPLE_LOGE("Failed to open PCM file");
        return 1; // return 1 is PCM file open failed
    }
    ret = encoder_.Create();
    if (ret != 0) {
        AVCODEC_SAMPLE_LOGE("Failed to create encoder");
        return 2; // return 2 is encoder create failed
    }
    ret = encoder_.Config(SAMPLE_RATE, CHANNEL_COUNT, CHANNEL_LAYOUT, BITRATE);
    if (ret != 0) {
        AVCODEC_SAMPLE_LOGE("Failed to config encoder");
        return 3; // return 3 is encoder config failed
    }

    muxer_ = OH_AVMuxer_Create(fd, AV_OUTPUT_FORMAT_FLV);
    if (muxer_ == nullptr) {
        AVCODEC_SAMPLE_LOGE("Failed to create OH_AVMuxer");
        return 4; // return 4 is OH_AVMuxer create failed
    }

    OH_AVFormat *format = OH_AVFormat_CreateAudioFormat(OH_AVCODEC_MIMETYPE_AUDIO_VIVID, SAMPLE_RATE, CHANNEL_COUNT);
    if (format == nullptr) {
        AVCODEC_SAMPLE_LOGE("Failed to create audio format");
        OH_AVMuxer_Destroy(muxer_);
        muxer_ = nullptr;
        return 5; // return 5 is audio format create failed
    }
    OH_AVFormat_SetIntValue(format, OH_MD_KEY_BITRATE, BITRATE);

    ret = OH_AVMuxer_AddTrack(muxer_, &trackIndex_, format);
    OH_AVFormat_Destroy(format);
    if (ret != AV_ERR_OK || trackIndex_ < 0) {
        AVCODEC_SAMPLE_LOGE("Failed to add audio track, ret=%{public}d", ret);
        OH_AVMuxer_Destroy(muxer_);
        muxer_ = nullptr;
        return 6; // return 6 is add audio track failed
    }

    ret = OH_AVMuxer_Start(muxer_);
    if (ret != AV_ERR_OK) {
        AVCODEC_SAMPLE_LOGE("Failed to start muxer, ret=%{public}d", ret);
        OH_AVMuxer_Destroy(muxer_);
        muxer_ = nullptr;
        return 7; // return 7 is muxer start failed
    }

    GenerateInitialMetadata();
    AVCODEC_SAMPLE_LOGI("AudioVividFlvMuxerManager initialized, trackIndex=%{public}d", trackIndex_);
    return 0;
}

int32_t AudioVividFlvMuxerManager::Start()
{
    if (isRunning_.load()) {
        AVCODEC_SAMPLE_LOGW("Muxing is already running");
        return 0;
    }
    shouldStop_.store(false);
    presentationTimeUs_.store(0);
    int32_t ret = encoder_.Start();
    if (ret != 0) {
        AVCODEC_SAMPLE_LOGE("Failed to start encoder");
        return -1;
    }
    encoderInputThread_ = std::thread(&AudioVividFlvMuxerManager::EncoderInputThread, this);
    encoderOutputThread_ = std::thread(&AudioVividFlvMuxerManager::EncoderOutputThread, this);
    isRunning_.store(true);
    AVCODEC_SAMPLE_LOGI("AudioVividFlvMuxerManager started");
    return 0;
}

int32_t AudioVividFlvMuxerManager::Stop()
{
    if (isRunning_.load()) {
        shouldStop_.store(true);
        encoder_.GetContext()->inputCond.notify_all();
        encoder_.GetContext()->outputCond.notify_all();
        if (encoderInputThread_.joinable()) {
            encoderInputThread_.join();
        }
        if (encoderOutputThread_.joinable()) {
            encoderOutputThread_.join();
        }
        encoder_.Stop();
        isRunning_.store(false);
    }
    if (muxer_) {
        OH_AVMuxer_Stop(muxer_);
        OH_AVMuxer_Destroy(muxer_);
        muxer_ = nullptr;
    }
    AVCODEC_SAMPLE_LOGI("AudioVividFlvMuxerManager stopped");
    return 0;
}

int32_t AudioVividFlvMuxerManager::Release()
{
    Stop();
    encoder_.Release();
    pcmReader_.Close();
    if (currentMetadata_) {
        delete[] currentMetadata_;
        currentMetadata_ = nullptr;
    }
    currentMetadataSize_ = 0;
    AVCODEC_SAMPLE_LOGI("AudioVividFlvMuxerManager released");
    return 0;
}

void AudioVividFlvMuxerManager::GenerateInitialMetadata()
{
    vividWrapper_.Initialize();
    vividWrapper_.UpdateObjectPosition(0, 0.3f, 0.0f, 0.5f);
    vividWrapper_.UpdateObjectPosition(1, -0.3f, 0.0f, 0.5f);
    vividWrapper_.UpdateObjectGain(0, 1.0f);
    vividWrapper_.UpdateObjectGain(1, 1.0f);
    std::vector<uint8_t> metadata = vividWrapper_.GetMetadata(true);
    {
        std::lock_guard<std::mutex> lock(metadataMutex_);
        if (currentMetadata_) {
            delete[] currentMetadata_;
        }
        currentMetadata_ = new uint8_t[metadata.size()];
        memcpy(currentMetadata_, metadata.data(), metadata.size());
        currentMetadataSize_ = metadata.size();
    }
    encoder_.UpdateMetadata(currentMetadata_, currentMetadataSize_);
}

void AudioVividFlvMuxerManager::UpdateMetadata(uint8_t *metadata, int32_t metadataSize)
{
    std::lock_guard<std::mutex> lock(metadataMutex_);
    if (currentMetadata_) {
        delete[] currentMetadata_;
    }
    if (metadataSize > MAX_METADATA_SIZE) {
        return;
    }
    currentMetadata_ = new uint8_t[metadataSize];
    memcpy(currentMetadata_, metadata, metadataSize);
    currentMetadataSize_ = metadataSize;
    encoder_.UpdateMetadata(metadata, metadataSize);
}

void AudioVividFlvMuxerManager::EncoderInputThread()
{
    uint8_t *pcmBuffer = new uint8_t[BYTES_PER_INPUT_FRAME];
    uint8_t *encoderBuffer = new uint8_t[BYTES_PER_INPUT_FRAME];
    while (!shouldStop_.load()) {
        int32_t bytesRead = pcmReader_.Read(pcmBuffer, BYTES_PER_INPUT_FRAME);
        if (bytesRead <= 0) {
            if (pcmReader_.IsEOF()) {
                AVCODEC_SAMPLE_LOGI("PCM file read complete, notifying EOS");
                encoder_.NotifyEndOfStream();
                break;
            }
            continue;
        }
        memset(encoderBuffer, 0, BYTES_PER_INPUT_FRAME);
        memcpy(encoderBuffer, pcmBuffer, bytesRead);
        int64_t pts = presentationTimeUs_.fetch_add(FRAME_SIZE * US_PER_SECOND / SAMPLE_RATE);
        uint8_t *metadata = nullptr;
        int32_t metadataSize = 0;
        {
            std::lock_guard<std::mutex> lock(metadataMutex_);
            if (currentMetadata_ && currentMetadataSize_ > 0) {
                metadata = currentMetadata_;
                metadataSize = currentMetadataSize_;
            }
        }
        int32_t retryCount = 0;
        while (!shouldStop_.load()) {
            int32_t ret = encoder_.PushInputBuffer(encoderBuffer, BYTES_PER_INPUT_FRAME, metadata, metadataSize, pts);
            if (ret == 0) {
                break;
            }
            retryCount++;
            if (retryCount > MAX_PUSH_RETRY_COUNT) {
                AVCODEC_SAMPLE_LOGE("Failed to push input buffer after retries");
                break;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(POLL_INTERVAL_MS));
        }
    }
    delete[] pcmBuffer;
    delete[] encoderBuffer;
    AVCODEC_SAMPLE_LOGI("EncoderInputThread exited");
}

void AudioVividFlvMuxerManager::WriteEncodedSample(AudioVividEncoderOutputData *outputData)
{
    if (outputData->encodedData == nullptr || outputData->encodedSize <= 0) {
        return;
    }
    std::lock_guard<std::mutex> lock(writeMutex_);
    OH_AVBuffer *buffer = OH_AVBuffer_Create(outputData->encodedSize);
    if (buffer == nullptr) {
        return;
    }
    memcpy(OH_AVBuffer_GetAddr(buffer), outputData->encodedData, outputData->encodedSize);
    OH_AVCodecBufferAttr attr;
    attr.pts = outputData->presentationTimeUs;
    attr.size = outputData->encodedSize;
    attr.offset = 0;
    attr.flags = AVCODEC_BUFFER_FLAGS_NONE;
    OH_AVBuffer_SetBufferAttr(buffer, &attr);
    int32_t ret = OH_AVMuxer_WriteSampleBuffer(muxer_, trackIndex_, buffer);
    if (ret != AV_ERR_OK) {
        AVCODEC_SAMPLE_LOGE("WriteSampleBuffer failed, ret=%{public}d", ret);
    }
    OH_AVBuffer_Destroy(buffer);
}

void AudioVividFlvMuxerManager::EncoderOutputThread()
{
    while (!shouldStop_.load()) {
        AudioVividEncoderOutputData *outputData = encoder_.GetOutputBuffer();
        if (outputData == nullptr) {
            std::this_thread::sleep_for(std::chrono::milliseconds(POLL_INTERVAL_MS));
            continue;
        }
        if (outputData->eos) {
            delete outputData;
            encoder_.FreeOutputBuffer();
            AVCODEC_SAMPLE_LOGI("EncoderOutputThread received EOS");
            break;
        }
        WriteEncodedSample(outputData);
        delete outputData;
        encoder_.FreeOutputBuffer();
    }
    // Drain remaining data in the output queue
    while (true) {
        AudioVividEncoderOutputData *outputData = encoder_.GetOutputBuffer();
        if (outputData == nullptr) {
            break;
        }
        if (outputData->eos) {
            delete outputData;
            encoder_.FreeOutputBuffer();
            break;
        }
        WriteEncodedSample(outputData);
        delete outputData;
        encoder_.FreeOutputBuffer();
    }
    AVCODEC_SAMPLE_LOGI("EncoderOutputThread exited");
}