#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
    if (!lq->ass) return false;
    lq->ass_renderer = ass_renderer_init(lq->ass);
    if (!lq->ass_renderer) {
        log_error("Failed to initialize ass renderer.");
        return false;
    }
    ass_set_shaper(lq->ass_renderer, ASS_SHAPING_COMPLEX);
    lq->ass_track = ass_new_track(lq->ass);
    if (!lq->ass_track) {
        log_error("Failed to create ass track.");
        return false;
    }
    ass_set_frame_size(lq->ass_renderer, width, height);
    ass_set_storage_size(lq->ass_renderer, width, height);
    ass_set_fonts(lq->ass_renderer, NULL, NULL, ASS_FONTPROVIDER_AUTODETECT, NULL, 0);
    ass_set_hinting(lq->ass_renderer, ASS_HINTING_NONE);
    AVCodecParameters *codecpar = stream->av.stream->codecpar;
    ass_process_codec_private(lq->ass_track, (const char *)codecpar->extradata,
        codecpar->extradata_size);
    nn_mutex_init(&lq->subtitle_lock);
    lq->has_subtitles = 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_GREYA:
        masks[0] = 0x000000ff;
        masks[1] = 0x0000ff00;
        pixel_stride = 2;
        break;
    case CAMU_PIXEL_FORMAT_GREY:
        masks[0] = 0x000000ff;
        pixel_stride = 1;
        break;
    default:
        log_error("Unsuppored pixel format.");
        al_assert_and_return(false);
    }

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

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

    out_frame->num_planes = 1;
    out_frame->planes[0] = plane;
    out_frame->repr = pl_color_repr_unknown;
    out_frame->color = pl_color_space_monitor;
    out_frame->crop = (pl_rect2df){ 0.f, 0.f, frame->video.width, frame->video.height };
    //out_frame->profile =

    camu_codec_frame_discard(frame);

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

    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)
{
    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
    });
}

#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->subtitles, 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.0;
        current_part->color[1] = ((c >> 16) & 0xff) / 255.0;
        current_part->color[2] = ((c >> 8) & 0xff) / 255.0;
        current_part->color[3] = 1.0 - (c & 0xff) / 255.0;

        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("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->has_subtitles) {
        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->subtitle_lock);
        ASS_Image *ass_frame = ass_render_frame(lq->ass_renderer, lq->ass_track, now, &change);
        al_assert(change >= 0);
        // ass_frame = NULL and lq->ass_track->n_events = 0 can also mean parsing completely failed.
        if (ass_frame) {
            struct camu_overlay *overlay = lq->last_subtitle;
            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->last_subtitle = overlay;
                if (overlay) {
                    if (reused != AL_ARRAY_NO_INDEX) {
                        // Pointer may have changed due to reallocation.
                        al_array_at(lq->subtitles, reused) = overlay;
                    } else { // New overlay object.
                        overlay->lock = &lq->subtitle_lock;
                        al_array_push(lq->subtitles, 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->subtitle_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)
{
    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)
    });
}
#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
    AVPacket *pkt = packet->av.pkt;
    nn_mutex_lock(&lq->subtitle_lock);
    al_assert(lq->ass_track && lq->ass_renderer);
    ass_process_chunk(lq->ass_track, (const char *)pkt->data, pkt->size, pkt->pts, pkt->duration);
    nn_mutex_unlock(&lq->subtitle_lock);
#else
    (void)lq;
    (void)packet;
#endif
}

static void queue_lp_flush(struct camu_frame_queue *queue)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)queue;
    pl_queue_push(lq->queue, NULL);
}

static s32 queue_lp_count(struct camu_frame_queue *queue)
{
    struct camu_frame_queue_lp *lq = (struct camu_frame_queue_lp *)queue;
    return pl_queue_num_frames(lq->queue);
}

static u8 queue_lp_read(struct camu_frame_queue *queue, f64 pts, void *out)
{
    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;
    switch (pl_queue_update(lq->queue, mix, &lq->params)) {
    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;
    pl_queue_push(lq->queue, NULL);
    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
    if (lq->has_subtitles) {
        struct camu_overlay *overlay;
        al_array_foreach_rev(lq->subtitles, i, overlay) {
            free_overlay(lq->gpu, overlay);
        }
        al_array_free(lq->subtitles);
        ass_free_track(lq->ass_track);
        ass_renderer_done(lq->ass_renderer);
        nn_mutex_destroy(&lq->subtitle_lock);
    }
#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->last_subtitle = NULL;
    al_array_init(lq->subtitles);
    lq->has_subtitles = false;
#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.count = queue_lp_count;
    lq->q.read = queue_lp_read;
    lq->q.reset = queue_lp_reset;
    lq->q.free = queue_lp_free;
    return (struct camu_frame_queue *)lq;
}