#define _DEFAULT_SOURCE // can be removed once glibc knows about <endian.h> in POSIX 2024
#include <stdio.h>
#include <stdlib.h>
#include <stdbit.h>
#include <endian.h>
#include <string.h>
#include "util/err.h"
#include "nut.h"

#define DATA_SIZE_MUL 16383
#define STATIC_OUT(buffer) { .type = NUT_OUT_STATIC, .ptr = buffer, .cap = sizeof buffer }

static uint32_t crc32(size_t len, char bytes[len])
{
    int32_t crc = 0;
    for (size_t i = 0; i < len; i++) {
        crc ^= bytes[i] << 24;
        for (size_t j = 0; j < 8; j++)
            crc = (crc << 1) ^ ((0x04C11DB7) & (crc >> 31));
    }
    return crc;
}

static void put_f(struct nut_output *out, size_t len, const void *data)
{
    if (out->type == NUT_OUT_FILE) {
        fwrite(data, 1, len, out->file);
    } else {
        if (out->type == NUT_OUT_HEAP)
            arraybuf_grow(out, len);
        else if (out->len + len > out->cap)
            eprintf("NUT buffer size exceeded\n");
        memcpy(out->ptr+out->len, data, len);
    }

    out->len += len;
}

static void put_u8(struct nut_output *out, uint8_t v)
{
    put_f(out, sizeof v, &v);
}

static void put_u32(struct nut_output *out, uint32_t v)
{
    v = htobe32(v);
    put_f(out, sizeof v, &v);
}

static void put_u64(struct nut_output *out, uint64_t v)
{
    v = htobe64(v);
    put_f(out, sizeof v, &v);
}

static void put_v(struct nut_output *out, uint64_t v)
{
    unsigned int bits = stdc_bit_width(v);
    if (!bits) bits = 1;
    for (int i = bits/7; i >= 0; i--)
        put_u8(out, ((v >> (i*7)) & 0x7f) | ((!!i)*0x80));
}

static size_t put_packet_header(struct nut_output *out, size_t data_len, uint64_t startcode)
{
    char buffer[BUFSIZ];
    struct nut_output header = STATIC_OUT(buffer);

    put_u64(&header, startcode);
    size_t forward_ptr = data_len+sizeof(uint32_t);
    put_v(&header, forward_ptr);
    if (data_len > 4096)
        put_u32(&header, crc32(header.len, header.ptr));
    put_f(out, header.len, header.ptr);
    return header.len;
}

static void put_packet_data(struct nut_output *out, size_t len, void *data)
{
    put_f(out, len, data);
    put_u32(out, crc32(len, data));
}

static void put_packet(struct nut_output *out, size_t len, void *data, uint64_t startcode)
{
    put_packet_header(out, len, startcode);
    put_packet_data(out, len, data);
}

static void put_syncpoint(struct nut_writer *wr)
{
    char buffer[BUFSIZ];

    size_t pos = wr->out.len;
    size_t offset = pos - wr->last_syncpoint;

    if (wr->syncp_offs.len > 0 && wr->syncp_offs.ptr[wr->syncp_offs.len-1].off == offset)
        wr->syncp_offs.ptr[wr->syncp_offs.len-1].num++;
    else
        arraybuf_insert(&wr->syncp_offs, ((struct nut_syncp_off) { .off = offset, .num = 1 }));

    struct nut_output syncp = STATIC_OUT(buffer);
    put_v(&syncp, wr->n_frame++); // timestamp
    put_v(&syncp, wr->last_syncpoint ? offset / 16 : 0); // back_ptr
    put_packet(&wr->out, syncp.len, syncp.ptr,
        0xE4ADEECA4569ULL + (((uint64_t)('N'<<8) + 'K')<<48));

    wr->last_syncpoint = pos;
}

void nut_write_start(struct nut_writer *wr, FILE *f, unsigned int fps, unsigned int size[2])
{
    char buffer[BUFSIZ];

    wr->out.type = NUT_OUT_FILE;
    wr->out.file = f;
    wr->out.len = 0;
    wr->frame_size = size[0] * size[1] * 4;
    wr->last_syncpoint = 0; // ???
    wr->n_frame = 0;
    wr->syncp_offs.len = 0;
    wr->syncp_offs.cap = 4;
    wr->syncp_offs.ptr = nullptr;

    const char id_string[] = "nut/multimedia container";
    put_f(&wr->out, sizeof id_string, id_string);

    // main header
    struct nut_output main = STATIC_OUT(buffer);

    put_v(&main, 3); // version
    put_v(&main, 1); // stream count
    put_v(&main, 65536); // max_distance
    put_v(&main, 1); // time_base_count
    put_v(&main, 1);   // time_base_denom
    put_v(&main, fps); // time_base_num

    // frame code 0
    put_v(&main, 1 << 13); // FLAG_INVALID
    put_v(&main, 0); // fields

    // frame code 1
    put_v(&main, (1 << 0) | (1 << 6) | (1 << 5)); // FLAG_KEY | FLAG_CHECKSUM | FLAG_SIZE_MSB
    put_v(&main, 6); // fields
    put_v(&main, 0); // pts
    put_v(&main, DATA_SIZE_MUL); // mul
    put_v(&main, 0); // stream
    put_v(&main, wr->frame_size % DATA_SIZE_MUL); // size
    put_v(&main, 0); // res
    put_v(&main, 1); // count

    // frame code 2
    put_v(&main, (1 << 0) | (1 << 1)); // FLAG_KEY | FLAG_EOR
    put_v(&main, 2); // fields
    put_v(&main, 0); // pts
    put_v(&main, 1); // mul

    // other frame codes
    put_v(&main, 1 << 13); // FLAG_INVALID
    put_v(&main, 6); // fields
    for (size_t i = 0; i < 5; i++)
        put_v(&main, 0);
    put_v(&main, 256-4); // count

    put_v(&main, 0); // header_count_minus1
    put_v(&main, 0); // main_flags

    put_packet(&wr->out, main.len, main.ptr,
        0x7A561F5F04ADULL + (((uint64_t)('N'<<8) + 'M')<<48));

    // stream header
    struct nut_output stream = STATIC_OUT(buffer);

    put_v(&stream, 0); // stream id
    put_v(&stream, 0); // stream class: video

    const char fourcc[] = "RGBA";
    put_v(&stream, sizeof fourcc-1);
    put_f(&stream, sizeof fourcc-1, fourcc);

    put_v(&stream, 0); // time base id

    put_v(&stream, 0); // msb_pts_shift
    put_v(&stream, 0); // max_pts_distance

    put_v(&stream, 0); // decode_delay
    put_v(&stream, 1); // stream flags: FLAG_FIXED_FPS
    put_v(&stream, 0); // length of codec_specific_data

    // video data
    put_v(&stream, size[0]); // width
    put_v(&stream, size[1]); // height
    put_v(&stream, 1); // sample_width
    put_v(&stream, 1); // sample_height
    put_v(&stream, 0); // colorspace_type: unknown (?)

    put_packet(&wr->out, stream.len, stream.ptr,
        0x11405BF2F9DBULL + (((uint64_t)('N'<<8) + 'S')<<48));
}

void nut_write_frame(struct nut_writer *wr, void *data)
{
    char buffer[BUFSIZ];

    put_syncpoint(wr);

    struct nut_output frame = STATIC_OUT(buffer);
    put_u8(&frame, 1); // frame code
    put_v(&frame, wr->frame_size / DATA_SIZE_MUL); // size msb
    put_f(&wr->out, frame.len, frame.ptr);
    put_u32(&wr->out, crc32(frame.len, frame.ptr));

    put_f(&wr->out, wr->frame_size, data);
}

void nut_write_end(struct nut_writer *wr)
{
    put_syncpoint(wr);
    put_u8(&wr->out, 2); // frame code

    struct nut_output index = { .type = NUT_OUT_HEAP };
    put_v(&index, wr->n_frame-1); // max_pts
    put_v(&index, wr->n_frame); // syncpoints
    size_t syncp_pos = 0;
    for (size_t i = 0; i < wr->syncp_offs.len; i++)
        for (size_t n = 0; n < wr->syncp_offs.ptr[i].num; n++) {
            size_t last = syncp_pos;
            syncp_pos += wr->syncp_offs.ptr[i].off;
            put_v(&index, syncp_pos / 16 - last / 16);
        }
    free(wr->syncp_offs.ptr);

    // first syncp: does not have keyframe before it: type=0 f0=0 term=1
    put_v(&index, 0b100);

    // other syncps: all have keyframes type=1 flag=1 n=(n_frame-1)
    put_v(&index, ((wr->n_frame-1)<<2) | 0b11);

    for (size_t i = 1; i < wr->n_frame; i++) {
        if (i == wr->n_frame-1) {
            put_v(&index, 0); // indicate eor
            put_v(&index, 1); // A
            put_v(&index, 1); // B
        } else
            put_v(&index, 1); // A
    }

    size_t header_len = put_packet_header(&wr->out, index.len+sizeof(uint64_t),
        0xDD672F23E64EULL + (((uint64_t)('N'<<8) + 'X')<<48));

    put_u64(&index, header_len+index.len+sizeof(uint64_t)+sizeof(uint32_t)); // index_ptr

    put_packet_data(&wr->out, index.len, index.ptr);

    free(index.ptr);
}