#include "CFFmpegAV1Encoder.h"
#include "common/config.h"
#include "common/logger.h"

// 合规守护：DISABLE_FFMPEG_FOR_TEST=1 时整个实现移出编译单元（FFmpeg lib 已在
// CFFmpegH264Encoder.cpp 用同条件链接，此处不重复 #pragma comment）
#if defined(_WIN64) && !DISABLE_FFMPEG_FOR_TEST

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavutil/opt.h>
#include <libavutil/imgutils.h>
#include <libyuv/libyuv.h>
}

#include <string.h>

// FFmpeg / 系统库已经由 CFFmpegH264Encoder.cpp 的 #pragma comment(lib) 引入。
// 这里不再重复声明（重复 #pragma comment 在同一 link 单元不冲突但冗余）。

// av_opt_set 包装：拼错的参数值会被 FFmpeg 静默忽略，包一层日志便于发现。
// 实现与 CFFmpegH264Encoder 内的 helper 相同；放成 static 文件内可见即可。
static void setOpt(void* obj, const char* name, const char* val, const char* backend) {
    int rc = av_opt_set(obj, name, val, 0);
    if (rc < 0) {
        char errbuf[128] = {0};
        av_strerror(rc, errbuf, sizeof(errbuf));
        Mprintf("[WARN] av_opt_set('%s'='%s') on %s failed (%d): %s\n",
                name, val, backend, rc, errbuf);
    }
}
static void setOptInt(void* obj, const char* name, int64_t val, const char* backend) {
    int rc = av_opt_set_int(obj, name, val, 0);
    if (rc < 0) {
        char errbuf[128] = {0};
        av_strerror(rc, errbuf, sizeof(errbuf));
        Mprintf("[WARN] av_opt_set_int('%s'=%lld) on %s failed (%d): %s\n",
                name, (long long)val, backend, rc, errbuf);
    }
}

// AV1 硬编后端探测顺序，没有 av1_mf 兜底（FFmpeg 7.1 不支持）。
// 全失败时 EncoderFactory 自动回退到 H.264 路径，行为对称。
static const char* kAV1Backends[] = {
    "av1_nvenc",    // NVIDIA RTX 40 / 50 系（Ada Lovelace+）
    "av1_amf",      // AMD RX 7000+（RDNA 3+）
    "av1_qsv",      // Intel Arc 独显 / 部分 11 代+ 核显
};

CFFmpegAV1Encoder::CFFmpegAV1Encoder() = default;

CFFmpegAV1Encoder::~CFFmpegAV1Encoder() {
    close();
}

void CFFmpegAV1Encoder::cleanupCodec() {
    if (m_packet) { av_packet_free(&m_packet); m_packet = nullptr; }
    if (m_frame)  { av_frame_free(&m_frame);   m_frame  = nullptr; }
    if (m_ctx)    { avcodec_free_context(&m_ctx); m_ctx = nullptr; }
}

void CFFmpegAV1Encoder::close() {
    cleanupCodec();
    m_backend.clear();
    m_pts = 0;
    m_forceIDR = false;
}

bool CFFmpegAV1Encoder::open(const EncoderParams& params) {
    close();
    for (const char* name : kAV1Backends) {
        if (tryOpenBackend(name, params)) {
            m_backend = name;
            return true;
        }
        cleanupCodec();
    }
    return false;
}

bool CFFmpegAV1Encoder::tryOpenBackend(const char* name, const EncoderParams& p) {
    const AVCodec* codec = avcodec_find_encoder_by_name(name);
    if (!codec) {
        // AV1 硬编没注册 = 老 ffmpeg lib 不含 AV1 encoder（compress\ffmpeg 没启用 av1）
        Mprintf("=> FFmpeg: AV1 encoder '%s' NOT in linked lib\n", name);
        return false;
    }

    m_ctx = avcodec_alloc_context3(codec);
    if (!m_ctx) {
        Mprintf("=> FFmpeg: avcodec_alloc_context3('%s') failed\n", name);
        return false;
    }

    m_ctx->width        = p.width  & ~1;
    m_ctx->height       = p.height & ~1;
    m_ctx->time_base    = AVRational{1, p.fps};
    m_ctx->framerate    = AVRational{p.fps, 1};
    m_ctx->pix_fmt      = AV_PIX_FMT_NV12;
    m_ctx->gop_size     = p.fps * (p.gop_seconds > 0 ? p.gop_seconds : 15);
    m_ctx->max_b_frames = 0;
    m_ctx->bit_rate     = (int64_t)p.bitrate_kbps * 1000;
    m_ctx->rc_max_rate  = (int64_t)p.bitrate_kbps * 1500;
    m_ctx->rc_buffer_size = (int)(p.bitrate_kbps * 1000);

    // RC 策略与 H.264 路径对齐：peak-constrained VBR，远控静态画面省带宽。
    if (strcmp(name, "av1_nvenc") == 0) {
        // av1_nvenc preset p1~p7；远控 p5 兼顾质量与速度。
        // tile-columns=1 把帧切两列，解码端并行更友好（浏览器 AV1 解码常用 SIMD/多线程）
        setOpt(m_ctx->priv_data, "preset", "p5", name);
        setOpt(m_ctx->priv_data, "tune", "ll", name);
        setOpt(m_ctx->priv_data, "rc", "vbr", name);
        setOpt(m_ctx->priv_data, "zerolatency", "1", name);
        setOptInt(m_ctx->priv_data, "tile-columns", 1, name);
    } else if (strcmp(name, "av1_amf") == 0) {
        // av1_amf 选项命名与 h264_amf 大体一致，rc 同样支持 vbr_peak
        // (见 ffmpeg -h encoder=av1_amf)。静态画面省码率四件套同 H.264 路径。
        setOpt(m_ctx->priv_data, "usage", "lowlatency", name);
        setOpt(m_ctx->priv_data, "quality", "quality", name);
        setOpt(m_ctx->priv_data, "rc", "vbr_peak", name);
        setOptInt(m_ctx->priv_data, "vbaq", 1, name);
        setOptInt(m_ctx->priv_data, "preanalysis", 1, name);
        setOptInt(m_ctx->priv_data, "filler_data", 0, name);
        setOptInt(m_ctx->priv_data, "enforce_hrd", 0, name);
    } else if (strcmp(name, "av1_qsv") == 0) {
        // av1_qsv：bit_rate < max_rate 时自动 VBR
        setOpt(m_ctx->priv_data, "preset", "slow", name);
        setOptInt(m_ctx->priv_data, "async_depth", 1, name);
        setOptInt(m_ctx->priv_data, "low_power", 0, name);
    }

    int ret = avcodec_open2(m_ctx, codec, nullptr);
    if (ret < 0) {
        // 找到了但开不起来：无对应 GPU / 驱动太旧 / 跨适配器
        char errbuf[128] = {0};
        av_strerror(ret, errbuf, sizeof(errbuf));
        Mprintf("=> FFmpeg: avcodec_open2('%s') failed (%d): %s\n", name, ret, errbuf);
        return false;
    }

    m_frame = av_frame_alloc();
    if (!m_frame) return false;
    m_frame->format = AV_PIX_FMT_NV12;
    m_frame->width  = m_ctx->width;
    m_frame->height = m_ctx->height;
    if (av_frame_get_buffer(m_frame, 32) < 0) {
        Mprintf("=> FFmpeg: av_frame_get_buffer failed\n");
        return false;
    }

    m_packet = av_packet_alloc();
    return m_packet != nullptr;
}

void CFFmpegAV1Encoder::setBitrate(int kbps) {
    if (!m_ctx) return;
    m_ctx->bit_rate     = (int64_t)kbps * 1000;
    m_ctx->rc_max_rate  = (int64_t)kbps * 1500;
    m_ctx->rc_buffer_size = (int)(kbps * 1000);
    // 同 H.264 路径：多数硬编不支持运行时改 bit_rate 让 ctx 立刻生效；
    // 这里仅更新数值，下次 open 时生效。
}

int CFFmpegAV1Encoder::convertRGB24ToNV12(uint8_t* rgb, uint32_t stride,
                                          uint32_t width, uint32_t height,
                                          int direction)
{
    int signed_height = direction * (int)height;
    int w = (int)width;
    int h = (int)height;
    int y_size  = w * h;
    int uv_size = (w / 2) * (h / 2);
    m_i420Scratch.resize(y_size + 2 * uv_size);

    uint8_t* y = m_i420Scratch.data();
    uint8_t* u = y + y_size;
    uint8_t* v = u + uv_size;

    if (libyuv::RGB24ToI420(rgb, stride, y, w, u, w / 2, v, w / 2, w, signed_height) != 0)
        return -1;
    if (libyuv::I420ToNV12(y, w, u, w / 2, v, w / 2,
                           m_frame->data[0], m_frame->linesize[0],
                           m_frame->data[1], m_frame->linesize[1],
                           w, h) != 0)
        return -1;
    return 0;
}

int CFFmpegAV1Encoder::encode(
    uint8_t* rgb, uint8_t bpp, uint32_t stride,
    uint32_t width, uint32_t height,
    uint8_t** lppData, uint32_t* lpSize, int direction)
{
    if (!m_ctx || !m_frame || !m_packet) return -1;
    if (av_frame_make_writable(m_frame) < 0) return -1;

    int w = (int)width;
    int h = (int)height;
    int signed_height = direction * h;

    if (bpp == 32) {
        if (libyuv::ARGBToNV12(
                rgb, stride,
                m_frame->data[0], m_frame->linesize[0],
                m_frame->data[1], m_frame->linesize[1],
                w, signed_height) != 0) {
            return -1;
        }
    } else if (bpp == 24) {
        if (convertRGB24ToNV12(rgb, stride, width, height, direction) != 0) {
            return -1;
        }
    } else {
        return -2;
    }

    m_frame->pts = m_pts++;
    if (m_forceIDR) {
        m_frame->pict_type = AV_PICTURE_TYPE_I;
        m_forceIDR = false;
    } else {
        m_frame->pict_type = AV_PICTURE_TYPE_NONE;
    }

    int ret = avcodec_send_frame(m_ctx, m_frame);
    if (ret < 0) return -3;

    ret = avcodec_receive_packet(m_ctx, m_packet);
    if (ret == AVERROR(EAGAIN)) {
        *lppData = nullptr;
        *lpSize = 0;
        return 0;
    }
    if (ret < 0) return -4;

    m_outputBuffer.assign(m_packet->data, m_packet->data + m_packet->size);
    *lppData = m_outputBuffer.data();
    *lpSize  = (uint32_t)m_outputBuffer.size();
    av_packet_unref(m_packet);
    return 0;
}

#endif // _WIN64 && !DISABLE_FFMPEG_FOR_TEST