#include "stream_video_encoder.h"
#include "debug_config.h"
#include <QPainter>
#include <algorithm>
extern "C" {
#include "../../../src/stream/stream_media.h"
#ifdef HAVE_FFMPEG
#include <libavcodec/avcodec.h>
#include <libavcodec/bsf.h>
#include <libavutil/hwcontext.h>
#include <libavutil/opt.h>
#include <libswscale/swscale.h>
#endif
}

struct StreamVideoEncoder::Impl {
    QString error, name;
    StreamVideoConfig config;
#ifdef HAVE_FFMPEG
    AVCodecContext* codec = nullptr;
    AVBSFContext* filter = nullptr;
    AVFrame* frame = nullptr;
    AVBufferRef* device = nullptr;
    SwsContext* scale = nullptr;
    ~Impl() {
        av_bsf_free(&filter); avcodec_free_context(&codec); av_frame_free(&frame);
        av_buffer_unref(&device); sws_freeContext(scale);
    }
#endif
};
StreamVideoEncoder::StreamVideoEncoder() : m(std::make_unique<Impl>()) {}
StreamVideoEncoder::~StreamVideoEncoder() = default;
void StreamVideoEncoder::close() { m = std::make_unique<Impl>(); }
QString StreamVideoEncoder::error() const { return m->error; }
QString StreamVideoEncoder::name() const { return m->name; }
QString StreamVideoEncoder::codecName(int codec) {
    static const char* names[] = {"", "H.264", "H.265", "VP9", "AV1", "H.266"};
    return codec >= 1 && codec <= 5 ? QString::fromLatin1(names[codec]) : QString();
}

#ifdef HAVE_FFMPEG
namespace {
QString ffError(int result) { char text[AV_ERROR_MAX_STRING_SIZE]; av_strerror(result, text, sizeof(text)); return QString::fromUtf8(text); }
QList<QByteArray> encoders(int codec) {
    switch (codec) {
    case STREAM_MEDIA_H264: return {"h264_nvenc", "h264_qsv", "h264_amf", "h264_vaapi", "libx264", "libopenh264"};
    case STREAM_MEDIA_H265: return {"hevc_nvenc", "hevc_qsv", "hevc_amf", "hevc_vaapi", "libx265"};
    case STREAM_MEDIA_VP9: return {"vp9_qsv", "vp9_vaapi", "libvpx-vp9"};
    case STREAM_MEDIA_AV1: return {"av1_nvenc", "av1_qsv", "av1_amf", "av1_vaapi", "libsvtav1", "libaom-av1"};
    case STREAM_MEDIA_H266: return {"libvvenc"};
    default: return {};
    }
}
AVPixelFormat softwareFormat(const AVCodec* codec) {
    const AVPixelFormat* formats = nullptr;
#if LIBAVCODEC_VERSION_MAJOR >= 61
    avcodec_get_supported_config(nullptr, codec, AV_CODEC_CONFIG_PIX_FORMAT, 0, reinterpret_cast<const void**>(&formats), nullptr);
#else
    formats = codec->pix_fmts;
#endif
    for (auto preferred : {AV_PIX_FMT_YUV420P, AV_PIX_FMT_NV12, AV_PIX_FMT_YUV420P10LE}) {
        if (!formats) return preferred;
        for (auto f = formats; *f != AV_PIX_FMT_NONE; f++) if (*f == preferred) return preferred;
    }
    return AV_PIX_FMT_NONE;
}
}
#endif

bool StreamVideoEncoder::open(const StreamVideoConfig& config) {
    close();
    if (config.width < 16 || config.height < 16 || config.width > 4096 || config.height > 4096
        || config.width % 2 || config.height % 2 || config.fps < 1 || config.fps > 60
        || config.bitrate < 64000 || config.bitrate > 16000000 || config.codec < 1 || config.codec > 5) {
        m->error = "Invalid video configuration"; DEBUG_ERROR(DEBUG_CATEGORY_VIDEO, "desktop encoder invalid configuration"); return false;
    }
#ifdef HAVE_FFMPEG
    for (const auto& name : encoders(config.codec)) {
        const AVCodec* encoder = avcodec_find_encoder_by_name(name.constData());
        if (!encoder) continue;
        close(); m->config = config; m->name = QString::fromLatin1(name);
        m->codec = avcodec_alloc_context3(encoder);
        if (!m->codec) { m->error = "Video encoder allocation failed"; break; }
        auto* c = m->codec;
        c->width = config.width; c->height = config.height; c->time_base = {1, 1000000}; c->framerate = {config.fps, 1};
        c->bit_rate = config.bitrate; c->rc_max_rate = config.bitrate; c->rc_buffer_size = config.bitrate / 4;
        c->gop_size = config.fps; c->max_b_frames = 0; c->thread_count = 2;
        c->flags |= AV_CODEC_FLAG_GLOBAL_HEADER | AV_CODEC_FLAG_CLOSED_GOP;
        c->color_range = AVCOL_RANGE_MPEG; c->colorspace = AVCOL_SPC_BT709;
        c->color_primaries = AVCOL_PRI_BT709; c->color_trc = AVCOL_TRC_BT709;
        AVPixelFormat sw = softwareFormat(encoder);
        int result = 0;
        if (name.endsWith("_vaapi")) {
            sw = AV_PIX_FMT_NV12; c->pix_fmt = AV_PIX_FMT_VAAPI;
            result = av_hwdevice_ctx_create(&m->device, AV_HWDEVICE_TYPE_VAAPI, nullptr, nullptr, 0);
            if (result >= 0) {
                c->hw_frames_ctx = av_hwframe_ctx_alloc(m->device);
                if (!c->hw_frames_ctx) result = AVERROR(ENOMEM);
                else {
                    auto* frames = reinterpret_cast<AVHWFramesContext*>(c->hw_frames_ctx->data);
                    frames->format = c->pix_fmt; frames->sw_format = sw; frames->width = c->width; frames->height = c->height;
                    frames->initial_pool_size = 4; result = av_hwframe_ctx_init(c->hw_frames_ctx);
                }
            }
        } else c->pix_fmt = sw;
        if (sw == AV_PIX_FMT_NONE) result = AVERROR(ENOSYS);
        AVDictionary* options = nullptr;
        if (name == "libx264" || name == "libx265") {
            av_dict_set(&options, "preset", "ultrafast", 0); av_dict_set(&options, "tune", "zerolatency", 0);
            if (name == "libx264") av_dict_set(&options, "x264-params", "scenecut=0:open-gop=0:repeat-headers=0", 0);
            else av_dict_set(&options, "x265-params", "bframes=0:rc-lookahead=0:open-gop=0:pools=1:frame-threads=1:log-level=error", 0);
        } else if (name.endsWith("_nvenc")) {
            av_dict_set(&options, "preset", "p1", 0); av_dict_set(&options, "tune", "ull", 0);
            av_dict_set(&options, "zerolatency", "1", 0); av_dict_set(&options, "rc-lookahead", "0", 0);
            av_dict_set(&options, "forced-idr", "1", 0);
        } else if (name == "libvpx-vp9" || name == "libaom-av1") {
            av_dict_set(&options, "deadline", "realtime", 0); av_dict_set(&options, "cpu-used", "8", 0);
            av_dict_set(&options, "lag-in-frames", "0", 0); av_dict_set(&options, "row-mt", "1", 0);
        } else if (name == "libsvtav1") {
            av_dict_set(&options, "preset", "12", 0); av_dict_set(&options, "svtav1-params", "pred-struct=1:lookahead=0:lp=2", 0);
        } else if (name.endsWith("_qsv")) {
            av_dict_set(&options, "async_depth", "1", 0); av_dict_set(&options, "look_ahead", "0", 0);
        } else if (name == "libvvenc") av_dict_set(&options, "preset", "fast", 0);
        if (result >= 0) result = avcodec_open2(c, encoder, &options);
        av_dict_free(&options);
        if (result >= 0) {
            m->frame = av_frame_alloc();
            if (!m->frame) result = AVERROR(ENOMEM);
            else {
                m->frame->format = sw; m->frame->width = c->width; m->frame->height = c->height;
                m->frame->color_range = c->color_range; m->frame->colorspace = c->colorspace;
                result = av_frame_get_buffer(m->frame, 32);
            }
        }
        const char* filter = config.codec == STREAM_MEDIA_H264 ? "h264_mp4toannexb"
            : config.codec == STREAM_MEDIA_H265 ? "hevc_mp4toannexb" : config.codec == STREAM_MEDIA_H266 ? "vvc_mp4toannexb" : nullptr;
        if (result >= 0 && filter) {
            auto* bsf = av_bsf_get_by_name(filter);
            result = bsf ? av_bsf_alloc(bsf, &m->filter) : AVERROR(ENOSYS);
            if (result >= 0) result = avcodec_parameters_from_context(m->filter->par_in, c);
            if (result >= 0) { m->filter->time_base_in = c->time_base; result = av_bsf_init(m->filter); }
        }
        if (result < 0) {
            m->error = ffError(result);
            DEBUG_DEBUG(DEBUG_CATEGORY_VIDEO, "desktop encoder unavailable name=%s size=%dx%d error=%s",
                name.constData(), config.width, config.height, qPrintable(m->error)); continue;
        }
        QImage probe(config.width, config.height, QImage::Format_RGB32); probe.fill(Qt::black);
        QList<StreamEncodedFrame> output;
        if (!encode(probe, 0, true, output) || output.isEmpty() || !output.front().key) {
            DEBUG_WARN(DEBUG_CATEGORY_VIDEO, "desktop encoder probe lacks immediate KEY name=%s error=%s", name.constData(), qPrintable(m->error));
            continue;
        }
        DEBUG_INFO(DEBUG_CATEGORY_VIDEO, "desktop encoder ready name=%s codec=%d size=%dx%d fps=%d bitrate=%d probe_KEY=%d",
            name.constData(), config.codec, config.width, config.height, config.fps, config.bitrate, output.front().data.size());
        return true;
    }
    if (m->error.isEmpty()) m->error = "No working encoder for " + codecName(config.codec);
#else
    m->error = "FFmpeg is unavailable";
#endif
    DEBUG_ERROR(DEBUG_CATEGORY_VIDEO, "desktop encoder open failed: %s", qPrintable(m->error)); return false;
}

bool StreamVideoEncoder::encode(const QImage& image, uint64_t pts, bool key, QList<StreamEncodedFrame>& output) {
#ifdef HAVE_FFMPEG
    if (!m->codec || !m->frame || image.isNull() || pts > INT64_MAX) {
        m->error = "Invalid video frame"; DEBUG_ERROR(DEBUG_CATEGORY_VIDEO, "desktop encode invalid frame/state"); return false;
    }
    QImage canvas(m->config.width, m->config.height, QImage::Format_RGB32); canvas.fill(Qt::black);
    QSize fitted = image.size().scaled(canvas.size(), Qt::KeepAspectRatio);
    { QPainter painter(&canvas); painter.drawImage(QRect(QPoint((canvas.width() - fitted.width()) / 2,
        (canvas.height() - fitted.height()) / 2), fitted), image); }
    int result = av_frame_make_writable(m->frame);
    m->scale = sws_getCachedContext(m->scale, canvas.width(), canvas.height(), AV_PIX_FMT_BGRA,
        m->frame->width, m->frame->height, static_cast<AVPixelFormat>(m->frame->format), SWS_FAST_BILINEAR, nullptr, nullptr, nullptr);
    if (!m->scale) result = AVERROR(ENOMEM);
    if (result >= 0) {
        const uint8_t* input[] = {canvas.constBits()}; int stride[] = {canvas.bytesPerLine()};
        result = sws_scale(m->scale, input, stride, 0, canvas.height(), m->frame->data, m->frame->linesize);
    }
    m->frame->pts = static_cast<int64_t>(pts); m->frame->pict_type = key ? AV_PICTURE_TYPE_I : AV_PICTURE_TYPE_NONE;
    AVFrame* hardware = nullptr;
    if (result >= 0 && m->codec->hw_frames_ctx) {
        hardware = av_frame_alloc();
        result = hardware ? av_hwframe_get_buffer(m->codec->hw_frames_ctx, hardware, 0) : AVERROR(ENOMEM);
        if (result >= 0) result = av_hwframe_transfer_data(hardware, m->frame, 0);
        if (result >= 0) result = av_frame_copy_props(hardware, m->frame);
    }
    if (result >= 0) result = avcodec_send_frame(m->codec, hardware ? hardware : m->frame);
    av_frame_free(&hardware);
    AVPacket* packet = av_packet_alloc();
    if (!packet) result = AVERROR(ENOMEM);
    while (result >= 0) {
        result = avcodec_receive_packet(m->codec, packet);
        if (result < 0) break;
        if (m->filter) {
            result = av_bsf_send_packet(m->filter, packet);
            if (result >= 0) result = av_bsf_receive_packet(m->filter, packet);
            if (result < 0) break;
        }
        StreamEncodedFrame encoded; encoded.key = (packet->flags & AV_PKT_FLAG_KEY) != 0;
        encoded.pts = packet->pts >= 0 ? static_cast<uint64_t>(packet->pts) : pts;
        if (encoded.key) {
            const uint8_t* config = m->filter ? m->filter->par_out->extradata : m->codec->extradata;
            int length = m->filter ? m->filter->par_out->extradata_size : m->codec->extradata_size;
            if (length < 0 || length > static_cast<int>(STREAM_MEDIA_MAX_CONFIG)) { result = AVERROR(EINVAL); break; }
            encoded.data.resize(4 + length); stream_write_be(reinterpret_cast<uint8_t*>(encoded.data.data()), length, 4);
            if (length) memcpy(encoded.data.data() + 4, config, static_cast<size_t>(length));
        }
        encoded.data.append(reinterpret_cast<const char*>(packet->data), packet->size);
        if (encoded.data.size() > STREAM_MEDIA_MAX_FRAME) { result = AVERROR(EOVERFLOW); break; }
        output.push_back(std::move(encoded)); av_packet_unref(packet);
    }
    av_packet_free(&packet);
    if (result == AVERROR(EAGAIN) || result == AVERROR_EOF) return true;
    if (result < 0) { m->error = ffError(result); DEBUG_ERROR(DEBUG_CATEGORY_VIDEO, "desktop encode failed name=%s PTS=%llu: %s",
        qPrintable(m->name), (unsigned long long)pts, qPrintable(m->error)); return false; }
    return true;
#else
    Q_UNUSED(image); Q_UNUSED(pts); Q_UNUSED(key); Q_UNUSED(output); return false;
#endif
}

QList<int> StreamVideoEncoder::available() {
    QList<int> result;
    for (int codec = 1; codec <= 5; codec++) {
        StreamVideoEncoder encoder; StreamVideoConfig config; config.width = 320; config.height = 180; config.codec = codec;
        if (encoder.open(config)) result.push_back(codec);
    }
    return result;
}