#define AL_LOG_SECTION "frame_queue_libplacebo"
#include <al/log.h>
#include <al/lib.h>

#ifdef CAMU_HAVE_FFMPEG
#include "../codec/ffmpeg/common.h"
#endif

#include "queue_libplacebo.h"

static bool queue_lp_configure_subtitles(struct camu_frame_queue *queue, u32 width, u32 height,
    struct camu_codec_stream *stream)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)queue;
#ifdef CAMU_HAVE_SUBTITLES
    ASS_Library *ass = lq->subs.ass;
    if (!(lq->subs.renderer = ass_renderer_init(ass))) {
        log_error("Failed to initialize ass renderer.");
        return false;
    }
    if (!(lq->subs.track = ass_new_track(ass))) {
        log_error("Failed to create ass track.");
        return false;
    }
    ASS_Renderer *renderer = lq->subs.renderer;
    ASS_Track *track = lq->subs.track;
    ass_set_shaper(renderer, ASS_SHAPING_COMPLEX);
    ass_set_frame_size(renderer, width, height);
    ass_set_storage_size(renderer, width, height);
    ass_set_fonts(renderer, NULL, NULL, ASS_FONTPROVIDER_AUTODETECT, NULL, 0);
    ass_set_hinting(renderer, ASS_HINTING_NONE);
    AVCodecParameters *codecpar = stream->av.stream->codecpar;
    lq->subs.conversion_needed = codecpar->codec_id != AV_CODEC_ID_ASS;
    if (lq->subs.conversion_needed) {
        const AVCodecDescriptor *desc = avcodec_descriptor_get(codecpar->codec_id);
        if (!(desc->props & AV_CODEC_PROP_TEXT_SUB)) {
            log_error("Cannot convert non-text (bitmap) subtitles.");
            return false;
        }
        const AVCodec *codec = avcodec_find_decoder(codecpar->codec_id);
        if (!codec) {
            log_error("Failed to find subtitle transcoder.");
            return false;
        }
        if (!(lq->subs.converter = camu_ff_alloc_codec_context(codec, codecpar))) {
            return false;
        }
        AVCodecContext *converter = lq->subs.converter;
        converter->sub_charenc_mode = FF_SUB_CHARENC_MODE_IGNORE;
        if (!camu_ff_open_avcodec(converter, codec, NULL)) {
            return false;
        }
        const AVCodecDescriptor *ass_desc = avcodec_descriptor_get(AV_CODEC_ID_ASS);
        log_info("Subtitle conversion: %s -> %s.", camu_ff_longest_name(desc), camu_ff_longest_name(ass_desc));
        // https://github.com/mpv-player/mpv/blob/266cb79f38fd1a5fd448b453dee5971795a145ca/sub/sd_ass.c#L634
#define MP_ASS_FONT_PLAYRESX 384.0
        track->PlayResX = track->PlayResY * (f64)width / MAX(height, (u32)1);
        f64 fix_margins = track->PlayResX / MP_ASS_FONT_PLAYRESX;
        f64 font_scale = 1.0;
        for (s32 n = 0; n < track->n_styles; n++) {
            track->styles[n].MarginL = track->styles[n].MarginL * fix_margins;
            track->styles[n].MarginR = track->styles[n].MarginR * fix_margins;
            track->styles[n].MarginV = track->styles[n].MarginV * font_scale;
        }
        ass_process_codec_private(track, (const char *)converter->subtitle_header,
            converter->subtitle_header_size);
    } else {
        ass_process_codec_private(track, (const char *)codecpar->extradata,
            codecpar->extradata_size);
    }
    nn_mutex_init(&lq->subs.lock);
    lq->subs.are_present = true;
    return true;
#else
    (void)lq;
    (void)width;
    (void)height;
    (void)stream;
    return false;
#endif
}

static bool map_frame(pl_gpu gpu, pl_tex *tex, const struct pl_source_frame *src,
    struct pl_frame *out_frame)
{
    struct camu_codec_frame *frame = (struct camu_codec_frame *)src->frame_data;

    u64 masks[4] = { 0 };
    s32 pixel_stride;
    switch (frame->format) {
    case CAMU_PIXEL_FORMAT_RGBA:
        masks[0] = 0x000000ff;
        masks[1] = 0x0000ff00;
        masks[2] = 0x00ff0000;
        masks[3] = 0xff000000;
        pixel_stride = 4;
        break;
    case CAMU_PIXEL_FORMAT_RGB24:
        masks[0] = 0x000000ff;
        masks[1] = 0x0000ff00;
        masks[2] = 0x00ff0000;
        pixel_stride = 3;
        break;
    case CAMU_PIXEL_FORMAT_GREY:
        masks[0] = 0x000000ff;
        pixel_stride = 1;
        break;
    default:
        log_error("Unsuppored pixel format.");
        al_assert_and_return(false);
    }

    if (frame->format == CAMU_PIXEL_FORMAT_GREY) {
        out_frame->repr.sys = PL_COLOR_SYSTEM_BT_601; // pl_system_from_av(AVCOL_SPC_BT470BG)
        out_frame->repr.levels = PL_COLOR_LEVELS_FULL; // pl_levels_from_av(AVCOL_RANGE_JPEG)
    } else {
        out_frame->repr = pl_color_repr_unknown;
    }
    out_frame->color = pl_color_space_unknown;
    out_frame->crop = (pl_rect2df){ 0.f, 0.f, frame->video.width, frame->video.height };

    struct pl_plane_data data[4] = { 0 };
    data[0].type = PL_FMT_UNORM;
    data[0].width = frame->video.width;
    data[0].height = frame->video.height;
    data[0].pixel_stride = pixel_stride;
    data[0].row_stride = frame->video.width * pixel_stride;
    data[0].pixels = frame->data;
    pl_plane_data_from_mask(&data[0], masks);

    struct pl_plane plane = { 0 };
    bool ok = pl_upload_plane(gpu, &plane, tex, &data[0]);

    camu_codec_frame_discard(frame);

    if (!ok) {
        log_error("Failed to upload texture.");
        return false;
    }

    out_frame->num_planes = 1;
    out_frame->planes[0] = plane;

    return true;
}

static void unmap_frame(pl_gpu gpu, struct pl_frame *frame, const struct pl_source_frame *src)
{
    (void)gpu;
    (void)frame;
    (void)src;
}

static void discard_frame(const struct pl_source_frame *src)
{
    struct camu_codec_frame *frame = (struct camu_codec_frame *)src->frame_data;
    camu_codec_frame_discard(frame);
}

static void queue_lp_push(struct camu_frame_queue *queue, struct camu_codec_frame *frame, f64 pts, u32 *count)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)queue;
    pl_queue_push(lq->queue, &(struct pl_source_frame){
        .pts = pts,
        .duration = 0.0,
        .map = map_frame,
        .unmap = unmap_frame,
        .discard = discard_frame,
        .frame_data = frame
    });
    *count = (u32)pl_queue_num_frames(lq->queue);
}

#ifdef CAMU_HAVE_FFMPEG
#ifdef CAMU_HAVE_SUBTITLES
static void free_overlay(pl_gpu gpu, struct camu_overlay *overlay)
{
    al_assert(!overlay->ref);
    struct pl_overlay *current = CAMU_OVERLAY_OFFSET(overlay, 0);
    for (u32 i = 0; i < overlay->alloc; i++) {
        if (current->tex) {
            pl_tex_destroy(gpu, &current->tex);
        }
        current++;
    }
    al_free(overlay);
}

static struct camu_overlay *first_nonref_overlay(struct camu_frame_queue_lp *lq, u32 *index)
{
    struct camu_overlay *overlay;
    al_array_foreach(lq->subs.overlays, i, overlay) {
        if (!overlay->ref) {
            *index = i;
            return overlay;
        }
    }
    return NULL;
}

static struct camu_overlay *create_subtitle_overlay(struct camu_overlay *prev, pl_gpu gpu, ASS_Image *ass_frame)
{
    u16 frame_count = 1;
    ASS_Image *head_ass_frame = ass_frame;
    while (ass_frame->next) { al_assert(frame_count < UINT16_MAX); frame_count++; ass_frame = ass_frame->next; }
    ass_frame = head_ass_frame;

    u16 prev_alloc = prev ? prev->alloc : 0;
    struct camu_overlay *overlay = camu_overlay_alloc(prev, frame_count,
        sizeof(struct pl_overlay) + sizeof(struct pl_overlay_part));
    if (frame_count > prev_alloc) {
        // Zero any newly allocated pl_overlays.
        al_memset(CAMU_OVERLAY_OFFSET(overlay, prev_alloc * sizeof(struct pl_overlay)), 0,
            (frame_count - prev_alloc) * sizeof(struct pl_overlay));
    }

    // The list of pl_overlay's need to be contiguous in memory.
    struct pl_overlay *current = CAMU_OVERLAY_OFFSET(overlay, 0);
    struct pl_overlay_part *current_part = CAMU_OVERLAY_OFFSET(overlay,
        sizeof(struct pl_overlay) * overlay->alloc);
    overlay->num = 0;

    for (; ass_frame; ass_frame = ass_frame->next) {
        bool ok = pl_tex_recreate(gpu, &current->tex, pl_tex_params(
            .w = ass_frame->w,
            .h = ass_frame->h,
            .format = pl_find_named_fmt(gpu, "r8"),
            .sampleable = true,
            .host_writable = true
        ));
        if (ok) {
            ok = pl_tex_upload(gpu, pl_tex_transfer_params(
                .tex = current->tex,
                .row_pitch = ass_frame->stride,
                .ptr = ass_frame->bitmap
            ));
        }
        if (!ok) {
            frame_count--;
            log_error("Failed to upload subtitle bitmap.");
            continue;
        }

        // https://github.com/mpv-player/mpv/blob/70aaba71d6e3071a732069a1d222d1eb4293faf2/video/out/vo_gpu_next.c#L289
        // https://github.com/mpv-player/mpv/blob/70aaba71d6e3071a732069a1d222d1eb4293faf2/sub/ass_mp.c#L195
        current->color = *pl_color_space(
            .primaries = PL_COLOR_PRIM_BT_709,
            .transfer = PL_COLOR_TRC_SRGB
        );
        current->mode = PL_OVERLAY_MONOCHROME;
        current->repr = pl_color_repr_unknown;
        current->repr.alpha = PL_ALPHA_INDEPENDENT;

        current->coords = PL_OVERLAY_COORDS_SRC_FRAME;
        current_part->src = (pl_rect2df){ 0.f, 0.f, ass_frame->w, ass_frame->h };
        current_part->dst = (pl_rect2df){
            ass_frame->dst_x,
            ass_frame->dst_y,
            ass_frame->dst_x + ass_frame->w,
            ass_frame->dst_y + ass_frame->h
        };
        u32 c = ass_frame->color;
        current_part->color[0] = (c >> 24) / 255.f;
        current_part->color[1] = ((c >> 16) & 0xff) / 255.f;
        current_part->color[2] = ((c >> 8) & 0xff) / 255.f;
        current_part->color[3] = 1.f - (c & 0xff) / 255.f;

        current->parts = current_part;
        current->num_parts = 1;

        overlay->num++;

        current++;
        current_part++;
    }

    al_assert((u32)overlay->num == frame_count);

    return overlay;
}
#endif

static bool map_av_frame(pl_gpu gpu, pl_tex *tex, const struct pl_source_frame *src,
    struct pl_frame *out_frame)
{
    AVFrame *frame = (AVFrame *)src->frame_data;
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)frame->opaque;
    AVStream *stream = lq->q.buf->stream->av.stream;

    bool ok = false;

    if (!lq->copy_frame_fallback) {
        ok = pl_map_avframe_ex(gpu, out_frame, pl_avframe_params(
            .frame = frame,
            .tex = tex
        ));
    }

    // https://github.com/streetpea/chiaki-ng/blob/dd145ec2c802814e98fb833d3a9692169aa67315/gui/src/qmlmainwindow.cpp#L860
    // https://github.com/streetpea/chiaki-ng/blob/dd145ec2c802814e98fb833d3a9692169aa67315/gui/src/qmlbackend.cpp#L786
    // pl_map_avframe_derived() will fail under Xwayland because of a lack of EGL_EXT_image_dma_buf_import.
    if (!ok && frame->hw_frames_ctx) {
        if (!lq->copy_frame_fallback) {
            log_warn("Failed to map hwframe, falling back to software copy.");
            lq->copy_frame_fallback = true;
            lq->swframe = av_frame_alloc();
            enum AVPixelFormat *fmts;
            av_hwframe_transfer_get_formats(frame->hw_frames_ctx, AV_HWFRAME_TRANSFER_DIRECTION_FROM, &fmts, 0);
            lq->swframe->format = fmts[0];
            av_free(fmts);
        }
        s32 ret = av_hwframe_transfer_data(lq->swframe, frame, 0);
        ok = ret == 0;
        if (ok) {
            av_frame_copy_props(lq->swframe, frame);
            ok = pl_map_avframe_ex(gpu, out_frame, pl_avframe_params(
                .frame = lq->swframe,
                .tex = tex
            ));
            av_frame_unref(lq->swframe);
        } else {
            log_error("Failed to transfer hwframe from the device (%s).", av_err2str(ret));
        }
    }

    ((struct pl_source_frame *)src)->frame_data = NULL;

#ifdef CAMU_HAVE_SUBTITLES
    if (ok && lq->subs.are_present) {
        s64 now = av_rescale_q(frame->best_effort_timestamp, stream->time_base, (AVRational){ 1, 1000 });
        s32 change = -1; // 1 = different position, 2 = different content.
        nn_mutex_lock(&lq->subs.lock);
        ASS_Image *ass_frame = ass_render_frame(lq->subs.renderer, lq->subs.track, now, &change);
        al_assert(change >= 0);
        // ass_frame = NULL and lq->subs.track->n_events = 0 can also mean parsing completely failed.
        if (ass_frame) {
            struct camu_overlay *overlay = lq->subs.last;
            if (change != 0) {
                u32 reused = AL_ARRAY_NO_INDEX;
                overlay = create_subtitle_overlay(first_nonref_overlay(lq, &reused), gpu, ass_frame);
                // If all texture uploads failed, overlay->num will be 0.
                lq->subs.last = overlay;
                if (overlay) {
                    if (reused != AL_ARRAY_NO_INDEX) {
                        // Pointer may have changed due to reallocation.
                        al_array_at(lq->subs.overlays, reused) = overlay;
                    } else { // New overlay object.
                        overlay->lock = &lq->subs.lock;
                        al_array_push(lq->subs.overlays, overlay);
                    }
                }
            }
            if (overlay) {
                overlay->ref++;
                out_frame->overlays = CAMU_OVERLAY_OFFSET(overlay, 0);
                ((struct pl_source_frame *)src)->frame_data = overlay;
            } else {
               out_frame->overlays = NULL;
            }
        }
        nn_mutex_unlock(&lq->subs.lock);
    }
#endif

    av_frame_free(&frame);

    if (!ok) {
        log_error("Failed to map AVFrame.");
        return false;
    }

    pl_frame_copy_stream_props(out_frame, stream);

    return true;
}

static void unmap_av_frame(pl_gpu gpu, struct pl_frame *frame, const struct pl_source_frame *src)
{
    if (src->frame_data) {
        struct camu_overlay *overlay = (struct camu_overlay *)src->frame_data;
        nn_mutex_lock(overlay->lock);
        overlay->ref--;
        nn_mutex_unlock(overlay->lock);
    }
    pl_unmap_avframe(gpu, frame);
}

static void discard_av_frame(const struct pl_source_frame *src)
{
    AVFrame *frame = (AVFrame *)src->frame_data;
    av_frame_free(&frame);
    log_warn("Dropped frame with PTS %.3f.", src->pts);
}

static void queue_lp_push_av_frame(struct camu_frame_queue *queue, AVFrame *frame, f64 pts, u32 *count)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)queue;
    AVStream *stream = lq->q.buf->stream->av.stream;
    frame->opaque = lq;
    pl_queue_push(lq->queue, &(struct pl_source_frame){
        .pts = pts,
        .duration = camu_ff_frame_duration(frame) * av_q2d(stream->time_base),
        .map = map_av_frame,
        .unmap = unmap_av_frame,
        .discard = discard_av_frame,
        .frame_data = frame,
        .first_field = pl_field_from_avframe(frame)
    });
    *count = (u32)pl_queue_num_frames(lq->queue);
}
#endif

#ifdef CAMU_HAVE_SUBTITLES
static bool convert_subtitle(struct camu_frame_queue_lp *lq, AVPacket *pkt)
{
    AVSubtitle sub = { 0 };
    s32 got_sub = 0;
    s32 ret = avcodec_decode_subtitle2(lq->subs.converter, &sub, &got_sub, pkt);
    if (ret < 0) {
        log_error("Failed to convert subtitle (%s).", av_err2str(ret));
    } else if (got_sub) {
        for (u32 i = 0; i < sub.num_rects; i++) {
            AVSubtitleRect *rect = sub.rects[i];
            size_t size = al_strlen(rect->ass);
            if (rect->ass[size - 1] == 0x02) size--; // Not sure what this is...
            ass_process_chunk(lq->subs.track, rect->ass, size, pkt->pts, pkt->duration);
        }
        avsubtitle_free(&sub);
        return true;
    }
    return false;
}
#endif

static void queue_lp_push_subtitle(struct camu_frame_queue *queue, struct camu_codec_packet *packet)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)queue;
#ifdef CAMU_HAVE_SUBTITLES
    if (!lq->subs.are_present) return;
    AVPacket *pkt = packet->av.pkt;
    nn_mutex_lock(&lq->subs.lock);
    al_assert(lq->subs.track && lq->subs.renderer);
    if (lq->subs.conversion_needed) {
        convert_subtitle(lq, pkt);
    } else {
        ass_process_chunk(lq->subs.track, (const char *)pkt->data, pkt->size, pkt->pts, pkt->duration);
    }
    nn_mutex_unlock(&lq->subs.lock);
#else
    (void)lq;
    (void)packet;
#endif
}

static void queue_lp_flush(struct camu_frame_queue *queue, u32 *count)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)queue;
    *count = (u32)pl_queue_num_frames(lq->queue);
    pl_queue_push(lq->queue, NULL);
#ifdef CAMU_HAVE_SUBTITLES
    if (lq->subs.are_present) {
        if (lq->subs.conversion_needed) {
            AVPacket *empty_pkt = av_packet_alloc();
            while (convert_subtitle(lq, empty_pkt)) { }
            av_packet_free(&empty_pkt);
        }
        ass_flush_events(lq->subs.track);
    }
#endif
}

static u8 queue_lp_read(struct camu_frame_queue *queue, f64 pts, void *out, u32 *count)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)queue;
    struct pl_frame_mix *mix = (struct pl_frame_mix *)out;
    lq->params.pts = pts;
    enum pl_queue_status ret = pl_queue_update(lq->queue, mix, &lq->params);
    *count = (u32)pl_queue_num_frames(lq->queue);
    switch (ret) {
    case PL_QUEUE_OK:
        break;
    case PL_QUEUE_MORE:
        return CAMU_QUEUE_MORE;
    case PL_QUEUE_EOF:
        return CAMU_QUEUE_EOF;
    case PL_QUEUE_ERR:
    default:
        return CAMU_QUEUE_ERR;
    }
    return CAMU_QUEUE_OK;
}

static void queue_lp_reset(struct camu_frame_queue *queue)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)queue;
    u32 count;
    queue_lp_flush(&lq->q, &count);
    pl_queue_reset(lq->queue);
}

static void queue_lp_free(struct camu_frame_queue **queue)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)*queue;
    pl_queue_destroy(&lq->queue);
#ifdef CAMU_HAVE_FFMPEG
    if (lq->copy_frame_fallback) {
        av_frame_free(&lq->swframe);
    }
#endif
#ifdef CAMU_HAVE_SUBTITLES
    struct camu_overlay *overlay;
    al_array_foreach_rev(lq->subs.overlays, i, overlay) {
        free_overlay(lq->gpu, overlay);
    }
    al_array_free(lq->subs.overlays);
    if (lq->subs.are_present) {
        nn_mutex_destroy(&lq->subs.lock);
    }
    if (lq->subs.converter) {
        avcodec_free_context(&lq->subs.converter);
    }
    if (lq->subs.track) {
        ass_free_track(lq->subs.track);
    }
    if (lq->subs.renderer) {
        ass_renderer_done(lq->subs.renderer);
    }
#endif
    al_free(lq);
    *queue = NULL;
}

struct camu_frame_queue *camu_frame_queue_lp_create(void)
{
    struct camu_frame_queue_lp *lq = al_alloc_object(struct camu_frame_queue_lp);
#ifdef CAMU_HAVE_SUBTITLES
    lq->subs.are_present = false;
    lq->subs.last = NULL;
    al_array_init(lq->subs.overlays);
#endif
    lq->q.configure_subtitles = queue_lp_configure_subtitles;
    lq->q.push = queue_lp_push;
#ifdef CAMU_HAVE_FFMPEG
    lq->q.push_av_frame = queue_lp_push_av_frame;
#endif
    lq->q.push_subtitle = queue_lp_push_subtitle;
    lq->q.flush = queue_lp_flush;
    lq->q.read = queue_lp_read;
    lq->q.reset = queue_lp_reset;
    lq->q.free = queue_lp_free;
    return (struct camu_frame_queue *)lq;
}