#include "screencast/codec/encoder.h"
#include "screencast/utils/clock.h"

#include "ffmpeg_raii.h"

#include <algorithm>
#include <array>
#include <atomic>
#include <cstdint>
#include <limits>
#include <optional>
#include <span>
#include <string>
#include <vector>

namespace sc {
namespace {

using namespace sc::detail;

AVPixelFormat to_ffmpeg_format(PixelFormat format) noexcept {
    switch (format) {
        case PixelFormat::Rgba:
            return AV_PIX_FMT_RGBA;
        case PixelFormat::Bgrx:
            // BGRx is BGR with an unused fourth byte; FFmpeg models that layout as BGRA.
            return AV_PIX_FMT_BGRA;
        case PixelFormat::Yuv420p:
            return AV_PIX_FMT_YUV420P;
    }
    return AV_PIX_FMT_NONE;
}

bool is_packed_rgb(PixelFormat format) noexcept {
    return format == PixelFormat::Rgba || format == PixelFormat::Bgrx;
}

int minimum_row_bytes(PixelFormat format, int width) noexcept {
    return is_packed_rgb(format) ? width * 4 : width;
}

enum class ReceiveStatus {
    Packet,
    Again,
    Eof,
    Error,
};

struct ReceiveResult {
    ReceiveStatus status = ReceiveStatus::Error;
    AvPacketPtr packet;
    std::string error;
};

// Convert an H.264 elementary stream to Annex-B form.
//
// FFmpeg's libx264 packets are already Annex-B, so this is primarily a
// normalization helper for future encoders that may emit AVCC length-prefixed
// NAL units. Only 4-byte length prefixes are supported.
std::vector<std::byte> to_annex_b_h264(std::span<const std::byte> in) {
    if (in.empty()) {
        return {};
    }

    // Try to parse the input as AVCC. If the length-prefixed NALs consume
    // exactly the whole buffer, treat it as AVCC and rewrite it as Annex-B.
    if (in.size() >= 4) {
        std::vector<std::byte> out;
        out.reserve(in.size()); // Annex-B is the same size or larger.
        std::size_t offset = 0;
        bool ok = true;

        while (offset < in.size()) {
            if (offset + 4 > in.size()) {
                ok = false;
                break;
            }

            const auto* p = in.data() + offset;
            const uint32_t length = (static_cast<uint32_t>(static_cast<unsigned char>(p[0])) << 24) |
                                    (static_cast<uint32_t>(static_cast<unsigned char>(p[1])) << 16) |
                                    (static_cast<uint32_t>(static_cast<unsigned char>(p[2])) << 8) |
                                    (static_cast<uint32_t>(static_cast<unsigned char>(p[3])));

            // Guard against malformed or pathological length prefixes.
            if (length == 0 || length > std::numeric_limits<int>::max() / 2 || offset + 4 + length > in.size()) {
                ok = false;
                break;
            }

            out.push_back(std::byte{0x00});
            out.push_back(std::byte{0x00});
            out.push_back(std::byte{0x00});
            out.push_back(std::byte{0x01});
            out.insert(out.end(),
                       in.begin() + static_cast<std::ptrdiff_t>(offset + 4),
                       in.begin() + static_cast<std::ptrdiff_t>(offset + 4 + length));
            offset += 4 + length;
        }

        if (ok && offset == in.size()) {
            return out;
        }
    }

    return std::vector<std::byte>(in.begin(), in.end());
}

} // namespace

class FfmpegEncoder final : public Encoder {
  public:
    FfmpegEncoder(AvCodecContextPtr ctx, EncoderConfig config) : ctx_(std::move(ctx)), config_(std::move(config)) {}

    CodecResult<std::vector<EncodedFrame>> encode(const CapturedFrame& frame) override {
        auto input_result = make_input_frame(frame);
        if (is_codec_error(input_result)) {
            return CodecError{codec_error(input_result).message};
        }
        auto input = std::move(codec_value(input_result));

        if (force_keyframe_.exchange(false)) {
            input->pict_type = AV_PICTURE_TYPE_I;
        }

        std::vector<EncodedFrame> out;
        int send_ret = avcodec_send_frame(ctx_.get(), input.get());
        while (send_ret == AVERROR(EAGAIN)) {
            auto received = receive_one_packet();
            if (received.status == ReceiveStatus::Error) {
                return CodecError{received.error};
            }
            if (received.status == ReceiveStatus::Packet) {
                out.push_back(to_encoded_frame(received.packet.get(), frame.timestamp_ns));
            }
            if (received.status == ReceiveStatus::Again) {
                return CodecError{"encoder stalled before producing output"};
            }
            if (received.status == ReceiveStatus::Eof) {
                return CodecError{"encoder reached end of stream before accepting the frame"};
            }
            send_ret = avcodec_send_frame(ctx_.get(), input.get());
        }

        if (send_ret < 0) {
            return CodecError{ffmpeg_error(send_ret)};
        }

        if (auto error = drain_packets(frame.timestamp_ns, out)) {
            return CodecError{std::move(*error)};
        }

        return out;
    }

    CodecResult<std::vector<EncodedFrame>> flush() override {
        std::vector<EncodedFrame> out;
        int ret = avcodec_send_frame(ctx_.get(), nullptr);
        while (ret == AVERROR(EAGAIN)) {
            auto received = receive_one_packet();
            if (received.status == ReceiveStatus::Error) {
                return CodecError{received.error};
            }
            if (received.status == ReceiveStatus::Packet) {
                out.push_back(to_encoded_frame(received.packet.get(), 0));
            }
            if (received.status == ReceiveStatus::Again) {
                return CodecError{"encoder stalled during flush"};
            }
            if (received.status == ReceiveStatus::Eof) {
                // The encoder is already fully drained.
                return out;
            }
            ret = avcodec_send_frame(ctx_.get(), nullptr);
        }

        if (ret < 0) {
            return CodecError{ffmpeg_error(ret)};
        }

        if (auto error = drain_packets(0, out)) {
            return CodecError{std::move(*error)};
        }

        return out;
    }

    void request_keyframe() override {
        force_keyframe_ = true;
    }

    std::vector<std::byte> get_extradata() const override {
        if (ctx_->extradata == nullptr || ctx_->extradata_size <= 0) {
            return {};
        }
        return std::vector<std::byte>(as_bytes(ctx_->extradata), as_bytes(ctx_->extradata) + ctx_->extradata_size);
    }

  private:
    ReceiveResult receive_one_packet() {
        AvPacketPtr packet(av_packet_alloc(), AvPacketDeleter{});
        if (packet == nullptr) {
            return {ReceiveStatus::Error, nullptr, "failed to allocate AVPacket"};
        }

        int ret = avcodec_receive_packet(ctx_.get(), packet.get());
        if (ret == AVERROR(EAGAIN)) {
            return {ReceiveStatus::Again, nullptr, {}};
        }
        if (ret == AVERROR_EOF) {
            return {ReceiveStatus::Eof, nullptr, {}};
        }
        if (ret < 0) {
            return {ReceiveStatus::Error, nullptr, ffmpeg_error(ret)};
        }

        return {ReceiveStatus::Packet, std::move(packet), {}};
    }

    std::optional<std::string> drain_packets(uint64_t fallback_timestamp_ns, std::vector<EncodedFrame>& out) {
        while (true) {
            auto received = receive_one_packet();
            if (received.status == ReceiveStatus::Again || received.status == ReceiveStatus::Eof) {
                return std::nullopt;
            }
            if (received.status == ReceiveStatus::Error) {
                return received.error;
            }
            out.push_back(to_encoded_frame(received.packet.get(), fallback_timestamp_ns));
        }
    }

    EncodedFrame to_encoded_frame(const AVPacket* packet, uint64_t fallback_timestamp_ns) const {
        EncodedFrame frame;
        if (packet->pts != AV_NOPTS_VALUE) {
            frame.capture_timestamp_ns =
                static_cast<uint64_t>(av_rescale_q(packet->pts, ctx_->time_base, AVRational{1, 1'000'000'000}));
        } else {
            frame.capture_timestamp_ns = fallback_timestamp_ns;
        }
        frame.rtp_timestamp = rtp_timestamp_from_ns(frame.capture_timestamp_ns);
        frame.is_keyframe = (packet->flags & AV_PKT_FLAG_KEY) != 0;

        if (packet->data != nullptr && packet->size > 0) {
            frame.data = to_annex_b_h264(as_byte_span(packet->data, static_cast<std::size_t>(packet->size)));
        }

        return frame;
    }

    CodecResult<AvFramePtr> make_input_frame(const CapturedFrame& frame) const {
        if (frame.width <= 0 || frame.height <= 0) {
            return CodecError{"frame dimensions must be positive"};
        }

        const AVPixelFormat input_format = to_ffmpeg_format(frame.pixel_format);
        if (input_format == AV_PIX_FMT_NONE) {
            return CodecError{"unsupported pixel format"};
        }

        const int row_bytes = minimum_row_bytes(frame.pixel_format, frame.width);
        std::size_t stride = static_cast<std::size_t>(row_bytes);
        if (frame.stride != 0) {
            if (frame.stride < row_bytes) {
                return CodecError{"input stride is smaller than one pixel row"};
            }
            if (!is_packed_rgb(frame.pixel_format) && frame.stride != row_bytes) {
                return CodecError{"unsupported input stride for planar format"};
            }
            stride = static_cast<std::size_t>(frame.stride);
        }

        std::size_t required_size = 0;
        if (is_packed_rgb(frame.pixel_format)) {
            required_size = stride * (static_cast<std::size_t>(frame.height) - 1) + static_cast<std::size_t>(row_bytes);
        } else {
            required_size =
                static_cast<std::size_t>(av_image_get_buffer_size(input_format, frame.width, frame.height, 1));
        }
        if (frame.pixels.size() < required_size) {
            return CodecError{"input pixel buffer is too small"};
        }

        AvFramePtr output(av_frame_alloc(), AvFrameDeleter{});
        if (output == nullptr) {
            return CodecError{"failed to allocate AVFrame"};
        }

        // The output frame is at the encoder's configured dimensions (which
        // may be smaller than the capture when downscaling to the receiver's
        // display); sws_scale handles both the format conversion and the
        // resolution change in one pass.
        output->width = config_.width;
        output->height = config_.height;
        output->format = AV_PIX_FMT_YUV420P;
        output->time_base = ctx_->time_base;
        output->pts =
            av_rescale_q(static_cast<int64_t>(frame.timestamp_ns), AVRational{1, 1'000'000'000}, ctx_->time_base);

        if (av_frame_get_buffer(output.get(), 0) < 0) {
            return CodecError{"failed to allocate AVFrame buffer"};
        }

        if (!ensure_scaler(frame.width, frame.height, input_format)) {
            return CodecError{"failed to create swscale context"};
        }

        std::array<const uint8_t*, 4> src{nullptr, nullptr, nullptr, nullptr};
        std::array<int, 4> src_lines{0, 0, 0, 0};
        if (is_packed_rgb(frame.pixel_format)) {
            // Packed RGB is a single plane whose row pitch may be padded beyond
            // the packed row size, so build the source arrays by hand.
            src[0] = as_u8(frame.pixels.data());
            src_lines[0] = static_cast<int>(stride);
        } else {
            std::array<uint8_t*, 4> planar_src{nullptr, nullptr, nullptr, nullptr};
            std::array<int, 4> planar_lines{0, 0, 0, 0};
            if (av_image_fill_arrays(planar_src.data(),
                                     planar_lines.data(),
                                     as_u8(frame.pixels.data()),
                                     input_format,
                                     frame.width,
                                     frame.height,
                                     1) < 0) {
                return CodecError{"failed to fill input pixel arrays"};
            }
            std::copy(planar_src.begin(), planar_src.end(), src.begin());
            std::copy(planar_lines.begin(), planar_lines.end(), src_lines.begin());
        }

        if (sws_scale(scaler_.get(), src.data(), src_lines.data(), 0, frame.height, output->data, output->linesize) <=
            0) {
            return CodecError{"failed to convert input frame to YUV420P"};
        }

        return output;
    }

    bool ensure_scaler(int width, int height, AVPixelFormat input_format) const {
        if (scaler_ != nullptr && scaler_input_width_ == width && scaler_input_height_ == height &&
            scaler_input_format_ == input_format) {
            return true;
        }

        // Scale to the encoder's configured output (which may be smaller
        // than the input when downscaling to the receiver's display).
        scaler_.reset(sws_getContext(width,
                                     height,
                                     input_format,
                                     config_.width,
                                     config_.height,
                                     AV_PIX_FMT_YUV420P,
                                     SWS_BILINEAR,
                                     nullptr,
                                     nullptr,
                                     nullptr));
        if (scaler_ == nullptr) {
            return false;
        }
        scaler_input_width_ = width;
        scaler_input_height_ = height;
        scaler_input_format_ = input_format;
        return true;
    }

    AvCodecContextPtr ctx_;
    EncoderConfig config_;
    std::atomic<bool> force_keyframe_{false};

    mutable SwsContextPtr scaler_;
    mutable int scaler_input_width_ = 0;
    mutable int scaler_input_height_ = 0;
    mutable AVPixelFormat scaler_input_format_ = AV_PIX_FMT_NONE;
};

CodecResult<std::unique_ptr<Encoder>> EncoderFactory::create(const EncoderConfig& config) {
    if (config.width <= 0 || config.height <= 0) {
        return CodecError{"encoder width and height must be positive"};
    }
    if (config.frame_rate_num <= 0 || config.frame_rate_den <= 0) {
        return CodecError{"encoder frame rate must be positive"};
    }
    if (config.bitrate_kbps <= 0) {
        return CodecError{"encoder VBV max bitrate must be positive"};
    }
    if (config.crf < 0 || config.crf > 51) {
        return CodecError{"encoder CRF must be between 0 and 51"};
    }
    if (config.codec_name != "h264" && config.codec_name != "libx264") {
        return CodecError{"only h264 is supported in phase 2"};
    }
    if (config.hardware_accel) {
        return CodecError{"hardware acceleration is not implemented in phase 2"};
    }

    const AVCodec* codec = avcodec_find_encoder_by_name("libx264");
    if (codec == nullptr) {
        return CodecError{"libx264 encoder not found"};
    }

    AvCodecContextPtr ctx(avcodec_alloc_context3(codec), AvCodecContextDeleter{});
    if (ctx == nullptr) {
        return CodecError{"failed to allocate encoder context"};
    }

    ctx->codec_type = AVMEDIA_TYPE_VIDEO;
    ctx->width = config.width;
    ctx->height = config.height;
    // Microsecond resolution: the capture rate varies at runtime (monitor
    // refresh) and must not be quantized to the configured frame rate, or
    // consecutive frames get duplicate timestamps.
    ctx->time_base = AVRational{1, 1'000'000};
    ctx->framerate = AVRational{config.frame_rate_num, config.frame_rate_den};
    ctx->pix_fmt = AV_PIX_FMT_YUV420P;

    // CRF rate control: target a constant visual quality level instead of
    // a fixed bitrate, and let the encoder use fewer bits on static screen
    // content and more on motion or text. The VBV max rate still caps the
    // peak so bursts cannot overflow the receiver's UDP buffers.
    // CRF rate control: set x264's CRF directly as a private option.
    // FFmpeg's global_quality + AV_CODEC_FLAG_QSCALE path divides by
    // FF_QP2LAMBDA, giving a wrong CRF value (16/118 ≈ 0.1, essentially
    // lossless) — hence the "-qscale is ignored" warning. Setting the
    // private "crf" option bypasses that and gives x264 the exact value.
    const std::string crf_value = std::to_string(config.crf);
    if (av_opt_set(ctx->priv_data, "crf", crf_value.c_str(), 0) < 0) {
        return CodecError{"failed to set CRF quality"};
    }
    ctx->rc_max_rate = static_cast<int64_t>(config.bitrate_kbps) * 1000;
    // VBV: one frame period of budget keeps bursts tight — a two-frame
    // buffer lets a keyframe spike beyond what a constrained link can
    // absorb in real time, causing packet loss that cascades into PLI
    // storms. The tighter buffer trades a small quality dip on keyframes
    // for much better behavior on slow paths.
    ctx->rc_buffer_size = static_cast<int>(ctx->rc_max_rate / config.frame_rate_num);

    // A long GOP saves the keyframe overhead for screen content (which
    // changes incrementally); PLI feedback recovers from loss within one
    // frame time regardless of the GOP length.
    ctx->gop_size = config.frame_rate_num * 5;
    ctx->max_b_frames = 0;
    ctx->thread_count = 1;
    ctx->profile = AV_PROFILE_H264_MAIN;
    // No GLOBAL_HEADER: SPS/PPS are repeated in-band at every keyframe so a
    // receiver that joins mid-stream (or recovers after loss) can decode
    // without out-of-band parameter negotiation.
    ctx->flags |= AV_CODEC_FLAG_LOW_DELAY;

    if (av_opt_set(ctx->priv_data, "preset", "faster", 0) < 0) {
        return CodecError{"failed to set libx264 preset"};
    }
    if (av_opt_set(ctx->priv_data, "tune", "zerolatency", 0) < 0) {
        return CodecError{"failed to set libx264 tune"};
    }
    if (av_opt_set(ctx->priv_data, "forced-idr", "1", 0) < 0) {
        return CodecError{"failed to enable forced IDR keyframes"};
    }
    // Screen-content tuning: auto-variance AQ allocates bits away from
    // flat areas and toward text edges; higher psy-rd preserves texture
    // sharpness at the cost of slight rate efficiency.
    if (av_opt_set(ctx->priv_data, "aq-mode", "2", 0) < 0) {
        return CodecError{"failed to set adaptive quantization mode"};
    }
    if (av_opt_set(ctx->priv_data, "psy-rd", "1.5", 0) < 0) {
        return CodecError{"failed to set psychovisual rate-distortion strength"};
    }

    int open_ret = avcodec_open2(ctx.get(), codec, nullptr);
    if (open_ret < 0) {
        return CodecError{std::string{"failed to open libx264 encoder: "} + ffmpeg_error(open_ret)};
    }

    return std::make_unique<FfmpegEncoder>(std::move(ctx), config);
}

} // namespace sc