/*
    * Sdr.cpp
    * Generic software-defined radio receive subsystem.
    * Copyright (c) 2026 Daniel Hammer
*/

#include "Sdr.hpp"
#include <Memory/Heap.hpp>
#include <Libraries/Memory.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>

using namespace Kt;

namespace Drivers::Radio::Sdr {

    // I/Q ring size per receiver.  256 KiB is ~62 ms of jitter buffer at
    // 2.048 Msps (2 bytes/sample), which comfortably absorbs scheduling gaps
    // between a userspace reader's polls.
    static constexpr uint32_t RING_BYTES = 256 * 1024;

    struct Receiver {
        bool          used;
        bool          opened;
        bool          streaming;

        char          name[64];
        char          tuner[32];
        char          serial[32];
        uint64_t      freqMin, freqMax;
        uint32_t      sampleRateMin, sampleRateMax;
        int           gains[MAX_GAINS];
        uint32_t      numGains;
        uint8_t       format;

        ReceiverOps   ops;
        void*         ctx;

        // Last-requested configuration (cached for GETPARAM readback).
        uint64_t      freq;
        uint32_t      sampleRate;
        int           gainMode;     // 0 = auto, 1 = manual
        int           gain;         // tenths of dB
        int           ppm;
        int           agc;
        int           directSamp;

        // I/Q ring buffer (byte FIFO).
        uint8_t*      ring;
        uint32_t      head;         // write position
        uint32_t      count;        // bytes currently queued
        uint64_t      totalBytes;   // lifetime sample bytes delivered
        uint64_t      droppedBytes; // bytes dropped on overflow
        kcp::Spinlock lock;
    };

    static Receiver g_rx[MAX_RECEIVERS];

    // -------------------------------------------------------------------------
    // Helpers
    // -------------------------------------------------------------------------

    static void CopyStr(char* dst, uint32_t cap, const char* src) {
        uint32_t i = 0;
        if (src) {
            for (; i < cap - 1 && src[i]; i++) dst[i] = src[i];
        }
        dst[i] = '\0';
    }

    static Receiver* Lookup(int handle, bool needOpen) {
        if (handle < 0 || handle >= MAX_RECEIVERS) return nullptr;
        Receiver& r = g_rx[handle];
        if (!r.used) return nullptr;
        if (needOpen && !r.opened) return nullptr;
        return &r;
    }

    // -------------------------------------------------------------------------
    // Driver-facing API
    // -------------------------------------------------------------------------

    int Register(const ReceiverDesc& desc) {
        for (int i = 0; i < MAX_RECEIVERS; i++) {
            if (g_rx[i].used) continue;
            Receiver& r = g_rx[i];

            // Reset everything except the (non-copyable) spinlock instance.
            r.opened = false;
            r.streaming = false;
            CopyStr(r.name, sizeof(r.name), desc.name);
            CopyStr(r.tuner, sizeof(r.tuner), desc.tuner);
            CopyStr(r.serial, sizeof(r.serial), desc.serial);
            r.freqMin = desc.freqMin;
            r.freqMax = desc.freqMax;
            r.sampleRateMin = desc.sampleRateMin;
            r.sampleRateMax = desc.sampleRateMax;
            r.numGains = desc.numGains > MAX_GAINS ? MAX_GAINS : desc.numGains;
            for (uint32_t g = 0; g < r.numGains; g++) r.gains[g] = desc.gains[g];
            r.format = desc.format;
            r.ops = desc.ops;
            r.ctx = desc.ctx;

            r.freq = (desc.freqMin + desc.freqMax) / 2;
            r.sampleRate = desc.sampleRateMax;
            r.gainMode = 0;
            r.gain = 0;
            r.ppm = 0;
            r.agc = 0;
            r.directSamp = 0;

            r.ring = nullptr;
            r.head = r.count = 0;
            r.totalBytes = r.droppedBytes = 0;

            r.used = true;   // publish last
            KernelLogStream(OK, "SDR") << "Registered receiver " << (uint64_t)i
                << ": " << r.name << " / " << r.tuner;
            return i;
        }
        KernelLogStream(WARNING, "SDR") << "No free receiver slot for " << desc.name;
        return -1;
    }

    void Unregister(int idx) {
        if (idx < 0 || idx >= MAX_RECEIVERS) return;
        Receiver& r = g_rx[idx];
        if (!r.used) return;

        if (r.streaming && r.ops.Stop) r.ops.Stop(r.ctx);

        r.lock.Acquire();
        r.streaming = false;
        r.opened = false;
        r.used = false;
        uint8_t* ring = r.ring;
        r.ring = nullptr;
        r.head = r.count = 0;
        r.lock.Release();

        if (ring) Memory::g_heap->Free(ring);
        KernelLogStream(INFO, "SDR") << "Unregistered receiver " << (uint64_t)idx;
    }

    void PushSamples(int idx, const uint8_t* data, uint32_t len) {
        if (idx < 0 || idx >= MAX_RECEIVERS || !data || len == 0) return;
        Receiver& r = g_rx[idx];

        r.lock.Acquire();
        // Re-validate under the lock: Unregister() clears these and frees the
        // ring while holding the same lock, so an in-flight USB completion can
        // never write into a freed buffer.
        if (!r.used || !r.ring) { r.lock.Release(); return; }

        uint32_t space = RING_BYTES - r.count;
        uint32_t n = len;
        uint32_t dropped = 0;
        if (n > space) {
            // Truncate to a whole number of I/Q byte pairs: dropping an odd
            // count would swap I and Q for the rest of the stream.
            n = space & ~1u;
            dropped = len - n;
        }

        uint32_t first = RING_BYTES - r.head;
        if (first > n) first = n;
        memcpy(r.ring + r.head, data, first);
        if (n > first) memcpy(r.ring, data + first, n - first);

        r.head = (r.head + n) % RING_BYTES;
        r.count += n;
        r.totalBytes += n;
        r.droppedBytes += dropped;
        r.lock.Release();
    }

    bool IsStreaming(int idx) {
        if (idx < 0 || idx >= MAX_RECEIVERS) return false;
        return g_rx[idx].used && g_rx[idx].streaming;
    }

    // -------------------------------------------------------------------------
    // Syscall-facing API
    // -------------------------------------------------------------------------

    int Count() {
        int n = 0;
        for (int i = 0; i < MAX_RECEIVERS; i++) if (g_rx[i].used) n++;
        return n;
    }

    bool GetInfo(int idx, montauk::abi::SdrDeviceInfo* out) {
        Receiver* r = Lookup(idx, false);
        if (!r || !out) return false;

        memset(out, 0, sizeof(*out));
        CopyStr(out->name, sizeof(out->name), r->name);
        CopyStr(out->tuner, sizeof(out->tuner), r->tuner);
        CopyStr(out->serial, sizeof(out->serial), r->serial);
        out->freqMin = r->freqMin;
        out->freqMax = r->freqMax;
        out->sampleRateMin = r->sampleRateMin;
        out->sampleRateMax = r->sampleRateMax;
        out->numGains = r->numGains;
        for (uint32_t g = 0; g < r->numGains && g < 32; g++) out->gains[g] = r->gains[g];
        out->sampleFormat = r->format;
        out->present = 1;
        out->streaming = r->streaming ? 1 : 0;
        return true;
    }

    int Open(int idx) {
        Receiver* r = Lookup(idx, false);
        if (!r) return -1;

        // Single-user OS: an Open always claims the device, reclaiming it from a
        // previous owner that exited without closing.
        if (r->streaming && r->ops.Stop) r->ops.Stop(r->ctx);

        if (!r->ring) {
            r->ring = (uint8_t*)Memory::g_heap->Request(RING_BYTES);
            if (!r->ring) {
                KernelLogStream(ERROR, "SDR") << "Ring alloc failed for receiver "
                    << (uint64_t)idx;
                return -1;
            }
        }

        r->lock.Acquire();
        r->head = r->count = 0;
        r->lock.Release();
        r->streaming = false;
        r->opened = true;
        return idx;   // handle == index
    }

    int Close(int handle) {
        Receiver* r = Lookup(handle, true);
        if (!r) return -1;
        if (r->streaming && r->ops.Stop) r->ops.Stop(r->ctx);
        r->streaming = false;
        r->opened = false;
        return 0;
    }

    int Start(int handle) {
        Receiver* r = Lookup(handle, true);
        if (!r) return -1;
        // Already streaming: a second Start must not re-arm the driver's
        // transfer pool (it would double-queue every buffer).
        if (r->streaming) return 0;

        r->lock.Acquire();
        r->head = r->count = 0;   // discard stale samples before (re)starting
        r->lock.Release();

        int rc = r->ops.Start ? r->ops.Start(r->ctx) : -1;
        if (rc == 0) r->streaming = true;
        return rc;
    }

    int Stop(int handle) {
        Receiver* r = Lookup(handle, true);
        if (!r) return -1;
        int rc = r->ops.Stop ? r->ops.Stop(r->ctx) : 0;
        r->streaming = false;
        return rc;
    }

    int Read(int handle, uint8_t* buf, uint32_t len) {
        Receiver* r = Lookup(handle, true);
        if (!r || !buf || !r->ring) return -1;
        if (len == 0) return 0;

        // Give the driver a process-context tick (e.g. USB stall recovery)
        // before draining; do this outside the ring lock since it may issue
        // blocking USB commands.
        if (r->streaming && r->ops.Service) r->ops.Service(r->ctx);

        r->lock.Acquire();
        uint32_t n = r->count < len ? r->count : len;
        uint32_t tail = (r->head + RING_BYTES - r->count) % RING_BYTES;
        uint32_t first = RING_BYTES - tail;
        if (first > n) first = n;
        memcpy(buf, r->ring + tail, first);
        if (n > first) memcpy(buf + first, r->ring, n - first);
        r->count -= n;
        r->lock.Release();
        return (int)n;
    }

    uint32_t Available(int handle) {
        Receiver* r = Lookup(handle, true);
        if (!r) return 0;
        return r->count;
    }

    int64_t SetParam(int handle, int param, uint64_t value) {
        Receiver* r = Lookup(handle, true);
        if (!r) return -1;

        switch (param) {
            case montauk::abi::SDR_PARAM_FREQ:
                if (!r->ops.SetFreq) return -1;
                if (r->ops.SetFreq(r->ctx, value) != 0) return -1;
                r->freq = value;
                return 0;
            case montauk::abi::SDR_PARAM_SAMPLE_RATE:
                if (!r->ops.SetSampleRate) return -1;
                if (r->ops.SetSampleRate(r->ctx, (uint32_t)value) != 0) return -1;
                r->sampleRate = (uint32_t)value;
                return 0;
            case montauk::abi::SDR_PARAM_GAIN_MODE:
                if (!r->ops.SetGainMode) return -1;
                if (r->ops.SetGainMode(r->ctx, (int)value) != 0) return -1;
                r->gainMode = (int)value ? 1 : 0;
                return 0;
            case montauk::abi::SDR_PARAM_GAIN:
                if (!r->ops.SetGain) return -1;
                if (r->ops.SetGain(r->ctx, (int)(int64_t)value) != 0) return -1;
                r->gain = (int)(int64_t)value;
                return 0;
            case montauk::abi::SDR_PARAM_FREQ_CORR:
                if (!r->ops.SetFreqCorrection) return -1;
                if (r->ops.SetFreqCorrection(r->ctx, (int)(int64_t)value) != 0) return -1;
                r->ppm = (int)(int64_t)value;
                return 0;
            case montauk::abi::SDR_PARAM_AGC:
                if (!r->ops.SetAgc) return -1;
                if (r->ops.SetAgc(r->ctx, (int)value) != 0) return -1;
                r->agc = (int)value ? 1 : 0;
                return 0;
            case montauk::abi::SDR_PARAM_DIRECT_SAMP:
                if (!r->ops.SetDirectSampling) return -1;
                if (r->ops.SetDirectSampling(r->ctx, (int)value) != 0) return -1;
                r->directSamp = (int)value;
                return 0;
            default:
                return -1;
        }
    }

    int64_t GetParam(int handle, int param) {
        Receiver* r = Lookup(handle, true);
        if (!r) return -1;

        switch (param) {
            case montauk::abi::SDR_PARAM_FREQ:        return (int64_t)r->freq;
            case montauk::abi::SDR_PARAM_SAMPLE_RATE: return (int64_t)r->sampleRate;
            case montauk::abi::SDR_PARAM_GAIN_MODE:   return r->gainMode;
            case montauk::abi::SDR_PARAM_GAIN:        return r->gain;
            case montauk::abi::SDR_PARAM_FREQ_CORR:   return r->ppm;
            case montauk::abi::SDR_PARAM_AGC:         return r->agc;
            case montauk::abi::SDR_PARAM_DIRECT_SAMP: return r->directSamp;
            default:                                  return -1;
        }
    }

}
