#include "clip_recorder.h"
#include "faststart.h"

#include <media/NdkMediaMuxer.h>
#include <media/NdkMediaFormat.h>
#include <media/NdkMediaCodec.h>
#include <android/log.h>

#include <algorithm>
#include <fcntl.h>
#include <unistd.h>
#include <cstdio>

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

void ClipRecorder::Configure(int width, int height,
                             uint32_t audioSampleRate, uint32_t audioChannels_,
                             int audioBitrate_) {
    std::lock_guard<std::mutex> lock(mutex);
    videoWidth = width;
    videoHeight = height;
    sampleRate = audioSampleRate;
    audioChannels = audioChannels_;
    audioBitrate = audioBitrate_;
    LOGI("Configured: %dx%d, audio %uHz %uch %dbps", width, height,
         audioSampleRate, audioChannels_, audioBitrate_);
}

void ClipRecorder::SetVideoFormat(const uint8_t* codecConfig, uint32_t size) {
    std::lock_guard<std::mutex> lock(mutex);
    videoCodecConfig.assign(codecConfig, codecConfig + size);
    LOGI("Video codec config set: %u bytes", size);
}

void ClipRecorder::SetAudioFormat(const uint8_t* codecConfig, uint32_t size) {
    std::lock_guard<std::mutex> lock(mutex);
    audioCodecConfig.assign(codecConfig, codecConfig + size);
    LOGI("Audio codec config set: %u bytes", size);
}

void ClipRecorder::Start() {
    std::lock_guard<std::mutex> lock(mutex);
    gopBuffer.clear();
    currentGop = GopBuffer();
    audioBuffer.clear();
    active = true;
    LOGI("Started");
}

void ClipRecorder::Stop() {
    std::lock_guard<std::mutex> lock(mutex);
    active = false;
    gopBuffer.clear();
    currentGop = GopBuffer();
    audioBuffer.clear();
    videoCodecConfig.clear();
    audioCodecConfig.clear();
    LOGI("Stopped");
}

void ClipRecorder::FeedVideoPacket(const uint8_t* data, uint32_t size,
                                   int64_t timestampUs, bool isKeyframe) {
    std::lock_guard<std::mutex> lock(mutex);
    if (!active) return;

    if (isKeyframe) {
        // Finalize current GOP and push to buffer
        if (!currentGop.samples.empty()) {
            gopBuffer.push_back(std::move(currentGop));
        }
        currentGop = GopBuffer();
        currentGop.startTimeUs = timestampUs;
    }

    VideoSample sample;
    sample.data.assign(data, data + size);
    sample.timestampUs = timestampUs;
    sample.isKeyframe = isKeyframe;
    currentGop.samples.push_back(std::move(sample));

    TrimBuffers();
}

void ClipRecorder::FeedAudioPacket(const uint8_t* data, uint32_t size,
                                   int64_t timestampUs) {
    std::lock_guard<std::mutex> lock(mutex);
    if (!active) return;

    AudioSample sample;
    sample.data.assign(data, data + size);
    sample.timestampUs = timestampUs;
    audioBuffer.push_back(std::move(sample));

    TrimBuffers();
}

void ClipRecorder::TrimBuffers() {
    // Trim oldest GOPs to stay under MAX_BUFFER_DURATION_US
    if (gopBuffer.size() < 2) return;

    int64_t newestTs = 0;
    if (!currentGop.samples.empty()) {
        newestTs = currentGop.samples.back().timestampUs;
    } else if (!gopBuffer.empty() && !gopBuffer.back().samples.empty()) {
        newestTs = gopBuffer.back().samples.back().timestampUs;
    }
    if (newestTs == 0) return;

    while (!gopBuffer.empty()) {
        int64_t oldestTs = gopBuffer.front().startTimeUs;
        if (newestTs - oldestTs > MAX_BUFFER_DURATION_US && gopBuffer.size() > 1) {
            gopBuffer.pop_front();
        } else {
            break;
        }
    }

    // Trim audio samples older than oldest remaining video
    if (!gopBuffer.empty()) {
        int64_t videoStart = gopBuffer.front().startTimeUs;
        while (!audioBuffer.empty() && audioBuffer.front().timestampUs < videoStart) {
            audioBuffer.pop_front();
        }
    }
}

bool ClipRecorder::FlushClip(const std::string& outputDir) {
    // Copy data under lock
    std::vector<GopBuffer> gopsCopy;
    std::vector<AudioSample> audioCopy;
    std::vector<uint8_t> videoCsdCopy;
    std::vector<uint8_t> audioCsdCopy;
    int w, h;
    uint32_t sr, ch;
    int abr;

    {
        std::lock_guard<std::mutex> lock(mutex);
        if (!active) return false;
        if (gopBuffer.empty()) {
            LOGW("FlushClip: no complete GOPs");
            return false;
        }
        if (videoCodecConfig.empty()) {
            LOGW("FlushClip: no video codec config");
            return false;
        }

        gopsCopy.assign(gopBuffer.begin(), gopBuffer.end());
        audioCopy.assign(audioBuffer.begin(), audioBuffer.end());
        videoCsdCopy = videoCodecConfig;
        audioCsdCopy = audioCodecConfig;
        w = videoWidth;
        h = videoHeight;
        sr = sampleRate;
        ch = audioChannels;
        abr = audioBitrate;
    }

    // Mux to temp file
    std::string tempPath = outputDir + "/preview_clip_tmp.mp4";
    std::string finalPath = outputDir + "/preview_clip.mp4";

    int fd = open(tempPath.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (fd < 0) {
        LOGE("FlushClip: failed to open temp file: %s", tempPath.c_str());
        return false;
    }

    AMediaMuxer* muxer = AMediaMuxer_new(fd, AMEDIAMUXER_OUTPUT_FORMAT_MPEG_4);
    if (!muxer) {
        LOGE("FlushClip: failed to create muxer");
        close(fd);
        return false;
    }

    // Video track — split Annex-B SPS+PPS into separate csd-0 (SPS) and csd-1 (PPS)
    // MediaCodec CODEC_CONFIG output: [00 00 00 01 SPS] [00 00 00 01 PPS]
    // AMediaMuxer expects csd-0 = SPS NAL (with start code), csd-1 = PPS NAL (with start code)
    const uint8_t* csd = videoCsdCopy.data();
    size_t csdSize = videoCsdCopy.size();
    size_t spsLen = csdSize;
    size_t ppsOffset = 0;
    size_t ppsLen = 0;

    // Find the second start code (00 00 00 01) to split SPS from PPS
    for (size_t i = 4; i + 3 < csdSize; i++) {
        if (csd[i] == 0x00 && csd[i+1] == 0x00 && csd[i+2] == 0x00 && csd[i+3] == 0x01) {
            spsLen = i;
            ppsOffset = i;
            ppsLen = csdSize - i;
            break;
        }
    }

    LOGI("FlushClip: csd total=%zu, SPS=%zu bytes, PPS=%zu bytes", csdSize, spsLen, ppsLen);

    AMediaFormat* videoFormat = AMediaFormat_new();
    AMediaFormat_setString(videoFormat, AMEDIAFORMAT_KEY_MIME, "video/avc");
    AMediaFormat_setInt32(videoFormat, AMEDIAFORMAT_KEY_WIDTH, w);
    AMediaFormat_setInt32(videoFormat, AMEDIAFORMAT_KEY_HEIGHT, h);
    AMediaFormat_setBuffer(videoFormat, "csd-0", csd, spsLen);
    if (ppsLen > 0) {
        AMediaFormat_setBuffer(videoFormat, "csd-1", csd + ppsOffset, ppsLen);
    }

    ssize_t videoTrack = AMediaMuxer_addTrack(muxer, videoFormat);
    AMediaFormat_delete(videoFormat);

    if (videoTrack < 0) {
        LOGE("FlushClip: failed to add video track");
        AMediaMuxer_delete(muxer);
        close(fd);
        return false;
    }

    // Audio track (optional — may not have audio config)
    ssize_t audioTrack = -1;
    if (!audioCsdCopy.empty()) {
        AMediaFormat* audioFormat = AMediaFormat_new();
        AMediaFormat_setString(audioFormat, AMEDIAFORMAT_KEY_MIME, "audio/mp4a-latm");
        AMediaFormat_setInt32(audioFormat, AMEDIAFORMAT_KEY_SAMPLE_RATE, sr);
        AMediaFormat_setInt32(audioFormat, AMEDIAFORMAT_KEY_CHANNEL_COUNT, ch);
        AMediaFormat_setInt32(audioFormat, AMEDIAFORMAT_KEY_BIT_RATE, abr);
        AMediaFormat_setBuffer(audioFormat, "csd-0", audioCsdCopy.data(), audioCsdCopy.size());

        audioTrack = AMediaMuxer_addTrack(muxer, audioFormat);
        AMediaFormat_delete(audioFormat);

        if (audioTrack < 0) {
            LOGW("FlushClip: failed to add audio track, continuing video-only");
            audioTrack = -1;
        }
    }

    media_status_t status = AMediaMuxer_start(muxer);
    if (status != AMEDIA_OK) {
        LOGE("FlushClip: failed to start muxer: %d", status);
        AMediaMuxer_delete(muxer);
        close(fd);
        return false;
    }

    // Compute base timestamp for zero-based output
    int64_t baseTs = gopsCopy.front().startTimeUs;

    // Write video samples
    int videoSamplesWritten = 0;
    for (const auto& gop : gopsCopy) {
        for (const auto& sample : gop.samples) {
            AMediaCodecBufferInfo info;
            info.offset = 0;
            info.size = static_cast<int32_t>(sample.data.size());
            info.presentationTimeUs = sample.timestampUs - baseTs;
            info.flags = sample.isKeyframe ? AMEDIACODEC_BUFFER_FLAG_KEY_FRAME : 0;

            AMediaMuxer_writeSampleData(muxer, videoTrack,
                                        sample.data.data(), &info);
            videoSamplesWritten++;
        }
    }

    // Write audio samples
    int audioSamplesWritten = 0;
    if (audioTrack >= 0) {
        for (const auto& sample : audioCopy) {
            // Only include audio within the video time range
            int64_t relTs = sample.timestampUs - baseTs;
            if (relTs < 0) continue;

            AMediaCodecBufferInfo info;
            info.offset = 0;
            info.size = static_cast<int32_t>(sample.data.size());
            info.presentationTimeUs = relTs;
            info.flags = 0;

            AMediaMuxer_writeSampleData(muxer, audioTrack,
                                        sample.data.data(), &info);
            audioSamplesWritten++;
        }
    }

    AMediaMuxer_stop(muxer);
    AMediaMuxer_delete(muxer);
    close(fd);

    LOGI("FlushClip: muxed %d video + %d audio samples to temp file",
         videoSamplesWritten, audioSamplesWritten);

    // Apply faststart (move moov atom to front)
    if (!MoovFastStart(tempPath, finalPath)) {
        LOGW("FlushClip: faststart failed, using non-optimized file");
        rename(tempPath.c_str(), finalPath.c_str());
    } else {
        unlink(tempPath.c_str());
    }

    LOGI("FlushClip: clip ready at %s", finalPath.c_str());

    if (clipReadyCallback) {
        clipReadyCallback(finalPath);
    }

    return true;
}

void ClipRecorder::SetClipReadyCallback(ClipReadyCallback cb) {
    std::lock_guard<std::mutex> lock(mutex);
    clipReadyCallback = std::move(cb);
}