/*
 * flac.cpp
 * Native FLAC decoder for the Music player
 * Copyright (c) 2026 Daniel Hammer
 */

#include "flac.hpp"

#include <montauk/heap.h>

extern "C" {
#include <string.h>
}

namespace flac {

// ============================================================================
// CRC tables (CRC-8 poly 0x07 over frame headers, CRC-16 poly 0x8005 over
// whole frames)
// ============================================================================

static uint8_t  g_crc8[256];
static uint16_t g_crc16[256];
static bool     g_crc_ready = false;

static void init_crc_tables() {
    if (g_crc_ready) return;
    for (int i = 0; i < 256; i++) {
        uint8_t c8 = (uint8_t)i;
        for (int b = 0; b < 8; b++)
            c8 = (c8 & 0x80) ? (uint8_t)((c8 << 1) ^ 0x07) : (uint8_t)(c8 << 1);
        g_crc8[i] = c8;

        uint16_t c16 = (uint16_t)(i << 8);
        for (int b = 0; b < 8; b++)
            c16 = (c16 & 0x8000) ? (uint16_t)((c16 << 1) ^ 0x8005) : (uint16_t)(c16 << 1);
        g_crc16[i] = c16;
    }
    g_crc_ready = true;
}

static uint8_t crc8(const uint8_t* p, uint64_t n) {
    uint8_t c = 0;
    for (uint64_t i = 0; i < n; i++) c = g_crc8[c ^ p[i]];
    return c;
}

static uint16_t crc16(const uint8_t* p, uint64_t n) {
    uint16_t c = 0;
    for (uint64_t i = 0; i < n; i++)
        c = (uint16_t)((c << 8) ^ g_crc16[(c >> 8) ^ p[i]]);
    return c;
}

// ============================================================================
// Bit reader (MSB first)
// ============================================================================

struct BitReader {
    const uint8_t* data;
    uint64_t size;   // bytes
    uint64_t pos;    // bits
    bool overrun;
};

// The next 64 bits at the cursor; at least the top 57 are real stream bits
// (or zero padding past the end, which callers catch through `overrun`).
static inline uint64_t peek64(const BitReader& br) {
    uint64_t byte = br.pos >> 3;
    uint64_t v = 0;
    if (byte + 8 <= br.size) {
        uint64_t raw;
        memcpy(&raw, br.data + byte, 8);
        v = __builtin_bswap64(raw);
    } else {
        for (int i = 0; i < 8; i++) {
            v <<= 8;
            if (byte + i < br.size) v |= br.data[byte + i];
        }
    }
    return v << (br.pos & 7);
}

static inline void advance(BitReader& br, uint64_t bits) {
    br.pos += bits;
    if (br.pos > br.size * 8) br.overrun = true;
}

static inline uint32_t read_bits(BitReader& br, int n) {
    if (n == 0) return 0;
    uint32_t v = (uint32_t)(peek64(br) >> (64 - n));
    advance(br, n);
    return v;
}

static inline int32_t read_signed(BitReader& br, int n) {
    if (n == 0) return 0;
    uint32_t v = read_bits(br, n);
    return (int32_t)(v << (32 - n)) >> (32 - n);
}

static inline int64_t read_signed64(BitReader& br, int n) {
    if (n <= 32) return read_signed(br, n);
    uint64_t hi = read_bits(br, n - 32);
    uint64_t v = (hi << 32) | read_bits(br, 32);
    return (int64_t)(v << (64 - n)) >> (64 - n);
}

static inline uint32_t read_unary(BitReader& br) {
    uint32_t count = 0;
    for (;;) {
        uint64_t v = peek64(br);
        if (v) {
            int z = __builtin_clzll(v);
            advance(br, (uint64_t)z + 1);
            return count + (uint32_t)z;
        }
        advance(br, 56);
        count += 56;
        if (br.overrun) return 0;
    }
}

// ============================================================================
// Frame header
// ============================================================================

static bool parse_header(const Decoder& dec, uint64_t offset, FrameHeader& h) {
    if (offset >= dec.size) return false;
    uint64_t avail = dec.size - offset;

    // A header is at most 16 bytes; parse from a zero-padded copy so the
    // field reads need no bounds checks, then make sure it really fit.
    uint8_t p[16] = {};
    memcpy(p, dec.data + offset, avail < 16 ? avail : 16);

    if (p[0] != 0xFF || (p[1] & 0xFE) != 0xF8) return false;
    bool variable = p[1] & 1;
    uint8_t bs_code = p[2] >> 4;
    uint8_t sr_code = p[2] & 0x0F;
    uint8_t ch_code = p[3] >> 4;
    uint8_t ss_code = (p[3] >> 1) & 7;
    if ((p[3] & 1) || bs_code == 0 || sr_code == 15 || ch_code > 10 || ss_code == 3)
        return false;

    // UTF-8-style coded frame number (fixed) or sample number (variable)
    int i = 4;
    uint8_t b0 = p[i++];
    uint64_t num;
    int extra;
    if (!(b0 & 0x80))              { num = b0;        extra = 0; }
    else if ((b0 & 0xE0) == 0xC0)  { num = b0 & 0x1F; extra = 1; }
    else if ((b0 & 0xF0) == 0xE0)  { num = b0 & 0x0F; extra = 2; }
    else if ((b0 & 0xF8) == 0xF0)  { num = b0 & 0x07; extra = 3; }
    else if ((b0 & 0xFC) == 0xF8)  { num = b0 & 0x03; extra = 4; }
    else if ((b0 & 0xFE) == 0xFC)  { num = b0 & 0x01; extra = 5; }
    else if (b0 == 0xFE)           { num = 0;         extra = 6; }
    else return false;
    if (!variable && extra > 5) return false;
    for (int e = 0; e < extra; e++) {
        uint8_t b = p[i++];
        if ((b & 0xC0) != 0x80) return false;
        num = (num << 6) | (b & 0x3F);
    }

    uint32_t bs;
    if (bs_code == 1)       bs = 192;
    else if (bs_code <= 5)  bs = 576u << (bs_code - 2);
    else if (bs_code == 6)  bs = (uint32_t)p[i++] + 1;
    else if (bs_code == 7)  { bs = (((uint32_t)p[i] << 8) | p[i + 1]) + 1; i += 2; }
    else                    bs = 256u << (bs_code - 8);

    if (sr_code == 12)                       i += 1;
    else if (sr_code == 13 || sr_code == 14) i += 2;

    if ((uint64_t)i + 1 > avail) return false;
    if (crc8(p, i) != p[i]) return false;
    i++;

    static constexpr uint8_t ss_bits[8] = {0, 8, 12, 0, 16, 20, 24, 32};
    uint8_t bps = ss_code ? ss_bits[ss_code] : dec.info.bits_per_sample;
    uint8_t channels = ch_code < 8 ? (uint8_t)(ch_code + 1) : 2;

    if (channels != dec.info.channels) return false;
    if (bs > dec.chan_cap) return false;

    h.block_size = bs;
    h.header_bytes = (uint32_t)i;
    h.channel_assign = ch_code;
    h.channels = channels;
    h.bits_per_sample = bps;
    if (variable) {
        h.sample_start = num;
    } else {
        // Fixed-blocksize streams number frames; every frame but the last
        // has the stream's block size.
        uint64_t fixed = (dec.info.min_block == dec.info.max_block) ? dec.info.max_block : bs;
        h.sample_start = num * fixed;
    }
    return true;
}

// ============================================================================
// Subframes
// ============================================================================

static bool decode_residual(BitReader& br, int32_t* out, uint32_t bs, uint32_t order) {
    uint32_t method = read_bits(br, 2);
    if (method > 1) return false;
    int param_bits = method == 0 ? 4 : 5;
    uint32_t escape = method == 0 ? 15 : 31;

    uint32_t po = read_bits(br, 4);
    uint32_t parts = 1u << po;
    if (bs & (parts - 1)) return false;
    uint32_t psize = bs >> po;
    if (psize < order) return false;

    int32_t* dst = out + order;
    for (uint32_t p = 0; p < parts; p++) {
        uint32_t n = psize - (p == 0 ? order : 0);
        uint32_t k = read_bits(br, param_bits);
        if (k == escape) {
            int nbits = (int)read_bits(br, 5);
            for (uint32_t j = 0; j < n; j++) *dst++ = read_signed(br, nbits);
        } else {
            for (uint32_t j = 0; j < n; j++) {
                uint32_t q = read_unary(br);
                uint32_t u = (q << k) | read_bits(br, (int)k);
                *dst++ = (int32_t)(u >> 1) ^ -(int32_t)(u & 1);
            }
        }
        if (br.overrun) return false;
    }
    return true;
}

// Samples are int32_t, except the 33-bit side channel of 32-bit streams
// (int64_t). Prediction arithmetic is 64-bit either way.
template <typename T>
static void restore_fixed(T* s, uint32_t bs, uint32_t order) {
    switch (order) {
    case 1:
        for (uint32_t i = 1; i < bs; i++) s[i] = (T)(s[i] + (int64_t)s[i - 1]);
        break;
    case 2:
        for (uint32_t i = 2; i < bs; i++)
            s[i] = (T)(s[i] + 2 * (int64_t)s[i - 1] - s[i - 2]);
        break;
    case 3:
        for (uint32_t i = 3; i < bs; i++)
            s[i] = (T)(s[i] + 3 * ((int64_t)s[i - 1] - s[i - 2]) + s[i - 3]);
        break;
    case 4:
        for (uint32_t i = 4; i < bs; i++)
            s[i] = (T)(s[i] + 4 * ((int64_t)s[i - 1] + s[i - 3])
                       - 6 * (int64_t)s[i - 2] - s[i - 4]);
        break;
    default:
        break;
    }
}

template <typename T>
static void restore_lpc(T* s, uint32_t bs, uint32_t order,
                        const int32_t* coef, int shift, bool wide) {
    if (!wide && sizeof(T) == 4) {
        // Small sample + coefficient widths: the sum provably fits 32 bits.
        for (uint32_t i = order; i < bs; i++) {
            int32_t sum = 0;
            for (uint32_t j = 0; j < order; j++) sum += coef[j] * (int32_t)s[i - 1 - j];
            s[i] = (T)(s[i] + (sum >> shift));
        }
    } else {
        for (uint32_t i = order; i < bs; i++) {
            int64_t sum = 0;
            for (uint32_t j = 0; j < order; j++) sum += (int64_t)coef[j] * s[i - 1 - j];
            s[i] = (T)(s[i] + (sum >> shift));
        }
    }
}

// Residuals always fit 32 bits; they are decoded into `scratch` (which is
// `out` itself for int32_t samples) and widened into `out` when T is wider.
template <typename T>
static bool decode_subframe(BitReader& br, T* out, int32_t* scratch, uint32_t bs, int bps) {
    if (read_bits(br, 1) != 0) return false;
    uint32_t type = read_bits(br, 6);

    int wasted = 0;
    if (read_bits(br, 1)) {
        wasted = (int)read_unary(br) + 1;
        if (wasted >= bps) return false;
        bps -= wasted;
    }
    if (bps < 1 || bps > (sizeof(T) == 4 ? 32 : 33)) return false;

    uint32_t order = 0;
    if (type == 0) {
        T v = (T)read_signed64(br, bps);
        for (uint32_t i = 0; i < bs; i++) out[i] = v;
    } else if (type == 1) {
        for (uint32_t i = 0; i < bs; i++) out[i] = (T)read_signed64(br, bps);
    } else if ((type >= 8 && type <= 12) || type >= 32) {
        bool lpc = type >= 32;
        order = lpc ? (type & 31) + 1 : type - 8;
        if (order > bs) return false;
        for (uint32_t i = 0; i < order; i++) out[i] = (T)read_signed64(br, bps);

        int precision = 0, shift = 0;
        int32_t coef[32];
        if (lpc) {
            precision = (int)read_bits(br, 4) + 1;
            if (precision == 16) return false;
            shift = read_signed(br, 5);
            if (shift < 0) return false;
            for (uint32_t i = 0; i < order; i++) coef[i] = read_signed(br, precision);
        }

        if (!decode_residual(br, scratch, bs, order)) return false;
        if ((void*)scratch != (void*)out) {
            for (uint32_t i = order; i < bs; i++) out[i] = scratch[i];
        }

        if (lpc) {
            int order_bits = 0;
            while ((1u << order_bits) < order) order_bits++;
            restore_lpc(out, bs, order, coef, shift, bps + precision + order_bits > 32);
        } else {
            restore_fixed(out, bs, order);
        }
    } else {
        return false;
    }

    if (wasted) {
        for (uint32_t i = 0; i < bs; i++) out[i] = (T)((uint64_t)(int64_t)out[i] << wasted);
    }
    return !br.overrun;
}

// Undo stereo decorrelation into a (left) and b (right). `side` is the side
// channel: b for left/side and mid/side, a for side/right, or the wide buffer.
template <typename T>
static void decorrelate(int32_t* a, int32_t* b, const T* side, uint8_t assign, uint32_t bs) {
    if (assign == 8) {         // left / side
        for (uint32_t i = 0; i < bs; i++) b[i] = (int32_t)(a[i] - (int64_t)side[i]);
    } else if (assign == 9) {  // side / right
        for (uint32_t i = 0; i < bs; i++) a[i] = (int32_t)(side[i] + (int64_t)b[i]);
    } else if (assign == 10) { // mid / side
        for (uint32_t i = 0; i < bs; i++) {
            int64_t sd = side[i];
            int64_t mid = ((int64_t)a[i] * 2) | (sd & 1);
            a[i] = (int32_t)((mid + sd) >> 1);
            b[i] = (int32_t)((mid - sd) >> 1);
        }
    }
}

// ============================================================================
// Output conversion
// ============================================================================

static inline int16_t to_s16(int64_t v, int bps) {
    if (bps > 16) v >>= (bps - 16);
    else if (bps < 16) v *= (int64_t)1 << (16 - bps);
    if (v > 32767) v = 32767;
    if (v < -32768) v = -32768;
    return (int16_t)v;
}

// Stereo downmix weights (Q14, left/right) for the standard FLAC channel
// orders of 3..8 channel streams. LFE is dropped.
static constexpr int16_t W1 = 16384, WH = 11585, WQ = 8192;
static constexpr int16_t g_downmix[9][MAX_CHANNELS][2] = {
    {}, {}, {},
    /* 3: FL FR FC */          {{W1,0},{0,W1},{WH,WH}},
    /* 4: FL FR BL BR */       {{W1,0},{0,W1},{WH,0},{0,WH}},
    /* 5: FL FR FC BL BR */    {{W1,0},{0,W1},{WH,WH},{WH,0},{0,WH}},
    /* 6: + LFE */             {{W1,0},{0,W1},{WH,WH},{0,0},{WH,0},{0,WH}},
    /* 7: FL FR FC LFE BC SL SR */
                               {{W1,0},{0,W1},{WH,WH},{0,0},{WQ,WQ},{WH,0},{0,WH}},
    /* 8: FL FR FC LFE BL BR SL SR */
                               {{W1,0},{0,W1},{WH,WH},{0,0},{WH,0},{0,WH},{WH,0},{0,WH}},
};

static void emit(const Decoder& dec, const FrameHeader& h, int16_t* out) {
    uint32_t bs = h.block_size;
    int bps = h.bits_per_sample;

    if (h.channels == 1) {
        for (uint32_t i = 0; i < bs; i++) out[i] = to_s16(dec.chan[0][i], bps);
        return;
    }
    if (h.channels == 2) {
        for (uint32_t i = 0; i < bs; i++) {
            out[2 * i]     = to_s16(dec.chan[0][i], bps);
            out[2 * i + 1] = to_s16(dec.chan[1][i], bps);
        }
        return;
    }

    const int16_t (*w)[2] = g_downmix[h.channels];
    int64_t wl = 0, wr = 0;
    for (int c = 0; c < h.channels; c++) { wl += w[c][0]; wr += w[c][1]; }
    for (uint32_t i = 0; i < bs; i++) {
        int64_t l = 0, r = 0;
        for (int c = 0; c < h.channels; c++) {
            l += (int64_t)w[c][0] * dec.chan[c][i];
            r += (int64_t)w[c][1] * dec.chan[c][i];
        }
        out[2 * i]     = to_s16(l / wl, bps);
        out[2 * i + 1] = to_s16(r / wr, bps);
    }
}

// ============================================================================
// Public API
// ============================================================================

int output_channels(const Decoder& dec) {
    return dec.info.channels == 1 ? 1 : 2;
}

uint64_t total_samples(const Decoder& dec) {
    return dec.info.total_samples;
}

uint64_t find_frame(const Decoder& dec, uint64_t offset) {
    FrameHeader h;
    for (uint64_t i = offset; i + 1 < dec.size; i++) {
        if (dec.data[i] != 0xFF || (dec.data[i + 1] & 0xFE) != 0xF8) continue;
        if (parse_header(dec, i, h)) return i;
    }
    return dec.size;
}

int decode_frame(Decoder& dec, uint64_t offset, int16_t* out, int out_cap_frames,
                 uint64_t* next, uint64_t* sample_start) {
    FrameHeader h;
    if (!parse_header(dec, offset, h)) return -1;
    if ((int)h.block_size > out_cap_frames) return -1;

    BitReader br = {dec.data + offset, dec.size - offset, (uint64_t)h.header_bytes * 8, false};
    int side = h.channel_assign == 9 ? 0 : (h.channel_assign >= 8 ? 1 : -1);
    bool wide_side = false;
    for (int c = 0; c < h.channels; c++) {
        // The side channel carries one extra bit.
        int bps = h.bits_per_sample + (c == side ? 1 : 0);
        bool ok;
        if (bps > 32) {
            if (!dec.wide) return -1;
            ok = decode_subframe(br, dec.wide, dec.chan[c], h.block_size, bps);
            wide_side = true;
        } else {
            ok = decode_subframe(br, dec.chan[c], dec.chan[c], h.block_size, bps);
        }
        if (!ok) return -1;
    }

    uint64_t len = (br.pos + 7) / 8;
    if (len + 2 > br.size) return -1;
    const uint8_t* p = dec.data + offset;
    if (crc16(p, len) != (uint16_t)((p[len] << 8) | p[len + 1])) return -1;

    if (side >= 0) {
        if (wide_side)
            decorrelate(dec.chan[0], dec.chan[1], dec.wide, h.channel_assign, h.block_size);
        else
            decorrelate(dec.chan[0], dec.chan[1], dec.chan[side], h.channel_assign, h.block_size);
    }

    emit(dec, h, out);
    *next = offset + len + 2;
    *sample_start = h.sample_start;
    return (int)h.block_size;
}

bool seek(const Decoder& dec, uint64_t sample, uint64_t* offset, uint64_t* frame_start) {
    if (dec.index_count <= 0) return false;
    int lo = 0, hi = dec.index_count - 1, best = 0;
    while (lo <= hi) {
        int mid = lo + (hi - lo) / 2;
        if (dec.index[mid].sample_start <= sample) {
            best = mid;
            lo = mid + 1;
        } else {
            hi = mid - 1;
        }
    }
    *offset = dec.index[best].offset;
    *frame_start = dec.index[best].sample_start;
    return true;
}

static bool push_index(Decoder& dec, int& cap, uint64_t sample_start, uint64_t offset) {
    if (dec.index_count >= cap) {
        int new_cap = cap ? cap * 2 : 1024;
        auto* grown = (SeekPoint*)montauk::realloc(dec.index, (uint64_t)new_cap * sizeof(SeekPoint));
        if (!grown) return false;
        dec.index = grown;
        cap = new_cap;
    }
    dec.index[dec.index_count++] = {sample_start, offset};
    return true;
}

// Walk the frame chain without decoding: each frame's successor is the next
// valid header whose first sample is exactly where this frame ends, which
// rejects sync patterns that happen to appear inside audio data.
static bool build_index(Decoder& dec) {
    int cap = 0;
    uint64_t off = find_frame(dec, dec.first_frame);
    uint64_t last_end = 0;
    FrameHeader h;

    while (off < dec.size && parse_header(dec, off, h)) {
        if (!push_index(dec, cap, h.sample_start, off)) return false;
        uint64_t want = h.sample_start + h.block_size;
        last_end = want;
        if (dec.info.total_samples && want >= dec.info.total_samples) break;

        uint64_t found = dec.size;
        FrameHeader nh;
        for (uint64_t i = off + h.header_bytes + 2; i + 1 < dec.size; i++) {
            if (dec.data[i] != 0xFF || (dec.data[i + 1] & 0xFE) != 0xF8) continue;
            if (parse_header(dec, i, nh) && nh.sample_start == want) {
                found = i;
                break;
            }
        }
        off = found;
    }

    if (dec.index_count <= 0) return false;
    if (!dec.info.total_samples) dec.info.total_samples = last_end;
    return true;
}

void close(Decoder& dec) {
    for (int c = 0; c < MAX_CHANNELS; c++) {
        if (dec.chan[c]) montauk::mfree(dec.chan[c]);
        dec.chan[c] = nullptr;
    }
    if (dec.wide) montauk::mfree(dec.wide);
    dec.wide = nullptr;
    if (dec.index) montauk::mfree(dec.index);
    dec.index = nullptr;
    dec.index_count = 0;
    dec.chan_cap = 0;
    dec.data = nullptr;
    dec.size = 0;
}

bool open(Decoder& dec, const uint8_t* data, uint64_t size) {
    init_crc_tables();
    close(dec);
    memset(&dec, 0, sizeof(dec));

    uint64_t pos = 0;

    // Some taggers prepend ID3v2 to FLAC files.
    if (size >= 10 && data[0] == 'I' && data[1] == 'D' && data[2] == '3') {
        pos = 10 + (((uint64_t)(data[6] & 0x7F) << 21) |
                    ((uint64_t)(data[7] & 0x7F) << 14) |
                    ((uint64_t)(data[8] & 0x7F) << 7)  |
                    ((uint64_t)(data[9] & 0x7F)));
        if (data[5] & 0x10) pos += 10; // footer
    }

    if (pos + 4 > size || memcmp(data + pos, "fLaC", 4) != 0) return false;
    pos += 4;

    bool have_info = false;
    for (;;) {
        if (pos + 4 > size) return false;
        uint8_t type = data[pos];
        uint32_t len = ((uint32_t)data[pos + 1] << 16) | ((uint32_t)data[pos + 2] << 8) | data[pos + 3];
        pos += 4;
        if (pos + len > size) return false;

        if ((type & 0x7F) == 0 && len >= 34) {
            const uint8_t* d = data + pos;
            StreamInfo& si = dec.info;
            si.min_block = ((uint32_t)d[0] << 8) | d[1];
            si.max_block = ((uint32_t)d[2] << 8) | d[3];
            si.min_frame_bytes = ((uint32_t)d[4] << 16) | ((uint32_t)d[5] << 8) | d[6];
            si.sample_rate = ((uint32_t)d[10] << 12) | ((uint32_t)d[11] << 4) | (d[12] >> 4);
            si.channels = (uint8_t)(((d[12] >> 1) & 7) + 1);
            si.bits_per_sample = (uint8_t)((((d[12] & 1) << 4) | (d[13] >> 4)) + 1);
            si.total_samples = ((uint64_t)(d[13] & 0x0F) << 32) |
                               ((uint64_t)d[14] << 24) | ((uint64_t)d[15] << 16) |
                               ((uint64_t)d[16] << 8) | d[17];
            have_info = true;
        }

        pos += len;
        if (type & 0x80) break;
    }

    StreamInfo& si = dec.info;
    if (!have_info || si.sample_rate == 0 || si.max_block == 0 ||
        si.bits_per_sample < 4 || si.bits_per_sample > 32)
        return false;

    dec.data = data;
    dec.size = size;
    dec.first_frame = pos;

    dec.chan_cap = si.max_block;
    for (int c = 0; c < si.channels; c++) {
        dec.chan[c] = (int32_t*)montauk::malloc((uint64_t)dec.chan_cap * sizeof(int32_t));
        if (!dec.chan[c]) { close(dec); return false; }
    }

    if (si.bits_per_sample == 32 && si.channels == 2) {
        dec.wide = (int64_t*)montauk::malloc((uint64_t)dec.chan_cap * sizeof(int64_t));
        if (!dec.wide) { close(dec); return false; }
    }

    if (!build_index(dec)) { close(dec); return false; }
    return true;
}

} // namespace flac