#include "dtob_internal.h"

/*
 * Byte-aligned trit encoding.
 *
 * Each byte (0-255) → 6 base-3 digits → 6 trit pairs (12 bits).
 * Total trit bits for N bytes = N * 12.
 * Pad to byte boundary with 11 pairs.
 *
 * Example: 1 byte → 12 bits (already byte-aligned at 1.5 bytes,
 *          so we need 2 bytes = 16 bits, pad with two 11 pairs)
 */

size_t trit_encode_padded(const uint8_t *bytes, size_t byte_len,
                          uint8_t **out_buf)
{
    if (byte_len == 0) { *out_buf = NULL; return 0; }

    /* each input byte → 6 trit pairs = 12 bits */
    size_t total_bits = byte_len * 12;
    /* pad to byte boundary */
    size_t padded_bits = (total_bits + 7) & ~(size_t)7;
    size_t out_len = padded_bits / 8;

    uint8_t *buf = calloc(out_len, 1);
    if (!buf) { *out_buf = NULL; return 0; }

    size_t bit_pos = 0;

    for (size_t i = 0; i < byte_len; i++) {
        uint8_t val = bytes[i];
        uint8_t trits[6];
        for (int t = 5; t >= 0; t--) {
            trits[t] = val % 3;
            val /= 3;
        }
        for (int t = 0; t < 6; t++) {
            size_t byte_idx = bit_pos / 8;
            int shift = 6 - (bit_pos % 8);
            buf[byte_idx] |= (trits[t] & 0x03) << shift;
            bit_pos += 2;
        }
    }

    /* fill remaining bits with 11 (pad) */
    while (bit_pos < padded_bits) {
        size_t byte_idx = bit_pos / 8;
        int shift = 6 - (bit_pos % 8);
        buf[byte_idx] |= 0x03 << shift;
        bit_pos += 2;
    }

    *out_buf = buf;
    return out_len;
}

size_t trit_decode_padded(const uint8_t *buf, size_t buf_len,
                          uint8_t **out_bytes)
{
    if (buf_len == 0) { *out_bytes = NULL; return 0; }

    /* read trit pairs until we hit a 11 pair (pad) or run out */
    uint8_t *trits = malloc(buf_len * 4);  /* max trit pairs */
    if (!trits) { *out_bytes = NULL; return 0; }
    size_t trit_count = 0;

    int in_padding = 0;
    for (size_t bit_pos = 0; bit_pos < buf_len * 8; bit_pos += 2) {
        size_t byte_idx = bit_pos / 8;
        int shift = 6 - (bit_pos % 8);
        uint8_t pair = (buf[byte_idx] >> shift) & 0x03;
        if (pair == 3) {
            in_padding = 1;
            continue;
        }
        if (in_padding) {
            /* data pair after padding = malformed */
            fprintf(stderr, "dtob: data after padding at byte %zu\n", byte_idx);
            free(trits);
            *out_bytes = NULL;
            return 0;
        }
        trits[trit_count++] = pair;
    }

    /* trit count must be a multiple of 6 (padding aligns to byte boundaries) */
    if (trit_count % 6 != 0) {
        fprintf(stderr, "dtob: trit count %zu is not a multiple of 6\n", trit_count);
        free(trits);
        *out_bytes = NULL;
        return 0;
    }

    /* decode: every 6 trits → 1 byte */
    size_t n_bytes = trit_count / 6;
    uint8_t *bytes = malloc(n_bytes ? n_bytes : 1);
    if (!bytes) { free(trits); *out_bytes = NULL; return 0; }

    for (size_t i = 0; i < n_bytes; i++) {
        unsigned int val = 0;
        for (int t = 0; t < 6; t++) {
            val = val * 3 + trits[i * 6 + t];
        }
        if (val > 255) {
            fprintf(stderr, "dtob: trit decode produced value %u > 255\n", val);
            free(trits);
            free(bytes);
            *out_bytes = NULL;
            return 0;
        }
        bytes[i] = (uint8_t)val;
    }

    free(trits);
    *out_bytes = bytes;
    return n_bytes;
}
