path: root/ytdl-storyboard/storyboard.c

#include <mpv/client.h>
#undef NDEBUG
#include <assert.h>
#include <stdatomic.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <stdio.h>
#include <math.h>
#include <limits.h>
#include <stdarg.h>
#include <cJSON.h>
#include <pthread.h>

#include <libavutil/pixfmt.h>
#include <libavutil/imgutils.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libswscale/swscale.h>

struct plugin;
struct plugin{
    mpv_handle *m;
    struct storyboard *sb;
    int64_t osd_remove_time_ns;
};

enum hooknumber{
    HOOK_PRELOADED,
    HOOK_UNLOAD
};
enum overlay_nums{
    OSD_STORYBOARD = 1,
};

#define Z(err) ({ \
    typeof(err) tmp = (err); \
    if(tmp) return -1; \
    tmp; \
})
#define NZ(err) ({ \
    typeof(err) tmp = (err); \
    if(!tmp) return -1; \
    tmp; \
})
#define TRY_goto(label, err, msg, ...) do{ \
    int tmp = (err); \
    if(tmp){ \
        fprintf(stderr, "error%d at %s:%d(%s): " msg "\n", \
            tmp, __FILE__, __LINE__, __func__, ## __VA_ARGS__); \
        goto label; \
    } \
}while(0)
#define TRY(err, msg, ...) do{ \
    int tmp = (err); \
    if(tmp){ \
        fprintf(stderr, "error%d at %s:%d(%s): " msg "\n", \
            tmp, __FILE__, __LINE__, __func__, ## __VA_ARGS__); \
        return -1; \
    } \
}while(0)

static int msg(struct plugin *self, const char *f, ...){
    int len;
    char msg_arr[1024];
    char *msg = msg_arr;
    char *freeme = NULL;

    va_list ap;
    va_start(ap, f);
    len = vsnprintf(msg_arr, sizeof(msg_arr), f, ap);
    if(sizeof(msg_arr) < len+1){
        freeme = msg = malloc(len+1);
        assert(msg);
        assert(vsnprintf(msg, len+1, f, ap) == len);
        msg[len] = 0;
    }
    va_end(ap);

    const char *command[] = { "print-text", msg, NULL };
    TRY(mpv_command(self->m, command), "print-text");

    free(freeme);
    return 0;
}
static int overlay_add(struct plugin *self,
    int id, int x, int y,
    uint8_t *buf, size_t off, size_t w, size_t h, size_t stride,
    size_t dw, size_t dh
){
    assert(off <= INT64_MAX);
    assert(w <= INT64_MAX);
    assert(h <= INT64_MAX);
    assert(stride <= INT64_MAX);
    assert(dw <= INT64_MAX);
    assert(dh <= INT64_MAX);

    char *command_keys[] = {
        "name", "id", "x", "y",
        "file", "offset", "fmt", "w", "h", "stride",
        "dw", "dh",
    };
    char filename[32];  // enough for 64bit int + 1byte prefix
    sprintf(filename, "&%tu", (void *)buf - NULL);
    mpv_node command_vals[] = {
        {
            .format = MPV_FORMAT_STRING,
            .u.string = (char *)"overlay-add",
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = id,
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = x,
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = y,
        }, {
            .format = MPV_FORMAT_STRING,
            .u.string = filename,
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = off,
        }, {
            .format = MPV_FORMAT_STRING,
            .u.string = (char *)"bgra",
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = w,
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = h,
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = stride,
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = dw,
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = dh,
        },
    };
    assert(sizeof(command_keys)/sizeof(command_keys[0])
        == sizeof(command_vals)/sizeof(command_vals[0]));
    mpv_node_list command_list = {
        .num = sizeof(command_keys)/sizeof(command_keys[0]),
        .values = command_vals,
        .keys = command_keys,
    };
    mpv_node command = {
        .format = MPV_FORMAT_NODE_MAP,
        .u.list = &command_list,
    };

    TRY(mpv_command_node(self->m, &command, NULL), "mpv_command_node");
    return 0;
}
static int overlay_remove(struct plugin *self, int id){
    char *command_keys[] = { "name", "id" };
    mpv_node command_vals[] = {
        {
            .format = MPV_FORMAT_STRING,
            .u.string = (char *)"overlay-remove",
        }, {
            .format = MPV_FORMAT_INT64,
            .u.int64 = id,
        },
    };
    assert(sizeof(command_keys)/sizeof(command_keys[0])
        == sizeof(command_vals)/sizeof(command_vals[0]));
    mpv_node_list command_list = {
        .num = sizeof(command_keys)/sizeof(command_keys[0]),
        .values = command_vals,
        .keys = command_keys,
    };
    mpv_node command = {
        .format = MPV_FORMAT_NODE_MAP,
        .u.list = &command_list,
    };

    TRY(mpv_command_node(self->m, &command, NULL), "mpv_command_node");
    return 0;
}
static int url_verifyhttp(const char *url){
    const char *protoend = strstr(url, "://");
    if(!protoend) return 0;
    if(protoend - url == 4){
        return !memcmp(url, "http", 4);
    }else if(protoend - url == 5){
        return !memcmp(url, "https", 5);
    }
    return 0;
}
static int load_image(AVFrame **framepp, const char *url){
    int retval = 0;

    const AVInputFormat *iformat = av_find_input_format("image2pipe");
    if(!iformat) return -1;
    const AVCodec *decoder = NULL;
    AVFormatContext *format_ctx = NULL;
    AVCodecContext *codec_ctx = NULL;
    AVFrame *frame = NULL;
    AVPacket pkt;

    // initialize demuxer
    TRY_goto(err, avformat_open_input(
            &format_ctx, url, iformat, NULL) < 0,
        "avformat_open_input");
    TRY_goto(err, avformat_find_stream_info(
            format_ctx, NULL) < 0,
        "avformat_find_stream_info");

    // initialize codec
    decoder = avcodec_find_decoder(
        format_ctx->streams[0]->codecpar->codec_id);
    if(!decoder) goto err;
    codec_ctx = avcodec_alloc_context3(decoder);
    if(!codec_ctx) goto err;
    TRY_goto(err, avcodec_parameters_to_context(
            codec_ctx, format_ctx->streams[0]->codecpar) < 0,
        "avcodec_parameters_to_context");

    // decode
    TRY_goto(err, avcodec_open2(codec_ctx, decoder, NULL), "avcodec_open2");
    frame = av_frame_alloc();
    if(!frame) goto err;
    TRY_goto(err, av_read_frame(format_ctx, &pkt), "av_read_frame");
    int sps = avcodec_send_packet(codec_ctx, &pkt);
    av_packet_unref(&pkt);
    TRY_goto(err, sps, "avcodec_send_packet");
    TRY_goto(err, avcodec_receive_frame(
        codec_ctx, frame), "avcodec_receive_frame");
    *framepp = frame;

    goto final;
err:
    retval = -1;
    av_frame_free(&frame);
final:
    avcodec_free_context(&codec_ctx);
    avformat_close_input(&format_ctx);
    return retval;
}
static int convert_image_rgb32(AVFrame *frame, uint8_t **bitmap){
    int retval = 0;
    uint8_t *dstbuf[4] = {0};
    int linesizes[4];
    struct SwsContext *sws_ctx = NULL;

    TRY_goto(err,
        av_image_alloc(dstbuf, linesizes,
            frame->width, frame->height, AV_PIX_FMT_RGB32, 16) < 0,
        "av_image_alloc");
    sws_ctx = sws_getContext(
        frame->width, frame->height, frame->format,
        frame->width, frame->height, AV_PIX_FMT_RGB32,
        SWS_POINT, NULL, NULL, NULL
    );
    TRY_goto(err, sws_ctx == NULL, "sws_getContext");
    sws_scale(sws_ctx,
        (const uint8_t *const *)frame->data, frame->linesize, 0, frame->height,
        dstbuf, linesizes);

    *bitmap = dstbuf[0];
    goto final;
err:
    retval = -1;
    av_freep(&dstbuf[0]);
final:
    sws_freeContext(sws_ctx);
    return retval;
}

struct storyboard_fragment{
    char url[4096];
    double offset;
    double duration;
    size_t frames;
    uint8_t *buf;
};
struct storyboard{
    size_t width;
    size_t height;
    size_t rows;
    size_t columns;
    double fps;
    size_t fragment_cnt;
    struct storyboard_fragment *fragments;

    int thread_initialized;
    atomic_int ready;
    pthread_t thread;
};
static int storyboard_init(struct storyboard *self, cJSON *src){
#define READSZ(name) do{ \
    cJSON *tmp = cJSON_GetObjectItem(src, #name); \
    if(!tmp) return -1; \
    double tmp_n = cJSON_GetNumberValue(tmp); \
    if(isnan(tmp_n) || tmp_n != (double)(size_t)tmp_n) return -1; \
    self->name = (size_t)tmp_n; \
}while(0);

    self->thread_initialized = 0;

    READSZ(width);
    READSZ(height);
    READSZ(columns);
    READSZ(rows);
    if((size_t)-1 / self->width / self->height / self->columns / self->rows < 1)
            return -1;

    cJSON *j_fps = cJSON_GetObjectItem(src, "fps");
    if(!j_fps) return -1;
    double fps = cJSON_GetNumberValue(j_fps);
    if(isnan(fps) || fps < 0) return -1;
    self->fps = fps;

    cJSON *j_fragments = cJSON_GetObjectItem(src, "fragments");
    if(!j_fragments) return -1;
    int fragment_cnt = cJSON_GetArraySize(j_fragments);
    if(fragment_cnt <= 0 || 1024 <= fragment_cnt) return -1;
    self->fragment_cnt = fragment_cnt;
    self->fragments = malloc(fragment_cnt * sizeof(self->fragments[0]));
    assert(self->fragments);
    int fi;
    double offset = 0;
    for(fi = 0; fi < fragment_cnt; fi++){
        cJSON *j_fragment = cJSON_GetArrayItem(j_fragments, fi);
        if(!j_fragment) goto err;
        cJSON *j_url = cJSON_GetObjectItem(j_fragment, "url");
        if(!j_url) goto err;
        const char *url = cJSON_GetStringValue(j_url);
        if(!url || sizeof(self->fragments[fi].url) <= strlen(url)) goto err;
        if(!url_verifyhttp(url)) goto err;
        cJSON *j_duration = cJSON_GetObjectItem(j_fragment, "duration");
        if(!j_duration) goto err;
        double duration = cJSON_GetNumberValue(j_duration);
        if(isnan(duration) || duration < 0) goto err;
        strcpy(self->fragments[fi].url, url);
        self->fragments[fi].offset = offset;
        self->fragments[fi].duration = duration;
        self->fragments[fi].buf = NULL;
        offset += duration;
    }

#undef READSZ
    return 0;
err:
    free(self->fragments);
    self->fragments = NULL;
    return -1;
}
static int storyboard_destroy(struct storyboard *self){
    if(self->thread_initialized){
        self->thread_initialized = 0;
        TRY(pthread_cancel(self->thread), "pthread_cancel");
        TRY(pthread_join(self->thread, NULL), "pthread_join");
    }
    if(self->fragments){
        for(int fi = 0; fi < self->fragment_cnt; fi++){
            free(self->fragments[fi].buf);
        }
        free(self->fragments);
    }
    return 0;
}
static int storyboard_thread_inner(struct storyboard *self){
    int retval = 0;
    int oldstate;
    pthread_setcancelstate(PTHREAD_CANCEL_DISABLE, &oldstate);

#define CANCELPOINT \
    pthread_setcancelstate(PTHREAD_CANCEL_ENABLE, &oldstate); \
    pthread_testcancel(); \
    pthread_setcancelstate(PTHREAD_CANCEL_DISABLE, &oldstate)

    for(size_t i = 0; i < self->fragment_cnt; i++){
        AVFrame *frame = NULL;
        uint8_t *rgb = NULL;
        pthread_cleanup_push((void (*)(void *))av_frame_free, &frame);
        pthread_cleanup_push((void (*)(void *))av_freep, &rgb);

        TRY_goto(err, load_image(&frame, self->fragments[i].url), "load_image");
        CANCELPOINT;

        if(
            frame->width < 0 || frame->height < 0
            || frame->width % self->width != 0
            || frame->height % self->height != 0
        ){
            goto err;
        }

        TRY_goto(err, convert_image_rgb32(frame, &rgb), "convert_image_rgb32");
        CANCELPOINT;

        size_t frag_w = frame->width / self->width;
        size_t frag_h = frame->height / self->height;
        size_t frag_s = self->width * frag_w;
        self->fragments[i].frames = frag_w * frag_h;
        uint8_t *buf = self->fragments[i].buf
            = malloc(frame->width * frame->height * 4);
        if(!buf) goto err;
        for(size_t f = 0; f < frag_w * frag_h; f++){
            size_t f_off = frag_s * self->height * (f / frag_w)
                + self->width * (f % frag_w);
            size_t of_off = self->width * self->height * f;
            for(size_t y = 0; y < self->height; y++){
                size_t off = f_off + frag_s * y;
                size_t o_off = of_off + self->width * y;
                memcpy(&buf[o_off*4], &rgb[off*4], self->width * 4);
            }
        }
        CANCELPOINT;

        goto final;
err:
        retval = -1;
final:
        pthread_cleanup_pop(1);
        pthread_cleanup_pop(1);
        if(retval) break;
    }
    return retval;
#undef CANCELPOINT
}
static void *storyboard_thread(void *self_v){
    struct storyboard *self = self_v;
    if(storyboard_thread_inner(self)){
        atomic_store(&self->ready, -1);
    }else{
        atomic_store(&self->ready, 1);
    }
    return NULL;
}
static int storyboard_spawn(struct storyboard *self){
    atomic_init(&self->ready, 0);
    TRY(pthread_create(&self->thread, NULL, storyboard_thread, self),
        "pthread_create");
    self->thread_initialized = 1;
    return 0;
}
static int storyboard_getimg(struct storyboard *self,
    uint8_t **ptr, double time
){
    if(!self->thread_initialized) return 1;
    if(atomic_load(&self->ready) != 1) return 1;

    struct storyboard_fragment *fragment = NULL;
    for(size_t i = 0; i < self->fragment_cnt; i++){
        struct storyboard_fragment *f = &self->fragments[i];
        if(f->offset <= time){
            fragment = f;
        }
    }
    if(!fragment || fragment->offset + fragment->duration < time){
        return 1;
    }

    assert(time - fragment->offset >= 0);
    size_t frame = (time - fragment->offset) * self->fps;
    if(fragment->frames <= frame) frame = fragment->frames - 1;
    size_t frame_bi = self->width * self->height * 4 * frame;

    *ptr = &fragment->buf[frame_bi];
    return 0;
}

static int plugin_init(struct plugin *self, mpv_handle *h){
    self->m = h;
    self->sb = NULL;
    self->osd_remove_time_ns = 0;

    TRY(mpv_hook_add(self->m, HOOK_PRELOADED, "on_preloaded", 0), "hook_add");
    TRY(mpv_hook_add(self->m, HOOK_UNLOAD, "on_unload", 0), "hook_add");

    return 0;
}
static int hook_on_preloaded(struct plugin *self){
    char *result_s = mpv_get_property_string(
        self->m,
        "user-data/mpv/ytdl/json-subprocess-result");
    if(!result_s){        // not ytdl video
        goto noytdl;
    }
    cJSON *result = cJSON_Parse(result_s);
    if(!result){
        msg(self, "ytdl subprocess-json parse error");
        goto resultperr;
    }

    cJSON *j_status = cJSON_GetObjectItem(result, "status");
    if(!j_status || cJSON_GetNumberValue(j_status) != 0.0){
        goto ytdlfail;
    }
    cJSON *j_ytdlout_s = cJSON_GetObjectItem(result, "stdout");
    if(!j_ytdlout_s) goto ytdlfail;
    const char *ytdlout_s = cJSON_GetStringValue(j_ytdlout_s);
    if(!ytdlout_s) goto ytdlfail;

    cJSON *ytdlout = cJSON_Parse(ytdlout_s);
    if(!ytdlout) goto ytdlfail;

    cJSON *formats = cJSON_GetObjectItem(ytdlout, "formats");
    if(!formats) goto err;
    int formats_len = cJSON_GetArraySize(formats);
    struct storyboard *best_sb = NULL;
    for(int i = 0; i < formats_len; i++){
        cJSON *format = cJSON_GetArrayItem(formats, i);
        if(!format) goto err;
        cJSON *j_format_id = cJSON_GetObjectItem(format, "format_id");
        if(!j_format_id) continue;
        const char *format_id = cJSON_GetStringValue(j_format_id);
        if(!format_id) continue;
        if(format_id[0] != 's' || format_id[1] != 'b') continue;
        struct storyboard *sb = malloc(sizeof(struct storyboard));
        assert(sb);
        if(storyboard_init(sb, format)) continue;
        if(!best_sb || best_sb->height < sb->height){
            if(best_sb) storyboard_destroy(best_sb);
            free(best_sb);
            best_sb = sb;
        }else{
            storyboard_destroy(sb);
            free(sb);
        }
    }
    if(!best_sb) goto err;
    storyboard_spawn(best_sb);
    if(self->sb){
        storyboard_destroy(self->sb);
        free(self->sb);
    }
    self->sb = best_sb;

err:
    cJSON_Delete(ytdlout);
ytdlfail:
    cJSON_Delete(result);
resultperr:
    mpv_free(result_s);
noytdl:
    return 0;
}
static int hook_on_unload(struct plugin *self){
    if(self->sb){
        storyboard_destroy(self->sb);
        free(self->sb);
        self->sb = NULL;
    }
    return 0;
}
static int event_on_seek(struct plugin *self){
    if(!self->sb) return 0;
    uint8_t *img;

    double pos;
    TRY(mpv_get_property(self->m, "time-pos", MPV_FORMAT_DOUBLE, &pos),
        "mpv_get_property");
    if(storyboard_getimg(self->sb, &img, pos)) return 0;

    overlay_add(self, OSD_STORYBOARD, 0, 0,
        img, 0, self->sb->width, self->sb->height, self->sb->width*4,
        self->sb->width, self->sb->height
    );
    self->osd_remove_time_ns = mpv_get_time_ns(self->m) + 1*1000*1000*1000;
    if(self->osd_remove_time_ns == 0) self->osd_remove_time_ns = 1;
    return 0;
}
static int timeout_handler_osd_remove(struct plugin *self, int64_t now){
    overlay_remove(self, OSD_STORYBOARD);
    return 0;
}
static int event_handler(struct plugin *self, const mpv_event *ev){
    switch(ev->event_id){
    case MPV_EVENT_HOOK:
        mpv_event_hook *hook = ev->data;
        switch(ev->reply_userdata){
        case HOOK_UNLOAD:
            Z(hook_on_unload(self));
            break;
        case HOOK_PRELOADED:
            Z(hook_on_preloaded(self));
            break;
        default:
            return -1;
        }
        TRY(mpv_hook_continue(self->m, hook->id), "hook continue");
        return 0;
    case MPV_EVENT_SEEK:
        Z(event_on_seek(self));
        return 0;
    default:
        return 0;
    }
}
static double event_wait_timeout(struct plugin *self){
    if(self->osd_remove_time_ns){
        int64_t now = mpv_get_time_ns(self->m);
        if(self->osd_remove_time_ns <= now){
            self->osd_remove_time_ns = 0;
            timeout_handler_osd_remove(self, now);
            return event_wait_timeout(self);
        }else{
            return (self->osd_remove_time_ns - now) / 1.0e9;
        }
    }else{
        return -1;
    }
}

int mpv_open_cplugin(mpv_handle *h){
    struct plugin self;
    Z(plugin_init(&self, h));

    mpv_event *ev;
    while(ev = mpv_wait_event(h, event_wait_timeout(&self))){
        Z(event_handler(&self, ev));
        if(ev->event_id == MPV_EVENT_SHUTDOWN){
            break;
        }
    }
    return 0;
}