#include "MovieEncoder.h"
#include "Logger.h"

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/opt.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
}

#include <string>

namespace UsdLayerManager {

// ---------------------------------------------------------------------------
// Impl holds all libav state.
// ---------------------------------------------------------------------------
struct MovieEncoder::Impl {
    AVFormatContext* fmtCtx  = nullptr;
    AVCodecContext*  encCtx  = nullptr;
    AVStream*        stream  = nullptr;
    SwsContext*      sws     = nullptr;
    AVFrame*         frame   = nullptr;
    AVPacket*        pkt     = nullptr;
    int64_t          pts     = 0;
    int              width   = 0;
    int              height  = 0;
};

// ---------------------------------------------------------------------------
MovieEncoder::MovieEncoder()  = default;
MovieEncoder::~MovieEncoder() { Close(); Reset(); }

// ---------------------------------------------------------------------------
static std::string AvError(int code) {
    char buf[128] = {};
    av_strerror(code, buf, sizeof(buf));
    return buf;
}

// ---------------------------------------------------------------------------
bool MovieEncoder::Open(const char* outputPath, int width, int height, double fps) {
    if (m_open) Close();
    Reset();

    m_impl = new Impl();
    m_impl->width  = width;
    m_impl->height = height;

    // ── Format context ────────────────────────────────────────────────────
    int ret = avformat_alloc_output_context2(&m_impl->fmtCtx,
                                             nullptr, nullptr, outputPath);
    if (ret < 0 || !m_impl->fmtCtx) {
        m_error = "avformat_alloc_output_context2: " + AvError(ret);
        Reset(); return false;
    }

    // ── Find encoder (prefer libx264, fall back to hardware or generic) ───
    const char* encoderOrder[] = {
        "libx264", "h264_nvenc", "h264_qsv", "h264_amf", "libopenh264", nullptr
    };
    const AVCodec* codec = nullptr;
    for (int i = 0; encoderOrder[i]; ++i) {
        codec = avcodec_find_encoder_by_name(encoderOrder[i]);
        if (codec) { LOG_INFO(std::string("MovieEncoder: using ") + codec->name); break; }
    }
    if (!codec) codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec) {
        m_error = "No H.264 encoder found (install ffmpeg GPL or NVIDIA/Intel drivers)";
        Reset(); return false;
    }

    // ── Stream ────────────────────────────────────────────────────────────
    m_impl->stream = avformat_new_stream(m_impl->fmtCtx, nullptr);
    if (!m_impl->stream) {
        m_error = "avformat_new_stream failed";
        Reset(); return false;
    }
    m_impl->stream->id = 0;

    // ── Codec context ─────────────────────────────────────────────────────
    m_impl->encCtx = avcodec_alloc_context3(codec);
    if (!m_impl->encCtx) {
        m_error = "avcodec_alloc_context3 failed";
        Reset(); return false;
    }

    AVRational fpsRat = av_d2q(fps, 65536);
    m_impl->encCtx->width     = width;
    m_impl->encCtx->height    = height;
    m_impl->encCtx->pix_fmt   = AV_PIX_FMT_YUV420P;
    m_impl->encCtx->framerate = fpsRat;
    m_impl->encCtx->time_base = av_inv_q(fpsRat);  // {1, fps_num}
    // Reasonable quality defaults for CRF-capable encoders.
    av_opt_set(m_impl->encCtx->priv_data, "crf",    "23",   AV_OPT_SEARCH_CHILDREN);
    av_opt_set(m_impl->encCtx->priv_data, "preset", "fast", AV_OPT_SEARCH_CHILDREN);
    // Ensure broad playback compatibility.
    av_opt_set(m_impl->encCtx->priv_data, "profile", "high", AV_OPT_SEARCH_CHILDREN);
    av_opt_set(m_impl->encCtx->priv_data, "level",   "4.1",  AV_OPT_SEARCH_CHILDREN);

    if (m_impl->fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
        m_impl->encCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    ret = avcodec_open2(m_impl->encCtx, codec, nullptr);
    if (ret < 0) {
        m_error = "avcodec_open2: " + AvError(ret);
        Reset(); return false;
    }

    ret = avcodec_parameters_from_context(m_impl->stream->codecpar, m_impl->encCtx);
    if (ret < 0) {
        m_error = "avcodec_parameters_from_context: " + AvError(ret);
        Reset(); return false;
    }
    m_impl->stream->time_base = m_impl->encCtx->time_base;

    // ── Open output file ──────────────────────────────────────────────────
    if (!(m_impl->fmtCtx->oformat->flags & AVFMT_NOFILE)) {
        ret = avio_open(&m_impl->fmtCtx->pb, outputPath, AVIO_FLAG_WRITE);
        if (ret < 0) {
            m_error = std::string("avio_open(") + outputPath + "): " + AvError(ret);
            Reset(); return false;
        }
    }

    ret = avformat_write_header(m_impl->fmtCtx, nullptr);
    if (ret < 0) {
        m_error = "avformat_write_header: " + AvError(ret);
        Reset(); return false;
    }

    // ── Frame + packet buffers ────────────────────────────────────────────
    m_impl->frame = av_frame_alloc();
    if (!m_impl->frame) {
        m_error = "av_frame_alloc failed";
        Reset(); return false;
    }
    m_impl->frame->width  = width;
    m_impl->frame->height = height;
    m_impl->frame->format = AV_PIX_FMT_YUV420P;
    ret = av_frame_get_buffer(m_impl->frame, 0);
    if (ret < 0) {
        m_error = "av_frame_get_buffer: " + AvError(ret);
        Reset(); return false;
    }

    m_impl->pkt = av_packet_alloc();
    if (!m_impl->pkt) {
        m_error = "av_packet_alloc failed";
        Reset(); return false;
    }

    // ── SwsContext: RGBA → YUV420P ────────────────────────────────────────
    m_impl->sws = sws_getContext(
        width, height, AV_PIX_FMT_RGBA,
        width, height, AV_PIX_FMT_YUV420P,
        SWS_BILINEAR, nullptr, nullptr, nullptr);
    if (!m_impl->sws) {
        m_error = "sws_getContext failed";
        Reset(); return false;
    }

    m_impl->pts = 0;
    m_open      = true;
    return true;
}

// ---------------------------------------------------------------------------
bool MovieEncoder::WriteFrame(const uint8_t* rgba) {
    if (!m_open || !m_impl) { m_error = "encoder not open"; return false; }

    int ret = av_frame_make_writable(m_impl->frame);
    if (ret < 0) { m_error = "av_frame_make_writable: " + AvError(ret); return false; }

    const uint8_t* srcSlice[1] = { rgba };
    int            srcStride[1] = { m_impl->width * 4 };
    sws_scale(m_impl->sws,
              srcSlice, srcStride, 0, m_impl->height,
              m_impl->frame->data, m_impl->frame->linesize);

    m_impl->frame->pts = m_impl->pts++;

    ret = avcodec_send_frame(m_impl->encCtx, m_impl->frame);
    if (ret < 0) { m_error = "avcodec_send_frame: " + AvError(ret); return false; }

    return DrainPackets();
}

// ---------------------------------------------------------------------------
bool MovieEncoder::Close() {
    if (!m_open || !m_impl) return true;

    // Flush
    avcodec_send_frame(m_impl->encCtx, nullptr);
    DrainPackets();

    av_write_trailer(m_impl->fmtCtx);

    if (m_impl->fmtCtx && !(m_impl->fmtCtx->oformat->flags & AVFMT_NOFILE))
        avio_closep(&m_impl->fmtCtx->pb);

    Reset();
    m_open = false;
    return true;
}

// ---------------------------------------------------------------------------
bool MovieEncoder::DrainPackets() {
    int ret = 0;
    while ((ret = avcodec_receive_packet(m_impl->encCtx, m_impl->pkt)) == 0) {
        av_packet_rescale_ts(m_impl->pkt,
                             m_impl->encCtx->time_base,
                             m_impl->stream->time_base);
        m_impl->pkt->stream_index = m_impl->stream->index;
        ret = av_interleaved_write_frame(m_impl->fmtCtx, m_impl->pkt);
        av_packet_unref(m_impl->pkt);
        if (ret < 0) { m_error = "av_interleaved_write_frame: " + AvError(ret); return false; }
    }
    return ret == AVERROR(EAGAIN) || ret == AVERROR_EOF;
}

// ---------------------------------------------------------------------------
void MovieEncoder::Reset() {
    if (!m_impl) return;
    if (m_impl->sws)    { sws_freeContext(m_impl->sws);         m_impl->sws    = nullptr; }
    if (m_impl->pkt)    { av_packet_free(&m_impl->pkt);         m_impl->pkt    = nullptr; }
    if (m_impl->frame)  { av_frame_free(&m_impl->frame);        m_impl->frame  = nullptr; }
    if (m_impl->encCtx) { avcodec_free_context(&m_impl->encCtx); m_impl->encCtx = nullptr; }
    if (m_impl->fmtCtx) { avformat_free_context(m_impl->fmtCtx); m_impl->fmtCtx = nullptr; }
    delete m_impl;
    m_impl = nullptr;
    m_open = false;
}

} // namespace UsdLayerManager
