//
// Created by bruno on 13. 7. 2026.
//

#include "decoder.h"

#include <SDL.h>
#include <SDL_audio.h>
#include <SDL_stdinc.h>
#include <math.h>
#include <stdio.h>
#include <string.h>

#include "texteditor.h"
#include "modem_settings.h"
#include "packet.h"
#include "protocol.h"

Decoder decoder;

float tone_energy(
    Decoder *dec,
    int index) {
    // Return the current accumulated Goertzel energy without clearing state.
    Goertzel *g = &dec->tones[index];
    return g->s1 * g->s1 + g->s2 * g->s2 - g->coeff * g->s1 * g->s2;
}

static void decoder_history_push(Decoder *dec, float sample) {
    dec->sampleHistory[dec->sampleHistoryWrite] = sample;
    dec->sampleHistoryWrite =
            (dec->sampleHistoryWrite + 1) % DECODER_SAMPLE_HISTORY;

    if (dec->sampleHistoryCount < DECODER_SAMPLE_HISTORY)
        dec->sampleHistoryCount++;

    dec->totalSamples++;
}

static bool decoder_history_get(const Decoder *dec, uint64_t sampleIndex, float *sample) {
    if (sampleIndex >= dec->totalSamples) {
        return false;
    }

    uint64_t oldest = dec->totalSamples - dec->sampleHistoryCount;

    if (sampleIndex < oldest) {
        return false;
    }

    uint32_t start =
            (dec->sampleHistoryWrite + DECODER_SAMPLE_HISTORY - dec->sampleHistoryCount)
            % DECODER_SAMPLE_HISTORY;

    uint32_t rel = (uint32_t) (sampleIndex - oldest);
    *sample = dec->sampleHistory[(start + rel) % DECODER_SAMPLE_HISTORY];
    return true;
}

void decoder_init(
    Decoder *dec) {
    memset(dec, 0, sizeof(*dec));


    dec->bitsPerSymbol = 3;
    dec->toneCount = 8;


    dec->startFreq = DEFAULT_START_FREQ;
    dec->endFreq = DEFAULT_END_FREQ;
    uint16_t freqRange = dec->endFreq - dec->startFreq;
    dec->dataSpacing =
            freqRange /
            (encoder.toneCount - 1);

    dec->sampleRate = SAMPLE_RATE;
    dec->symbolSamples =
            SAMPLE_RATE / DEFAULT_SYMBOL_RATE;


    dec->clock0 = goertzel_bin_freq(
        DEFAULT_CLOCK0_FREQ,
        dec->sampleRate,
        dec->symbolSamples
    );

    dec->clock1 = goertzel_bin_freq(
        DEFAULT_CLOCK1_FREQ,
        dec->sampleRate,
        dec->symbolSamples
    );

    decoder_rebuild_tones(dec);
}

void decoder_rebuild_tones(Decoder *dec) {
    dec->toneCount = 1 << dec->bitsPerSymbol;

    printf("Decoder tones: %d\n", dec->toneCount);


    // clock tones
    goertzel_init(
        &dec->tones[0],
        dec->clock0,
        dec->sampleRate,
        dec->symbolSamples
    );

    goertzel_init(
        &dec->tones[1],
        dec->clock1,
        dec->sampleRate,
        dec->symbolSamples
    );


    // data tones
    for (int i = 0; i < dec->toneCount; i++) {
        uint16_t freq =
                dec->startFreq +
                i * dec->dataSpacing;


        // MUST match encoder
        freq = goertzel_bin_freq(
            freq,
            dec->sampleRate,
            dec->symbolSamples
        );

        goertzel_init(
            &dec->tones[i + 2],
            freq,
            dec->sampleRate,
            dec->symbolSamples
        );


        printf(
            "Tone %d = %dHz\n",
            i,
            freq
        );
    }


    memset(
        dec->energies,
        0,
        sizeof(dec->energies)
    );


    dec->rxState = RX_SEARCH;
}

static int detect_clock(Decoder *dec) {
    // Use instantaneous per-symbol energies (dec->energies) for fast detection
    float e0 = dec->energies[0];
    float e1 = dec->energies[1];

    // Simple silence check
    float peak = fmaxf(e0, e1);
    if (peak < 1e-7f)
        return -1;

    // Determine which clock tone is stronger with a small hysteresis
    if (e1 > e0 * 1.05f)
        return 1;
    if (e0 > e1 * 1.05f)
        return 0;

    return -1; // uncertain
}

static void clock_search(
    Decoder *dec,
    int clock) {
    if (clock < 0) {
        dec->clockHistory = 0;
        dec->clockCount = 0;
        return;
    }


    dec->clockHistory =
            (dec->clockHistory << 1) |
            (clock & 1);


    if (dec->clockCount < 8)
        dec->clockCount++;


    if (dec->clockCount >= 8) {
        uint8_t h = dec->clockHistory & 0xff;


        if (h == 0x55 || h == 0xAA) {
            printf("CLOCK LOCK pattern=%02X\n", h);
            dec->goodClocks = 8;
            dec->clockHistory = 0;
            dec->clockCount = 0;

            return;
        }

        dec->goodClocks = 0;
    }
}

static void update_all_energies(Decoder *dec) {
    const int tones = dec->toneCount + 2;

    uint64_t oldest = dec->totalSamples - dec->sampleHistoryCount;
    if (dec->totalSamples < dec->symbolSamples)
        return; // not enough samples yet

    uint64_t start = dec->totalSamples - dec->symbolSamples;
    if (start < oldest) start = oldest;

    Goertzel work[MAX_FSK_TONES + 2];
    memcpy(work, dec->tones, sizeof(Goertzel) * tones);
    for (int i = 0; i < tones; ++i) {
        work[i].s1 = 0.0f;
        work[i].s2 = 0.0f;
    }

    for (uint32_t s = 0; s < dec->symbolSamples; ++s) {
        float sample;
        if (!decoder_history_get(dec, start + s, &sample)) { sample = 0.0f; }
        for (int t = 0; t < tones; ++t)
            goertzel_add(&work[t], sample);
    }

    float maxE = 0.0f;
    for (int t = 0; t < tones; ++t) {
        float e = goertzel_energy(&work[t]);
        // Instantaneous per-symbol energy
        dec->energies[t] = e;
        if (t == 0) dec->clockEnergy0 = e;
        if (t == 1) dec->clockEnergy1 = e;
        if (e > maxE) maxE = e;
    }

    // Faster decay so visualization follows symbol changes
    dec->maxEnergy = fmaxf(dec->maxEnergy * 0.8f, maxE);

    // Maintain a short visual smoothing for clocks so bars don't flicker
    if (!isfinite(dec->clockEnergy0)) dec->clockEnergy0 = dec->energies[0];
    if (!isfinite(dec->clockEnergy1)) dec->clockEnergy1 = dec->energies[1];
    dec->clockEnergy0 = dec->clockEnergy0 * 0.7f + dec->energies[0] * 0.3f;
    dec->clockEnergy1 = dec->clockEnergy1 * 0.7f + dec->energies[1] * 0.3f;

#ifdef DEBUG_ENERGIES
    int show = dec->toneCount + 2;
    int upto = show < 12 ? show : 12;
    printf("ENERGIES(start=%llu):", (unsigned long long) start);
    for (int i = 0; i < upto; ++i) {
        printf(" %d:%.6f", i, dec->energies[i]);
    }
    printf(" maxE=%.6f clk0=%.6f clk1=%.6f\n", dec->maxEnergy, dec->clockEnergy0, dec->clockEnergy1);
#endif
}

static int decode_symbol(
    Decoder *dec) {
    float bestEnergy = 0;
    int bestTone = 0;


    /*
       data tones start after clock tones
    */
    for (int i = 0; i < dec->toneCount; i++) {
        float e =
                dec->energies[i + 2];


        if (e > bestEnergy) {
            bestEnergy = e;
            bestTone = i;
        }
    }


    return bestTone;
}

bool decoder_push_byte(
    Decoder *dec,
    uint8_t b) {
    if (dec->rxState == RX_MAGIC) {
        dec->magicBuffer[dec->magicIndex++] = b;

        printf("MAGIC %02X\n", b);

        if (dec->magicIndex >= MAGIC_SIZE) {
            if (memcmp(dec->magicBuffer,
                       PACKET_MAGIC,
                       MAGIC_SIZE) == 0) {
                printf("MAGIC OK\n");

                dec->rxState = RX_DATA;

                dec->bufferSize = 0;
                dec->currentByte = 0;
                dec->bitIndex = 0;
            }

            printf("BAD MAGIC\n");

            dec->magicIndex = 0;
            dec->preambleCount = 0;
            dec->rxState = RX_SEARCH;
        }

        return false;
    }


    if (dec->rxState == RX_DATA) {
        if (dec->bufferSize < sizeof(dec->buffer)) {
            dec->buffer[dec->bufferSize++] = b;
        }

        if(decoder_parse_packet(dec))
        {
            printf("PACKET OK\n");

            decoder_reset_rx(dec);
        }
    }

    return false;
}

static void decoder_push_symbol(
    Decoder *dec,
    unsigned symbol) {
    for (int i = 0; i < dec->bitsPerSymbol; i++) {
        int bit = (symbol >> i) & 1;

        dec->currentByte |= bit << dec->bitIndex;

        if (++dec->bitIndex == 8) {
            printf("BYTE %02X\n", dec->currentByte);

            decoder_push_byte(
                dec,
                dec->currentByte);

            dec->currentByte = 0;
            dec->bitIndex = 0;
        }
    }
}

static bool clock_valid(Decoder *dec) {
    float e0 = dec->energies[0];
    float e1 = dec->energies[1];

    float peak = fmaxf(e0, e1);

    if (peak < dec->maxEnergy * 0.10f)
        return false;

    float ratio = fmaxf(e0, e1) /
                  (fminf(e0, e1) + 1e-12f);

    return ratio > 1.3f;
}

void decoder_reset_rx(Decoder *dec) {
    dec->rxState = RX_SEARCH;
    dec->locked = false;

    dec->preambleCount = 0;
    dec->magicIndex = 0;

    dec->bufferSize = 0;
    dec->currentByte = 0;
    dec->bitIndex = 0;

    dec->clockHistory = 0;
    dec->clockCount = 0;
    dec->goodClocks = 0;
}

void decoder_process(
    Decoder *dec,
    float *samples,
    uint32_t count) {
    const int tones = dec->toneCount + 2;

    for (uint32_t i = 0; i < count; i++) {
        decoder_history_push(dec, samples[i]);

        // Always feed samples to Goertzel for continuous analysis
        for (int t = 0; t < tones; t++)
            goertzel_add(&dec->tones[t], samples[i]);

        dec->sampleCount += 1;

        if (dec->sampleCount != dec->symbolSamples)
            continue;

        dec->sampleCount = 0;

        switch (dec->rxState) {
            case RX_SEARCH: {
                // Update all energies for real-time visual feedback
                update_all_energies(dec);

                int clock = detect_clock(dec);
                clock_search(dec, clock);

                if (dec->goodClocks >= 8) {
                    printf("CLOCK LOCKED\n");

                    dec->rxState = RX_SYNC;
                    dec->locked = true;

                    dec->clockHistory = 0;
                    dec->clockCount = 0;

                    dec->goodClocks = 0;
                }

                // DON'T reset filters - keep energy visible
                break;
            }

            case RX_SYNC: {
                update_all_energies(dec);


                int symbol = decode_symbol(dec);

                if (clock_valid(dec)) {
                    dec->lostClockCount = 0;
                    printf("Symbol %d  energy=%f\n", symbol, dec->energies[symbol + 2]);
                } else {
                    if (++dec->lostClockCount > 8) {
                        printf("CLOCK LOST\n");

                        decoder_reset_rx(dec);

                        break;
                    }
                }


                if (symbol == (dec->preambleCount & 7)) {
                    if (dec->preambleCount < 255)
                        dec->preambleCount++;

                    if (dec->preambleCount >= 16) {
                        printf("PREAMBLE OK\n");

                        dec->rxState = RX_MAGIC;

                        dec->magicIndex = 0;
                        dec->currentByte = 0;
                        dec->bitIndex = 0;

                        break;
                    }
                } else {
                    if (dec->preambleCount > 4)
                        dec->preambleCount -= 4;
                    else
                        dec->preambleCount = 0;
                }

                break;
            }
            case RX_MAGIC:
            case RX_DATA: {
                update_all_energies(dec);

                if (!clock_valid(dec)) {
                    if (++dec->lostClockCount > 8) {
                        printf("CLOCK LOST\n");
                        decoder_reset_rx(dec);
                        break;
                    }
                } else {
                    dec->lostClockCount = 0;
                }

                int symbol = decode_symbol(dec);

                decoder_push_symbol(dec, symbol);

                break;
            }
        }
    }
}

void input_callback(
    void *userdata,
    Uint8 *stream,
    int len) {
    Decoder *dec = userdata;

    // stream contains audio captured from sound card input
    float *samples = (float *) stream;
    int count = len / sizeof(float);

    decoder_process(
        dec,
        samples,
        count
    );
}

SDL_AudioDeviceID inputDev;
AudioData decodedAudioData;

void initAudioDecoder(SDL_Renderer *renderer) {
    memset(&decodedAudioData, 0, sizeof(AudioData));
    decoder_init(&decoder);
    decoder.renderer = renderer;

    SDL_AudioSpec inSpec = {0};
    inSpec.freq = SAMPLE_RATE;
    inSpec.format = AUDIO_F32SYS;
    inSpec.channels = 1;
    inSpec.samples = 960; // Match symbolSamples (48000/50)
    inSpec.callback = input_callback;
    inSpec.userdata = &decoder;

    SDL_AudioSpec realSpec = {0};

    inputDev = SDL_OpenAudioDevice(NULL, 1, &inSpec, &realSpec, 0);
    printf("Decoder device opened: inputDev=%u\n", inputDev);
    if (inputDev == 0) {
        printf("Failed to open audio decoder: %s\n", SDL_GetError());
        SDL_Quit();
    }
    printf("Input device spec: freq=%d channels=%d samples=%d format=%d\n",
           realSpec.freq, realSpec.channels, realSpec.samples, realSpec.format);

    // Update decoder sample rate and symbolSamples based on actual device
    decoder.sampleRate = realSpec.freq;
    decoder.symbolSamples = realSpec.freq / DEFAULT_SYMBOL_RATE;
    // Rebuild tones to match actual sample rate
    decoder_rebuild_tones(&decoder);

    SDL_PauseAudioDevice(inputDev, 0);
}