#pragma once

/*
 * modem.h -- inline UART encode/decode over PulseAudio volume
 *
 * Shared by tx.c, rx.c, chat.c.
 * Depends on zebra.h (zebra_pulse_t, zebra_set_volume, zebra_get_volume,
 * zebra_bit_to_vol, zebra_vol_to_bit, ZEBRA_VOL_THRESHOLD).
 */

#include <time.h>
#include <stdint.h>
#include "zebra.h"

/* callback invoked by zebra_rx_run for each decoded byte */
typedef void (*zebra_byte_cb)(uint8_t byte, void *userdata);

/* ------------------------------------------------------------------ *
 * timing helper                                                        *
 * ------------------------------------------------------------------ */

static inline void ts_add_ns(struct timespec *ts, long ns) {
    ts->tv_nsec += ns;
    if (ts->tv_nsec >= 1000000000L) {
        ts->tv_sec  += ts->tv_nsec / 1000000000L;
        ts->tv_nsec %= 1000000000L;
    }
}

/* ------------------------------------------------------------------ *
 * TX: UART framing (start + 8 data LSB-first + stop)                  *
 * ------------------------------------------------------------------ */

/* channels: from zebra_sink_t.channels — avoids PA channel-count query per symbol */
static inline void zebra_send_symbol(zebra_pulse_t *z, uint32_t sink,
                                      uint8_t channels, int bit,
                                      struct timespec *next, long period_ns) {
    zebra_set_volume_noack(z, sink, channels, zebra_bit_to_vol(bit));
    clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, next, NULL);
    ts_add_ns(next, period_ns);
}

static inline void zebra_send_byte(zebra_pulse_t *z, uint32_t sink,
                                    uint8_t channels, uint8_t byte,
                                    struct timespec *next, long period_ns) {
    zebra_send_symbol(z, sink, channels, 0, next, period_ns);          /* start */
    for (int i = 0; i < 8; i++)
        zebra_send_symbol(z, sink, channels, (byte >> i) & 1, next, period_ns);
    zebra_send_symbol(z, sink, channels, 1, next, period_ns);          /* stop  */
}

/* ------------------------------------------------------------------ *
 * RX: 2x oversampled UART decoder — runs forever, calls cb per byte   *
 * ------------------------------------------------------------------ */

static inline void zebra_rx_run(zebra_pulse_t *z, uint32_t sink, int baud,
                                 zebra_byte_cb cb, void *userdata) {
    const int oversample = 4;
    long quarter_ns = 1000000000L / ((long)baud * oversample);
    long half_ns    = 2 * quarter_ns;
    long full_ns    = 4 * quarter_ns;

    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    int prev = 1;   /* assume MARK (idle) */

    for (;;) {
        ts_add_ns(&ts, quarter_ns);
        clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, &ts, NULL);

        uint8_t vol;
        if (zebra_get_volume(z, sink, &vol) < 0) {
            /* sink gone — wait and retry */
            struct timespec retry = {1, 0};
            nanosleep(&retry, NULL);
            prev = 1;
            continue;
        }

        int cur = zebra_vol_to_bit(vol);

        /* MARK→SPACE falling edge = start bit */
        if (prev == 1 && cur == 0) {
            /* advance to center of start bit and confirm */
            ts_add_ns(&ts, half_ns);
            clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, &ts, NULL);
            if (zebra_get_volume(z, sink, &vol) < 0)  { prev = 1; continue; }
            if (zebra_vol_to_bit(vol) != 0)            { prev = 1; continue; }

            /* sample 8 data bits */
            uint8_t byte = 0;
            int ok = 1;
            for (int i = 0; i < 8; i++) {
                ts_add_ns(&ts, full_ns);
                clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, &ts, NULL);
                if (zebra_get_volume(z, sink, &vol) < 0) { ok = 0; break; }
                int b = zebra_vol_to_bit(vol);
                if (b < 0) b = (vol >= ZEBRA_VOL_THRESHOLD) ? 1 : 0;
                byte |= (uint8_t)(b << i);
            }
            if (ok) cb(byte, userdata);
            /* Resync sample clock to real time after each byte.
             * Prevents accumulated edge-detection error from shifting
             * data bit samples in subsequent bytes.                   */
            clock_gettime(CLOCK_MONOTONIC, &ts);
            prev = 1;

        } else if (cur >= 0) {
            prev = cur;
        }
    }
}

/* ------------------------------------------------------------------ *
 * PURE HELPERS — no PA, fully testable                                *
 * ------------------------------------------------------------------ */

/* Compute max safe baud from avg PA poll latency (nanoseconds).
 * Formula: 1e9 / (avg_ns * 4x_oversample) * 0.8_safety = 2e8 / avg_ns */
static inline int zebra_baud_from_avg_ns(long avg_ns) {
    if (avg_ns <= 0) return ZEBRA_BAUD_DEFAULT;
    int baud = (int)(200000000L / avg_ns);
    if (baud < ZEBRA_BAUD_MIN) baud = ZEBRA_BAUD_MIN;
    if (baud > ZEBRA_BAUD_MAX) baud = ZEBRA_BAUD_MAX;
    return baud;
}

/* Build a BAUD_OFFER handshake frame into buf[ZEBRA_HS_FRAME_LEN]. */
static inline void zebra_hs_build(uint8_t frame[ZEBRA_HS_FRAME_LEN], uint16_t baud) {
    frame[0] = ZEBRA_HS_MAGIC_0;
    frame[1] = ZEBRA_HS_MAGIC_1;
    frame[2] = ZEBRA_HS_TYPE_OFFER;
    frame[3] = (uint8_t)(baud & 0xFF);
    frame[4] = (uint8_t)(baud >> 8);
    frame[5] = frame[0] ^ frame[1] ^ frame[2] ^ frame[3] ^ frame[4];
}

/* Build a READY frame (type=0x02, baud=0) into buf[ZEBRA_HS_FRAME_LEN]. */
static inline void zebra_hs_build_ready(uint8_t frame[ZEBRA_HS_FRAME_LEN]) {
    frame[0] = ZEBRA_HS_MAGIC_0;
    frame[1] = ZEBRA_HS_MAGIC_1;
    frame[2] = ZEBRA_HS_TYPE_READY;
    frame[3] = 0;
    frame[4] = 0;
    frame[5] = frame[0] ^ frame[1] ^ frame[2] ^ frame[3] ^ frame[4];
}

/* Validate and parse a handshake frame.
 * Returns 0 and sets *baud on success; -1 on bad magic, type, or checksum. */
static inline int zebra_hs_parse(const uint8_t frame[ZEBRA_HS_FRAME_LEN],
                                   uint16_t *baud) {
    if (frame[0] != ZEBRA_HS_MAGIC_0 || frame[1] != ZEBRA_HS_MAGIC_1) return -1;
    if (frame[2] != ZEBRA_HS_TYPE_OFFER) return -1;
    uint8_t ck = frame[0] ^ frame[1] ^ frame[2] ^ frame[3] ^ frame[4];
    if (ck != frame[5]) return -1;
    *baud = (uint16_t)(frame[3] | ((uint16_t)frame[4] << 8));
    return 0;
}

/* ------------------------------------------------------------------ *
 * BENCHMARK: measure PA poll latency, derive max safe baud            *
 * ------------------------------------------------------------------ */

/* Runs N zebra_get_volume calls and measures average round-trip time.
 * Returns max baud receiver can sustain via zebra_baud_from_avg_ns.  */
static inline int zebra_benchmark_baud(zebra_pulse_t *z, uint32_t sink) {
    const int N = 100;
    uint8_t vol;
    struct timespec t0, t1;

    clock_gettime(CLOCK_MONOTONIC, &t0);
    for (int i = 0; i < N; i++)
        zebra_get_volume(z, sink, &vol);
    clock_gettime(CLOCK_MONOTONIC, &t1);

    long elapsed_ns = (t1.tv_sec - t0.tv_sec) * 1000000000L
                    + (t1.tv_nsec - t0.tv_nsec);
    return zebra_baud_from_avg_ns(elapsed_ns / N);
}

/* ------------------------------------------------------------------ *
 * HANDSHAKE TX: send negotiation frame at ZEBRA_BAUD_HANDSHAKE        *
 * ------------------------------------------------------------------ */

/* Frame layout (ZEBRA_HS_FRAME_LEN = 6 bytes):
 *   [0] 0x5A  magic 'Z'
 *   [1] 0x42  magic 'B'
 *   [2] 0x01  type: BAUD_OFFER
 *   [3] baud low byte   (uint16 little-endian)
 *   [4] baud high byte
 *   [5] XOR of bytes 0-4 (checksum)                                  */
static inline int zebra_send_handshake(zebra_pulse_t *z, uint32_t sink,
                                        uint8_t channels, uint16_t baud) {
    uint8_t frame[ZEBRA_HS_FRAME_LEN];
    zebra_hs_build(frame, baud);

    long period_ns = 1000000000L / ZEBRA_BAUD_HANDSHAKE;
    zebra_set_volume_fast(z, sink, channels, ZEBRA_VOL_MARK);

    struct timespec next;
    clock_gettime(CLOCK_MONOTONIC, &next);
    ts_add_ns(&next, period_ns);

    for (int i = 0; i < ZEBRA_HS_FRAME_LEN; i++)
        zebra_send_byte(z, sink, channels, frame[i], &next, period_ns);

    zebra_set_volume_fast(z, sink, channels, ZEBRA_VOL_MARK);
    return 0;
}

/* ------------------------------------------------------------------ *
 * HANDSHAKE RX: listen for any HS frame type with timeout            *
 * ------------------------------------------------------------------ */

/* General frame receiver: listens at ZEBRA_BAUD_HANDSHAKE for a frame
 * whose type byte matches expected_type. Returns 0 and sets *out_baud
 * (may be NULL for READY frames where baud=0) on success; -1 on timeout.
 * Uses sliding-window magic-byte sync so partial frame receipt is OK.  */
static inline int zebra_recv_hs_frame(zebra_pulse_t *z, uint32_t sink,
                                       int timeout_ms, uint8_t expected_type,
                                       uint16_t *out_baud) {
    const int oversample = 4;
    long quarter_ns = 1000000000L / ((long)ZEBRA_BAUD_HANDSHAKE * oversample);
    long half_ns    = 2 * quarter_ns;
    long full_ns    = 4 * quarter_ns;

    struct timespec deadline, ts, now;
    clock_gettime(CLOCK_MONOTONIC, &deadline);
    ts_add_ns(&deadline, (long)timeout_ms * 1000000L);
    clock_gettime(CLOCK_MONOTONIC, &ts);

    int prev = 1;
    uint8_t frame[ZEBRA_HS_FRAME_LEN];
    int fpos = 0;

    for (;;) {
        clock_gettime(CLOCK_MONOTONIC, &now);
        if (now.tv_sec > deadline.tv_sec ||
            (now.tv_sec == deadline.tv_sec && now.tv_nsec >= deadline.tv_nsec))
            return -1;

        ts_add_ns(&ts, quarter_ns);
        clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, &ts, NULL);

        uint8_t vol;
        if (zebra_get_volume(z, sink, &vol) < 0) {
            struct timespec r = {0, 10000000L};   /* 10ms retry */
            nanosleep(&r, NULL);
            prev = 1;
            continue;
        }

        int cur = zebra_vol_to_bit(vol);

        if (prev == 1 && cur == 0) {
            /* start bit — confirm at center */
            ts_add_ns(&ts, half_ns);
            clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, &ts, NULL);
            if (zebra_get_volume(z, sink, &vol) < 0) { prev = 1; continue; }
            if (zebra_vol_to_bit(vol) != 0)           { prev = 1; continue; }

            /* decode 8 data bits */
            uint8_t byte = 0;
            int ok = 1;
            for (int i = 0; i < 8; i++) {
                ts_add_ns(&ts, full_ns);
                clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, &ts, NULL);
                if (zebra_get_volume(z, sink, &vol) < 0) { ok = 0; break; }
                int b = zebra_vol_to_bit(vol);
                if (b < 0) b = (vol >= ZEBRA_VOL_THRESHOLD) ? 1 : 0;
                byte |= (uint8_t)(b << i);
            }
            if (!ok) { clock_gettime(CLOCK_MONOTONIC, &ts); prev = 1; continue; }

            /* sliding frame sync on magic bytes */
            if (fpos == 0) {
                if (byte == ZEBRA_HS_MAGIC_0) frame[fpos++] = byte;
            } else if (fpos == 1) {
                if      (byte == ZEBRA_HS_MAGIC_1) frame[fpos++] = byte;
                else if (byte == ZEBRA_HS_MAGIC_0) { fpos = 1; frame[0] = byte; }
                else                                fpos = 0;
            } else {
                frame[fpos++] = byte;
                if (fpos == ZEBRA_HS_FRAME_LEN) {
                    /* validate: magic + expected_type + checksum */
                    uint8_t ck = frame[0]^frame[1]^frame[2]^frame[3]^frame[4];
                    if (frame[2] == expected_type && ck == frame[5]) {
                        if (out_baud)
                            *out_baud = (uint16_t)(frame[3]|((uint16_t)frame[4]<<8));
                        return 0;
                    }
                    fpos = 0;
                }
            }
            clock_gettime(CLOCK_MONOTONIC, &ts);   /* resync after each byte */
            prev = 1;
        } else if (cur >= 0) {
            prev = cur;
        }
    }
}

/* Receive BAUD_OFFER frame. Wrapper around zebra_recv_hs_frame. */
static inline int zebra_recv_handshake(zebra_pulse_t *z, uint32_t sink,
                                        int timeout_ms, uint16_t *out_baud) {
    return zebra_recv_hs_frame(z, sink, timeout_ms, ZEBRA_HS_TYPE_OFFER, out_baud);
}

/* ------------------------------------------------------------------ *
 * HANDSHAKE READY: 3-way handshake completion                        *
 *                                                                      *
 * After sending BAUD_OFFER, RX sends 3x READY frames then enters its  *
 * receive loop immediately. TX waits for READY before sending data.   *
 * This eliminates the settle-timer race at high baud rates.           *
 * ------------------------------------------------------------------ */

/* Send READY frame 3x on sink so TX catches it even with scheduling jitter.
 * RX calls this immediately after zebra_send_handshake, then enters rx loop. */
static inline int zebra_send_ready(zebra_pulse_t *z, uint32_t sink,
                                    uint8_t channels) {
    uint8_t frame[ZEBRA_HS_FRAME_LEN];
    zebra_hs_build_ready(frame);
    long period_ns = 1000000000L / ZEBRA_BAUD_HANDSHAKE;

    for (int rep = 0; rep < 3; rep++) {
        zebra_set_volume_fast(z, sink, channels, ZEBRA_VOL_MARK);
        struct timespec next;
        clock_gettime(CLOCK_MONOTONIC, &next);
        ts_add_ns(&next, period_ns);
        for (int i = 0; i < ZEBRA_HS_FRAME_LEN; i++)
            zebra_send_byte(z, sink, channels, frame[i], &next, period_ns);
    }
    zebra_set_volume_fast(z, sink, channels, ZEBRA_VOL_MARK);
    return 0;
}

/* Wait for READY frame from RX. Returns 0 on success, -1 on timeout.
 * TX calls this after receiving BAUD_OFFER. Data send follows immediately. */
static inline int zebra_recv_ready(zebra_pulse_t *z, uint32_t sink,
                                    int timeout_ms) {
    return zebra_recv_hs_frame(z, sink, timeout_ms, ZEBRA_HS_TYPE_READY, NULL);
}