#include "streaming_engine.h"

#include <android/log.h>
#include <GLES3/gl3.h>
#include <GLES2/gl2ext.h>
#include <unistd.h>
#include <cstring>
#include <algorithm>

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

// Shader source for blitting OES texture (kept for legacy/direct path)
static const char* BLIT_VERTEX_SHADER = R"(#version 300 es
layout(location = 0) in vec2 aPos;
layout(location = 1) in vec2 aTexCoord;
out vec2 vTexCoord;
void main() {
    gl_Position = vec4(aPos, 0.0, 1.0);
    vTexCoord = aTexCoord;
}
)";

static const char* BLIT_FRAGMENT_SHADER = R"(#version 300 es
#extension GL_OES_EGL_image_external_essl3 : require
precision mediump float;
in vec2 vTexCoord;
out vec4 fragColor;
uniform samplerExternalOES uTexture;
void main() {
    fragColor = texture(uTexture, vTexCoord);
}
)";

// Blit FBO program: renders composed GL_TEXTURE_2D to a surface
static const char* BLIT_FBO_FRAGMENT_SHADER = R"(#version 300 es
precision mediump float;
in vec2 vTexCoord;
out vec4 fragColor;
uniform sampler2D uTexture;
void main() {
    fragColor = texture(uTexture, vTexCoord);
}
)";

static GLuint CompileShader(GLenum type, const char* source) {
    GLuint shader = glCreateShader(type);
    glShaderSource(shader, 1, &source, nullptr);
    glCompileShader(shader);

    GLint status;
    glGetShaderiv(shader, GL_COMPILE_STATUS, &status);
    if (!status) {
        char log[512];
        glGetShaderInfoLog(shader, sizeof(log), nullptr, log);
        LOGE("Shader compile error: %s", log);
        glDeleteShader(shader);
        return 0;
    }
    return shader;
}

StreamingEngine::StreamingEngine() {}

StreamingEngine::~StreamingEngine() {
    Stop();
    for (auto* sink : sinks) {
        delete sink;
    }
    sinks.clear();
}

bool StreamingEngine::Configure(int w, int h, int vBitrate, int aBitrate,
                                 int sr, int ch, int kfi) {
    width = w;
    height = h;
    videoBitrate = vBitrate;
    audioBitrate = aBitrate;
    sampleRate = sr;
    channels = ch;
    keyframeInterval = kfi;
    return true;
}

int StreamingEngine::AddDestination(const std::string& rtmpUrl) {
    auto* sink = new RtmpSink();
    sink->SetRtmpUrl(rtmpUrl);
    sinks.push_back(sink);
    return static_cast<int>(sinks.size() - 1);
}

bool StreamingEngine::InitVideoEncoder() {
    videoEncoder = AMediaCodec_createEncoderByType("video/avc");
    if (!videoEncoder) {
        LOGE("Failed to create video encoder");
        return false;
    }

    AMediaFormat* format = AMediaFormat_new();
    AMediaFormat_setString(format, AMEDIAFORMAT_KEY_MIME, "video/avc");
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_WIDTH, width);
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_HEIGHT, height);
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_BIT_RATE, videoBitrate);
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_FRAME_RATE, framerate);
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_I_FRAME_INTERVAL, keyframeInterval);
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_COLOR_FORMAT, 0x7F000789); // COLOR_FormatSurface
    AMediaFormat_setInt32(format, "profile", 8);   // AVCProfileHigh
    AMediaFormat_setInt32(format, "level", 2048);   // AVCLevel42
    AMediaFormat_setInt32(format, "bitrate-mode", 2); // CBR

    media_status_t status = AMediaCodec_configure(videoEncoder, format, nullptr, nullptr,
                                                   AMEDIACODEC_CONFIGURE_FLAG_ENCODE);
    AMediaFormat_delete(format);

    if (status != AMEDIA_OK) {
        LOGE("Video encoder configure failed: %d", status);
        AMediaCodec_delete(videoEncoder);
        videoEncoder = nullptr;
        return false;
    }

    status = AMediaCodec_createInputSurface(videoEncoder, &encoderSurface);
    if (status != AMEDIA_OK || !encoderSurface) {
        LOGE("Failed to create encoder input surface: %d", status);
        AMediaCodec_delete(videoEncoder);
        videoEncoder = nullptr;
        return false;
    }

    status = AMediaCodec_start(videoEncoder);
    if (status != AMEDIA_OK) {
        LOGE("Video encoder start failed: %d", status);
        ANativeWindow_release(encoderSurface);
        encoderSurface = nullptr;
        AMediaCodec_delete(videoEncoder);
        videoEncoder = nullptr;
        return false;
    }

    LOGI("Video encoder started: %dx%d @ %d bps", width, height, videoBitrate);
    return true;
}

bool StreamingEngine::InitAudioEncoder() {
    audioEncoder = AMediaCodec_createEncoderByType("audio/mp4a-latm");
    if (!audioEncoder) {
        LOGE("Failed to create audio encoder");
        return false;
    }

    AMediaFormat* format = AMediaFormat_new();
    AMediaFormat_setString(format, AMEDIAFORMAT_KEY_MIME, "audio/mp4a-latm");
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_AAC_PROFILE, 2); // AAC-LC
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_BIT_RATE, audioBitrate);
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_SAMPLE_RATE, sampleRate);
    AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_CHANNEL_COUNT, channels);

    media_status_t status = AMediaCodec_configure(audioEncoder, format, nullptr, nullptr,
                                                   AMEDIACODEC_CONFIGURE_FLAG_ENCODE);
    AMediaFormat_delete(format);

    if (status != AMEDIA_OK) {
        LOGE("Audio encoder configure failed: %d", status);
        AMediaCodec_delete(audioEncoder);
        audioEncoder = nullptr;
        return false;
    }

    status = AMediaCodec_start(audioEncoder);
    if (status != AMEDIA_OK) {
        LOGE("Audio encoder start failed: %d", status);
        AMediaCodec_delete(audioEncoder);
        audioEncoder = nullptr;
        return false;
    }

    LOGI("Audio encoder started: %dHz %dch @ %d bps", sampleRate, channels, audioBitrate);
    return true;
}

bool StreamingEngine::InitBlitResources() {
    GLuint vs = CompileShader(GL_VERTEX_SHADER, BLIT_VERTEX_SHADER);
    GLuint fs = CompileShader(GL_FRAGMENT_SHADER, BLIT_FRAGMENT_SHADER);
    if (!vs || !fs) return false;

    blitProgram = glCreateProgram();
    glAttachShader(blitProgram, vs);
    glAttachShader(blitProgram, fs);
    glLinkProgram(blitProgram);
    glDeleteShader(vs);
    glDeleteShader(fs);

    GLint linkStatus;
    glGetProgramiv(blitProgram, GL_LINK_STATUS, &linkStatus);
    if (!linkStatus) {
        LOGE("Blit program link failed");
        glDeleteProgram(blitProgram);
        blitProgram = 0;
        return false;
    }

    // Compile blit FBO program (sampler2D for composed texture → surface)
    {
        GLuint fboVs = CompileShader(GL_VERTEX_SHADER, BLIT_VERTEX_SHADER);
        GLuint fboFs = CompileShader(GL_FRAGMENT_SHADER, BLIT_FBO_FRAGMENT_SHADER);
        if (!fboVs || !fboFs) return false;

        blitFboProgram = glCreateProgram();
        glAttachShader(blitFboProgram, fboVs);
        glAttachShader(blitFboProgram, fboFs);
        glLinkProgram(blitFboProgram);
        glDeleteShader(fboVs);
        glDeleteShader(fboFs);

        GLint fboLinkStatus;
        glGetProgramiv(blitFboProgram, GL_LINK_STATUS, &fboLinkStatus);
        if (!fboLinkStatus) {
            LOGE("Blit FBO program link failed");
            glDeleteProgram(blitFboProgram);
            blitFboProgram = 0;
            return false;
        }
    }

    // Full-screen quad: pos(x,y) + texcoord(u,v)
    float quad[] = {
        -1.0f, -1.0f,  0.0f, 0.0f,
         1.0f, -1.0f,  1.0f, 0.0f,
        -1.0f,  1.0f,  0.0f, 1.0f,
         1.0f,  1.0f,  1.0f, 1.0f,
    };

    glGenVertexArrays(1, &blitVao);
    glGenBuffers(1, &blitVbo);
    glBindVertexArray(blitVao);
    glBindBuffer(GL_ARRAY_BUFFER, blitVbo);
    glBufferData(GL_ARRAY_BUFFER, sizeof(quad), quad, GL_STATIC_DRAW);
    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)(2 * sizeof(float)));
    glEnableVertexAttribArray(1);
    glBindVertexArray(0);

    // Initialize composition pipeline at encoder resolution
    if (!compositionPipeline.Init(width, height)) {
        LOGE("Composition pipeline init failed");
        return false;
    }

    return true;
}

void StreamingEngine::ReleaseBlitResources() {
    compositionPipeline.Release();
    if (blitVao) { glDeleteVertexArrays(1, &blitVao); blitVao = 0; }
    if (blitVbo) { glDeleteBuffers(1, &blitVbo); blitVbo = 0; }
    if (blitProgram) { glDeleteProgram(blitProgram); blitProgram = 0; }
    if (blitFboProgram) { glDeleteProgram(blitFboProgram); blitFboProgram = 0; }
}

bool StreamingEngine::Start() {
    if (running.load()) return true;

    if (width <= 0 || height <= 0) {
        LOGE("Invalid dimensions: %dx%d", width, height);
        return false;
    }

    if (sinks.empty()) {
        LOGE("No destinations configured");
        return false;
    }

    running.store(true);
    firstVideoFrame = true;
    startTimestampNs = 0;
    lastComposeTimeNs = 0;
    lastGameVideoFrameNs = 0;
    hasExternalAudioSource.store(false);
    audioPcmBuffer.clear();
    audioBufferSamplesWritten = 0;
    statsVideoBytes = 0;
    statsAudioBytes = 0;
    statsFrameCount = 0;
    statsLastUpdateNs = 0;

    encoderThread = std::thread(&StreamingEngine::EncoderThreadFunc, this);
    return true;
}

void StreamingEngine::EncoderThreadFunc() {
    LOGI("Encoder thread started");

    // Init EGL
    if (!eglContext.Init()) {
        LOGE("EGL init failed");
        running.store(false);
        if (errorCallback) errorCallback(1, "EGL initialization failed");
        return;
    }

    // Init video encoder (creates input surface)
    if (!InitVideoEncoder()) {
        LOGE("Video encoder init failed");
        eglContext.Release();
        running.store(false);
        if (errorCallback) errorCallback(2, "Video encoder initialization failed");
        return;
    }

    // Create EGL window surface from encoder input surface
    if (!eglContext.CreateWindowSurface(encoderSurface)) {
        LOGE("EGL window surface creation failed");
        AMediaCodec_stop(videoEncoder);
        AMediaCodec_delete(videoEncoder);
        videoEncoder = nullptr;
        ANativeWindow_release(encoderSurface);
        encoderSurface = nullptr;
        eglContext.Release();
        running.store(false);
        if (errorCallback) errorCallback(3, "EGL window surface creation failed");
        return;
    }

    if (!eglContext.MakeCurrent()) {
        LOGE("EGL make current failed");
        running.store(false);
        if (errorCallback) errorCallback(4, "EGL make current failed");
        return;
    }

    // Init blit resources
    if (!InitBlitResources()) {
        LOGE("Blit resources init failed");
        running.store(false);
        if (errorCallback) errorCallback(5, "Blit resources initialization failed");
        return;
    }

    // Init audio encoder
    if (!InitAudioEncoder()) {
        LOGW("Audio encoder init failed, continuing without audio");
    }

    // Open RTMP sinks
    for (auto* sink : sinks) {
        if (!sink->Open(width, height, framerate, sampleRate, channels)) {
            LOGE("Failed to open RTMP sink");
            if (errorCallback) errorCallback(6, "RTMP connection failed");
        }
    }

    LOGI("Streaming engine fully initialized");

    // Main encoder loop
    while (running.load()) {
        // Process pending preview and layer ops (must run on GL thread)
        ProcessPendingPreviewOps();
        ProcessPendingLayerOps();

        // Process video frames
        bool hadVideoFrames = false;
        {
            std::lock_guard<std::mutex> lock(videoMutex);
            hadVideoFrames = !videoQueue.empty();
            if (hadVideoFrames && statsFrameCount % 30 == 0) {
                LOGI("Processing %zu game video frames", videoQueue.size());
            }
            for (auto& frame : videoQueue) {
                ProcessVideoFrame(frame);
                // Release buffer back to the game's pool
                if (bufferReleasedCallback) {
                    bufferReleasedCallback(frame.bufferIndex);
                }
            }
            videoQueue.clear();
        }

        // Generate standby frames only when no game input has arrived for STANDBY_TIMEOUT_NS.
        // This prevents standby frames from being interleaved with game frames.
        if (!hadVideoFrames && compositionPipeline.IsInitialized()) {
            auto now = std::chrono::steady_clock::now().time_since_epoch();
            int64_t nowNs = std::chrono::duration_cast<std::chrono::nanoseconds>(now).count();
            int64_t frameIntervalNs = 1000000000LL / framerate;
            bool gameActive = lastGameVideoFrameNs > 0 &&
                              (nowNs - lastGameVideoFrameNs) < STANDBY_TIMEOUT_NS;
            if (!gameActive && nowNs - lastComposeTimeNs >= frameIntervalNs) {
                // Compose standby frame (dark background + overlays, no game texture)
                compositionPipeline.Compose(0);
                GLuint composedTex = compositionPipeline.GetComposedTexture();

                eglContext.MakeEncoderCurrent();
                BlitComposedToSurface(composedTex, width, height);
                if (firstVideoFrame) {
                    startTimestampNs = nowNs;
                    firstVideoFrame = false;
                }
                eglContext.SetPresentationTime(nowNs - startTimestampNs);
                eglContext.SwapBuffers();

                if (hasPreview && eglContext.HasPreviewSurface()) {
                    eglContext.MakePreviewCurrent();
                    BlitComposedToSurface(composedTex,
                                          eglContext.GetPreviewWidth(),
                                          eglContext.GetPreviewHeight());
                    eglContext.SwapPreviewBuffers();
                    eglContext.MakeEncoderCurrent();
                }

                lastComposeTimeNs = nowNs;
            }
        }

        // Process audio frames from external sources (accumulate into PCM buffer)
        {
            std::lock_guard<std::mutex> lock(audioMutex);
            for (auto& frame : audioQueue) {
                ProcessAudioFrame(frame);
            }
            audioQueue.clear();
        }

        // Drain encoders — always drain both regardless of pending input
        DrainVideoEncoder();
        if (audioEncoder) {
            DrainAudioEncoder();
        }

        // Flush accumulated audio in AAC-aligned chunks
        if (audioEncoder) {
            FlushAudioBuffer();
        }

        // Update stats every second regardless of frame output
        UpdateStats();

        // Don't spin-wait
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }

    // Cleanup
    LOGI("Encoder thread shutting down");

    ReleaseBlitResources();
    eglContext.DestroyPreviewSurface();
    hasPreview = false;

    for (auto* sink : sinks) {
        sink->Close();
    }

    if (videoEncoder) {
        AMediaCodec_stop(videoEncoder);
        AMediaCodec_delete(videoEncoder);
        videoEncoder = nullptr;
    }
    if (encoderSurface) {
        ANativeWindow_release(encoderSurface);
        encoderSurface = nullptr;
    }
    if (audioEncoder) {
        AMediaCodec_stop(audioEncoder);
        AMediaCodec_delete(audioEncoder);
        audioEncoder = nullptr;
    }

    eglContext.Release();
    LOGI("Encoder thread stopped");
}

void StreamingEngine::ProcessVideoFrame(const VideoFrame& frame) {
    if (!frame.buffer) return;

    // Use wall-clock relative timestamps (matching standby frames) to ensure
    // monotonically increasing PTS. The game's own timestamp is relative to
    // AppStreamingTime which starts at 0, but standby frames may have already
    // advanced the encoder's PTS — backward timestamps cause MediaCodec to drop frames.
    auto now = std::chrono::steady_clock::now().time_since_epoch();
    int64_t nowNs = std::chrono::duration_cast<std::chrono::nanoseconds>(now).count();

    if (firstVideoFrame) {
        startTimestampNs = nowNs;
        firstVideoFrame = false;
    }

    int64_t presentationNs = nowNs - startTimestampNs;
    lastGameVideoFrameNs = nowNs;

    // Wait on GPU fence
    eglContext.WaitFence(frame.fenceFd);

    // Import HardwareBuffer as OES texture
    GLuint texture = eglContext.ImportHardwareBuffer(frame.buffer);
    if (texture == 0) {
        LOGW("Failed to import HardwareBuffer as texture");
        return;
    }

    static int sProcessCount = 0;
    if (++sProcessCount <= 3 || sProcessCount % 300 == 0) {
        LOGI("ProcessVideoFrame: #%d tex=%u pts=%lldms buf=%p idx=%d",
             sProcessCount, texture, (long long)(presentationNs / 1000000),
             frame.buffer, frame.bufferIndex);
    }

    // Compose: game frame + overlay layers → FBO
    compositionPipeline.Compose(texture);
    GLuint composedTex = compositionPipeline.GetComposedTexture();

    // Blit composed texture → encoder surface
    eglContext.MakeEncoderCurrent();
    BlitComposedToSurface(composedTex, width, height);
    eglContext.SetPresentationTime(presentationNs);
    eglContext.SwapBuffers();

    // Blit composed texture → preview surface (if active)
    if (hasPreview && eglContext.HasPreviewSurface()) {
        eglContext.MakePreviewCurrent();
        BlitComposedToSurface(composedTex, eglContext.GetPreviewWidth(),
                              eglContext.GetPreviewHeight());
        eglContext.SwapPreviewBuffers();
        eglContext.MakeEncoderCurrent();
    }

    // Clean up imported texture
    glDeleteTextures(1, &texture);

    // Track compose time so standby frames don't overlap
    lastComposeTimeNs = nowNs;
}

void StreamingEngine::BlitComposedToSurface(GLuint composedTex, int viewportW, int viewportH) {
    glViewport(0, 0, viewportW, viewportH);

    glUseProgram(blitFboProgram);
    glActiveTexture(GL_TEXTURE0);
    glBindTexture(GL_TEXTURE_2D, composedTex);
    glUniform1i(glGetUniformLocation(blitFboProgram, "uTexture"), 0);

    glBindVertexArray(blitVao);
    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
    glBindVertexArray(0);
}

void StreamingEngine::ProcessAudioFrame(const AudioFrame& frame) {
    if (!audioEncoder || frame.pcmData.empty()) return;

    // Accumulate PCM data; FlushAudioBuffer will submit in AAC-aligned chunks
    audioPcmBuffer.insert(audioPcmBuffer.end(), frame.pcmData.begin(), frame.pcmData.end());
}

void StreamingEngine::FlushAudioBuffer() {
    if (!audioEncoder || audioPcmBuffer.empty()) return;

    // Bytes per AAC frame: 1024 samples * channels * 2 bytes (16-bit PCM)
    const size_t aacFrameBytes = AAC_FRAME_SAMPLES * channels * 2;

    while (audioPcmBuffer.size() >= aacFrameBytes) {
        ssize_t inputIndex = AMediaCodec_dequeueInputBuffer(audioEncoder, 5000);
        if (inputIndex < 0) {
            // No buffer available — leave data for next loop iteration
            break;
        }

        size_t bufferSize;
        uint8_t* inputBuffer = AMediaCodec_getInputBuffer(audioEncoder, inputIndex, &bufferSize);
        if (!inputBuffer) break;

        size_t copySize = std::min(aacFrameBytes, bufferSize);
        memcpy(inputBuffer, audioPcmBuffer.data(), copySize);

        // Timestamp based on total samples submitted (continuous, no jitter)
        int64_t timestampUs = audioBufferSamplesWritten * 1000000LL / sampleRate;
        AMediaCodec_queueInputBuffer(audioEncoder, inputIndex, 0, copySize,
                                      timestampUs, 0);

        audioBufferSamplesWritten += AAC_FRAME_SAMPLES;

        // Remove consumed data from front
        audioPcmBuffer.erase(audioPcmBuffer.begin(), audioPcmBuffer.begin() + copySize);
    }

    // Prevent unbounded accumulation if encoder is stuck
    const size_t maxBufferBytes = aacFrameBytes * 16; // ~340ms of audio
    if (audioPcmBuffer.size() > maxBufferBytes) {
        size_t excess = audioPcmBuffer.size() - maxBufferBytes;
        audioPcmBuffer.erase(audioPcmBuffer.begin(), audioPcmBuffer.begin() + excess);
        LOGW("Audio buffer overflow, dropped %zu bytes", excess);
    }
}

void StreamingEngine::DrainVideoEncoder() {
    if (!videoEncoder) return;

    AMediaCodecBufferInfo info;
    ssize_t outputIndex;

    while ((outputIndex = AMediaCodec_dequeueOutputBuffer(videoEncoder, &info, 0)) >= 0) {
        if (info.size > 0) {
            size_t outSize;
            uint8_t* outputData = AMediaCodec_getOutputBuffer(videoEncoder, outputIndex, &outSize);
            if (outputData) {
                bool isKeyframe = (info.flags & AMEDIACODEC_BUFFER_FLAG_KEY_FRAME) != 0;
                int64_t timestampMs = info.presentationTimeUs / 1000;

                for (auto* sink : sinks) {
                    sink->SendVideoPacket(outputData + info.offset, info.size,
                                          timestampMs, isKeyframe);
                }

                std::lock_guard<std::mutex> lock(statsMutex);
                statsVideoBytes += info.size;
                statsFrameCount++;
            }
        }

        if (info.flags & AMEDIACODEC_BUFFER_FLAG_CODEC_CONFIG) {
            // Sequence header (SPS/PPS) — forward to sinks
            size_t outSize;
            uint8_t* configData = AMediaCodec_getOutputBuffer(videoEncoder, outputIndex, &outSize);
            if (configData) {
                for (auto* sink : sinks) {
                    sink->OnVideoFormatReady(configData + info.offset, info.size);
                }
            }
        }

        AMediaCodec_releaseOutputBuffer(videoEncoder, outputIndex, false);
    }

    if (outputIndex == AMEDIACODEC_INFO_OUTPUT_FORMAT_CHANGED) {
        AMediaFormat* newFormat = AMediaCodec_getOutputFormat(videoEncoder);
        if (newFormat) {
            LOGI("Video encoder output format changed");
            AMediaFormat_delete(newFormat);
        }
    }
}

void StreamingEngine::DrainAudioEncoder() {
    if (!audioEncoder) return;

    AMediaCodecBufferInfo info;
    ssize_t outputIndex;

    while ((outputIndex = AMediaCodec_dequeueOutputBuffer(audioEncoder, &info, 0)) >= 0) {
        if (info.size > 0 && !(info.flags & AMEDIACODEC_BUFFER_FLAG_CODEC_CONFIG)) {
            size_t outSize;
            uint8_t* outputData = AMediaCodec_getOutputBuffer(audioEncoder, outputIndex, &outSize);
            if (outputData) {
                int64_t timestampMs = info.presentationTimeUs / 1000;

                for (auto* sink : sinks) {
                    sink->SendAudioPacket(outputData + info.offset, info.size, timestampMs);
                }

                std::lock_guard<std::mutex> lock(statsMutex);
                statsAudioBytes += info.size;
            }
        }

        if (info.flags & AMEDIACODEC_BUFFER_FLAG_CODEC_CONFIG) {
            // AAC config — sinks handle audio format via Open()
        }

        AMediaCodec_releaseOutputBuffer(audioEncoder, outputIndex, false);
    }
}

void StreamingEngine::UpdateStats() {
    auto now = std::chrono::steady_clock::now().time_since_epoch();
    int64_t nowNs = std::chrono::duration_cast<std::chrono::nanoseconds>(now).count();

    std::lock_guard<std::mutex> lock(statsMutex);

    if (statsLastUpdateNs == 0) {
        statsLastUpdateNs = nowNs;
        return;
    }

    int64_t elapsedNs = nowNs - statsLastUpdateNs;
    if (elapsedNs >= 1000000000LL) { // Every second
        double elapsedSec = elapsedNs / 1000000000.0;
        currentStats.videoBitrate = static_cast<int64_t>(statsVideoBytes * 8 / elapsedSec);
        currentStats.audioBitrate = static_cast<int64_t>(statsAudioBytes * 8 / elapsedSec);
        currentStats.fps = static_cast<int>(statsFrameCount / elapsedSec);

        statsVideoBytes = 0;
        statsAudioBytes = 0;
        statsFrameCount = 0;
        statsLastUpdateNs = nowNs;

        if (statsCallback) {
            statsCallback(currentStats);
        }
    }
}

static int sVideoSubmitCount = 0;
void StreamingEngine::SubmitVideoFrame(AHardwareBuffer* buffer, int64_t timestampNs, int fenceFd, int bufferIndex) {
    if (!running.load()) {
        if (fenceFd >= 0) close(fenceFd);
        // Release buffer immediately if not running
        if (bufferReleasedCallback) bufferReleasedCallback(bufferIndex);
        return;
    }

    if (++sVideoSubmitCount % 30 == 1) {
        LOGI("SubmitVideoFrame: frame #%d idx=%d buffer=%p ts=%lld fence=%d",
             sVideoSubmitCount, bufferIndex, buffer, (long long)timestampNs, fenceFd);
    }

    VideoFrame frame;
    frame.buffer = buffer;
    frame.timestampNs = timestampNs;
    frame.fenceFd = fenceFd;
    frame.bufferIndex = bufferIndex;

    std::lock_guard<std::mutex> lock(videoMutex);
    videoQueue.push_back(frame);
}

void StreamingEngine::SubmitAudioFrame(const uint8_t* pcmData, size_t pcmSize, int64_t timestampNs) {
    if (!running.load()) return;

    hasExternalAudioSource.store(true);

    AudioFrame frame;
    frame.pcmData.assign(pcmData, pcmData + pcmSize);
    frame.timestampNs = timestampNs;

    std::lock_guard<std::mutex> lock(audioMutex);
    audioQueue.push_back(std::move(frame));
}

void StreamingEngine::Stop() {
    if (!running.load()) return;

    LOGI("Stopping streaming engine");
    running.store(false);

    if (encoderThread.joinable()) {
        encoderThread.join();
    }
    LOGI("Streaming engine stopped");
}

void StreamingEngine::SetStatsCallback(StatsCallback callback) {
    statsCallback = std::move(callback);
}

void StreamingEngine::SetErrorCallback(ErrorCallback callback) {
    errorCallback = std::move(callback);
}

void StreamingEngine::SetBufferReleasedCallback(BufferReleasedCallback callback) {
    bufferReleasedCallback = std::move(callback);
}

// --- Preview surface ---

void StreamingEngine::SetPreviewSurface(ANativeWindow* window) {
    if (!window) return;
    ANativeWindow_acquire(window);
    std::lock_guard<std::mutex> lock(previewMutex);
    pendingPreviewOps.push_back(PreviewSetOp{window});
}

void StreamingEngine::RemovePreviewSurface() {
    std::lock_guard<std::mutex> lock(previewMutex);
    pendingPreviewOps.push_back(PreviewRemoveOp{});
}

void StreamingEngine::ProcessPendingPreviewOps() {
    std::vector<PreviewOp> ops;
    {
        std::lock_guard<std::mutex> lock(previewMutex);
        ops.swap(pendingPreviewOps);
    }

    for (auto& op : ops) {
        if (auto* setOp = std::get_if<PreviewSetOp>(&op)) {
            eglContext.DestroyPreviewSurface();
            if (eglContext.CreatePreviewSurface(setOp->window)) {
                hasPreview = true;
                LOGI("Preview surface set");
            } else {
                ANativeWindow_release(setOp->window);
                hasPreview = false;
            }
            // MakeEncoderCurrent since CreatePreviewSurface may change current
            eglContext.MakeEncoderCurrent();
        } else if (std::get_if<PreviewRemoveOp>(&op)) {
            eglContext.DestroyPreviewSurface();
            hasPreview = false;
            eglContext.MakeEncoderCurrent();
            LOGI("Preview surface removed");
        }
    }
}

// --- Composition layer management ---

int StreamingEngine::AddCompositionLayer(const uint8_t* rgbaData, int w, int h,
                                          float posX, float posY,
                                          float scaleX, float scaleY,
                                          float rotation, float opacity, int zOrder,
                                          const std::string& tag) {
    int id = nextLayerId.fetch_add(1);
    LayerAddOp addOp;
    addOp.rgbaData.assign(rgbaData, rgbaData + (w * h * 4));
    addOp.w = w;
    addOp.h = h;
    addOp.posX = posX;
    addOp.posY = posY;
    addOp.scaleX = scaleX;
    addOp.scaleY = scaleY;
    addOp.rotation = rotation;
    addOp.opacity = opacity;
    addOp.zOrder = zOrder;
    addOp.tag = tag;
    addOp.assignedId = id;

    std::lock_guard<std::mutex> lock(layerOpMutex);
    pendingLayerOps.push_back(std::move(addOp));
    return id;
}

void StreamingEngine::RemoveCompositionLayer(int layerId) {
    std::lock_guard<std::mutex> lock(layerOpMutex);
    pendingLayerOps.push_back(LayerRemoveOp{layerId});
}

void StreamingEngine::UpdateCompositionLayerTransform(int layerId, float posX, float posY,
                                                       float scaleX, float scaleY,
                                                       float rotation) {
    std::lock_guard<std::mutex> lock(layerOpMutex);
    pendingLayerOps.push_back(LayerTransformOp{layerId, posX, posY, scaleX, scaleY, rotation});
}

void StreamingEngine::UpdateCompositionLayerOpacity(int layerId, float opacity) {
    std::lock_guard<std::mutex> lock(layerOpMutex);
    pendingLayerOps.push_back(LayerOpacityOp{layerId, opacity});
}

void StreamingEngine::SetCompositionLayerEnabled(int layerId, bool enabled) {
    std::lock_guard<std::mutex> lock(layerOpMutex);
    pendingLayerOps.push_back(LayerEnabledOp{layerId, enabled});
}

void StreamingEngine::ProcessPendingLayerOps() {
    std::vector<LayerOp> ops;
    {
        std::lock_guard<std::mutex> lock(layerOpMutex);
        ops.swap(pendingLayerOps);
    }

    for (auto& op : ops) {
        if (auto* addOp = std::get_if<LayerAddOp>(&op)) {
            GLuint tex = CompositionPipeline::UploadTexture(
                addOp->rgbaData.data(), addOp->w, addOp->h);
            if (tex) {
                CompositionTransform transform;
                transform.posX = addOp->posX;
                transform.posY = addOp->posY;
                transform.scaleX = addOp->scaleX;
                transform.scaleY = addOp->scaleY;
                transform.rotation = addOp->rotation;
                compositionPipeline.AddLayer(tex, addOp->w, addOp->h, transform,
                                              addOp->opacity, addOp->zOrder, addOp->tag);
            }
        } else if (auto* removeOp = std::get_if<LayerRemoveOp>(&op)) {
            compositionPipeline.RemoveLayer(removeOp->layerId);
        } else if (auto* transformOp = std::get_if<LayerTransformOp>(&op)) {
            CompositionTransform t;
            t.posX = transformOp->posX;
            t.posY = transformOp->posY;
            t.scaleX = transformOp->scaleX;
            t.scaleY = transformOp->scaleY;
            t.rotation = transformOp->rotation;
            compositionPipeline.UpdateLayerTransform(transformOp->layerId, t);
        } else if (auto* opacityOp = std::get_if<LayerOpacityOp>(&op)) {
            compositionPipeline.UpdateLayerOpacity(opacityOp->layerId, opacityOp->opacity);
        } else if (auto* enabledOp = std::get_if<LayerEnabledOp>(&op)) {
            compositionPipeline.SetLayerEnabled(enabledOp->layerId, enabledOp->enabled);
        }
    }
}
