#include <spng.h>

#include "spng.h"
#include "common.h"

#ifdef LIANA_SERVER
#include "../cache/entry.h"

static bool spng_demuxer_init(struct camu_demuxer *demux, struct cch_handle *handle)
{
    struct camu_spng_demuxer *spng = (struct camu_spng_demuxer *)demux;

    spng->handle = handle;
    nn_buffer_init(&spng->buffer);
    spng->eof = false;

    cch_handle_seek(spng->handle, 0, SEEK_SET);
    struct cch_backing *backing = spng->handle->entry->backing;
    off_t size = backing->get_size_if_known(backing);
    nn_buffer_ensure_space(&spng->buffer, size);
    cch_handle_read(spng->handle, nn_buffer_get_ptr(&spng->buffer, 0), size);
    spng->buffer.size = size;

    spng_ctx *ctx = spng_ctx_new(SPNG_CTX_IGNORE_ADLER32);
    spng_set_png_buffer(ctx, nn_buffer_get_ptr(&spng->buffer, 0), spng->buffer.size);
    spng_set_option(ctx, SPNG_CHUNK_COUNT_LIMIT, 6250);
    spng_set_crc_action(ctx, SPNG_CRC_DISCARD, SPNG_CRC_DISCARD);

    struct spng_ihdr ihdr;
    spng_get_ihdr(ctx, &ihdr);

    struct camu_codec_stream stream = { 0 };
    stream.mode = CAMU_NORMAL;
    stream.type = CAMU_STREAM_VIDEO;
    struct camu_video_format *fmt = &stream.video.fmt;
    fmt->width = ihdr.width;
    fmt->height = ihdr.height;
    fmt->format = CAMU_PIXEL_FORMAT_RGBA;

    al_array_push(spng->demux.streams, stream);

    spng_ctx_free(ctx);

    return true;
}

static s32 spng_demuxer_get_packet(struct camu_demuxer *demux, struct camu_codec_packet *packet)
{
    struct camu_spng_demuxer *spng = (struct camu_spng_demuxer *)demux;
    if (!spng->eof) {
        packet->mode = CAMU_NORMAL;
        packet->buffer = &spng->buffer;
        spng->eof = true;
        return CAMU_OK;
    }
    return CAMU_ERR_EOF;
}

static u64 spng_demuxer_get_duration(struct camu_demuxer *demux)
{
    (void)demux;
    return 0;
}

static bool spng_demuxer_seek(struct camu_demuxer *demux, u64 pos, bool bytes)
{
    (void)demux;
    (void)pos;
    (void)bytes;
    return true;
}

static void spng_demuxer_free(struct camu_demuxer **demux)
{
    struct camu_spng_demuxer *spng = (struct camu_spng_demuxer *)*demux;
    nn_buffer_free(&spng->buffer);
    al_array_free(spng->demux.streams);
    al_free(spng);
    *demux = NULL;
}

struct camu_demuxer *camu_spng_demuxer_create(void)
{
    struct camu_spng_demuxer *spng = al_alloc_object(struct camu_spng_demuxer);
    spng->demux.mode = CAMU_NORMAL;
    al_array_init(spng->demux.streams);
    spng->demux.init = spng_demuxer_init;
    spng->demux.get_packet = spng_demuxer_get_packet;
    spng->demux.get_duration = spng_demuxer_get_duration;
    spng->demux.seek = spng_demuxer_seek;
    spng->demux.free = spng_demuxer_free;
    return (struct camu_demuxer *)spng;
}
#endif

#ifdef LIANA_CLIENT
static bool spng_decoder_init(struct camu_decoder *dec, struct camu_renderer *renderer, struct camu_codec_stream *stream,
    void (*callback)(void *, struct camu_codec_frame *), void *userdata)
{
    struct camu_spng_decoder *spng = (struct camu_spng_decoder *)dec;
    (void)renderer;
    spng->stream = stream;
    spng->callback = callback;
    spng->userdata = userdata;
    return true;
}

static s32 spng_decoder_push(struct camu_decoder *dec, struct camu_codec_packet *packet)
{
    struct camu_spng_decoder *spng = (struct camu_spng_decoder *)dec;
    if (!packet) return CAMU_OK;

    spng_ctx *ctx = spng_ctx_new(SPNG_CTX_IGNORE_ADLER32);
    spng_set_png_buffer(ctx, nn_buffer_get_ptr(packet->buffer, 0), packet->buffer->size);
    spng_set_option(ctx, SPNG_CHUNK_COUNT_LIMIT, 6250);
    spng_set_crc_action(ctx, SPNG_CRC_DISCARD, SPNG_CRC_DISCARD);

    size_t out_size;
    spng_decoded_image_size(ctx, SPNG_FMT_RGBA8, &out_size);

    struct camu_codec_frame *frame = al_alloc_object(struct camu_codec_frame);
    frame->mode = CAMU_NORMAL;
    frame->data = camu_page_alloc(out_size);
    spng_decode_chunks(ctx);
    spng_decode_image(ctx, frame->data, out_size, SPNG_FMT_RGBA8, 0);
    spng_ctx_free(ctx);
    frame->flags = 0;
    frame->format = CAMU_PIXEL_FORMAT_RGBA;
    struct camu_video_format *fmt = &spng->stream->video.fmt;
    frame->video.width = fmt->width;
    frame->video.height = fmt->height;
    frame->pts = 0.0;

    spng->callback(spng->userdata, frame);

    return CAMU_OK;
}

static void spng_decoder_flush(struct camu_decoder *dec)
{
    (void)dec;
}

static s32 spng_decoder_process(struct camu_decoder *dec)
{
    (void)dec;
    return CAMU_ERR_EOF;
}

static void spng_decoder_free(struct camu_decoder **dec)
{
    struct camu_spng_decoder *spng = (struct camu_spng_decoder *)*dec;
    al_free(spng);
    *dec = NULL;
}

struct camu_decoder *camu_spng_decoder_create(void)
{
    struct camu_spng_decoder *spng = al_alloc_object(struct camu_spng_decoder);
    spng->dec.mode = CAMU_NORMAL;
    spng->dec.init = spng_decoder_init;
    spng->dec.push = spng_decoder_push;
    spng->dec.process = spng_decoder_process;
    spng->dec.flush = spng_decoder_flush;
    spng->dec.free = spng_decoder_free;
    return (struct camu_decoder *)spng;
}
#endif