/*
    * Mixer.cpp
    * Software audio mixer: routes multiple per-process audio streams to a single
    * hardware output (Intel HDA), with per-stream volume, mute, and pause.
    * Copyright (c) 2026 Daniel Hammer
*/

#include "Mixer.hpp"
#include "IntelHda.hpp"
#include <Memory/PageFrameAllocator.hpp>
#include <Memory/HHDM.hpp>
#include <Sched/Scheduler.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <CppLib/Spinlock.hpp>
#include <Libraries/Memory.hpp>

namespace Drivers::Audio::Mixer {

    using namespace Kt;

    // =========================================================================
    // Per-stream state
    // =========================================================================
    //
    // Each virtual stream stores its input as 16-bit stereo at the *input*
    // sample rate. The pump consumes input frames at the input rate, linearly
    // interpolates between adjacent frames using a Q32 fractional cursor, and
    // produces output frames at MIX_RATE.

    static constexpr uint32_t INPUT_RING_FRAMES = 16384;  // ~341 ms at 48 kHz
    static constexpr int      INPUT_RING_BYTES  = INPUT_RING_FRAMES * 4;
    static constexpr int      INPUT_RING_PAGES  = (INPUT_RING_BYTES + 0xFFF) / 0x1000;

    // Cap one pump cycle to keep loop bounded under heavy write bursts.
    static constexpr uint32_t MAX_PUMP_FRAMES = 4096;     // ~85 ms at 48 kHz

    struct VirtualStream {
        bool active;
        int  ownerPid;
        char name[64];

        uint32_t inputRate;
        uint8_t  inputChannels;   // 1 or 2 (mono is upmixed at conversion time)
        uint8_t  inputBits;       // currently must be 16
        uint8_t  volume;          // 0-100
        bool     muted;
        bool     paused;

        // Input ring: int16 stereo at inputRate.
        int16_t* ring;
        uint32_t writeFrame;       // monotonically increasing, modulo'd at access
        uint32_t readFrame;
        uint64_t consumedInputBytes;  // total input bytes consumed by pump

        // Resample cursor (Q32 fixed point, fractional position within next input frame).
        uint64_t posQ32;
        uint64_t stepQ32;          // (inputRate << 32) / MIX_RATE
    };

    // =========================================================================
    // Module state
    // =========================================================================

    static kcp::Spinlock g_lock;
    static VirtualStream g_streams[MAX_STREAMS] = {};
    static bool g_hdaOpened = false;
    static int  g_hdaHandle = -1;
    static int  g_masterVolume = 80;
    static bool g_masterMute = false;
    static int  g_activeCount = 0;
    static uint64_t g_masterHwSeq = 0;

    // Monotonically increasing serial. Bumped (and waiters woken) on every
    // mutation of mixer state. Clients use it to detect changes without
    // re-reading the whole snapshot. The address of the serial doubles as the
    // wait-object passed to BlockOnObject / WakeObjectWaiters.
    static volatile uint64_t g_serial = 0;

    // Caller must hold g_lock.
    static void BumpSerialLocked() {
        g_serial++;
    }

    // Scratch mix buffer (int32 stereo, to avoid clipping during accumulation).
    // Sized for MAX_PUMP_FRAMES * 2 channels.
    static int32_t g_mixScratch[MAX_PUMP_FRAMES * 2];
    static int16_t g_outScratch[MAX_PUMP_FRAMES * 2];

    // =========================================================================
    // Helpers
    // =========================================================================

    static inline int16_t SatI16(int32_t v) {
        if (v >  32767) return 32767;
        if (v < -32768) return -32768;
        return (int16_t)v;
    }

    static int16_t* AllocRing() {
        void* p = Memory::g_pfa->ReallocConsecutive(nullptr, INPUT_RING_PAGES);
        if (!p) return nullptr;
        memset(p, 0, INPUT_RING_PAGES * 0x1000);
        return (int16_t*)p;
    }

    static void FreeRing(int16_t* ring) {
        if (!ring) return;
        Memory::g_pfa->ReallocConsecutive(ring, 0);
    }

    static bool EnsureHdaOpen() {
        if (g_hdaOpened) return true;
        if (!IntelHda::IsInitialized()) return false;
        g_hdaHandle = IntelHda::Open(MIX_RATE, MIX_CHANNELS, MIX_BITS);
        if (g_hdaHandle < 0) return false;
        // Master volume is applied in software during mixdown. Keep the codec
        // amp at unity and use it only as an immediate hard-mute gate.
        IntelHda::Control(g_hdaHandle, IntelHda::AUDIO_CTL_SET_VOLUME, g_masterMute ? 0 : 100);
        g_hdaOpened = true;
        return true;
    }

    static void SyncHdaMasterMute() {
        for (;;) {
            g_lock.Acquire();
            bool hdaOpen = g_hdaOpened;
            int handle = g_hdaHandle;
            bool muted = g_masterMute;
            uint64_t seq = g_masterHwSeq;
            g_lock.Release();

            if (!hdaOpen) return;

            IntelHda::Control(handle, IntelHda::AUDIO_CTL_SET_VOLUME,
                              muted ? 0 : 100);

            g_lock.Acquire();
            bool done = (seq == g_masterHwSeq);
            g_lock.Release();
            if (done) return;
        }
    }

    // Convert one chunk of raw input bytes from a stream into int16 stereo at
    // the stream's native rate, then push into its ring. Caller holds g_lock.
    // Returns the number of *input bytes* successfully ingested.
    static uint32_t ConvertAndPush(VirtualStream& s, const uint8_t* data, uint32_t size) {
        // Only 16-bit input is supported. Caller validates this in Open().
        const uint32_t inFrameBytes = (uint32_t)s.inputChannels * 2;
        uint32_t inFrames = size / inFrameBytes;
        if (inFrames == 0) return 0;

        // Clamp to whatever space the ring has left without growing past the
        // configured capacity. Excess input is dropped — callers know to write
        // again later.
        uint32_t inFlight = s.writeFrame - s.readFrame;
        uint32_t free = (inFlight >= INPUT_RING_FRAMES) ? 0 : (INPUT_RING_FRAMES - inFlight);
        if (inFrames > free) inFrames = free;
        if (inFrames == 0) return 0;

        const int16_t* src = (const int16_t*)data;
        for (uint32_t i = 0; i < inFrames; i++) {
            int16_t l, r;
            if (s.inputChannels == 1) {
                l = r = src[i];
            } else {
                l = src[i * 2 + 0];
                r = src[i * 2 + 1];
            }
            uint32_t idx = (s.writeFrame + i) % INPUT_RING_FRAMES;
            s.ring[idx * 2 + 0] = l;
            s.ring[idx * 2 + 1] = r;
        }
        s.writeFrame += inFrames;
        return inFrames * inFrameBytes;
    }

    // =========================================================================
    // Mixer pump
    // =========================================================================
    //
    // Compute the number of frames we can safely write to the HDA DMA buffer
    // (free space in the ring, minus a small guard), then cap that to the
    // amount of stream data actually available. Muted and zero-volume streams
    // still advance their read cursors so they cannot build up stale audio.
    // Finally hand the result to IntelHda::Write().

    static uint32_t ProducibleFrames(const VirtualStream& s, uint32_t maxFrames) {
        if (!s.active || s.paused) return 0;

        uint32_t available = s.writeFrame - s.readFrame;
        if (available < 2) return 0;

        uint64_t pos = s.posQ32;
        uint64_t step = s.stepQ32;
        uint32_t frames = 0;
        while (frames < maxFrames) {
            uint64_t intPart = pos >> 32;
            if (intPart + 1 >= available) break;
            frames++;
            pos += step;
        }
        return frames;
    }

    static void Pump() {
        if (!g_hdaOpened) return;

        // Produce only as many frames as the HDA DMA ring has room for right
        // now. Producing more would mean the surplus is silently dropped by
        // IntelHda::Write while the per-stream resamplers had already
        // advanced — that's what scrambles speech into a sequence of
        // unrelated chunks. Clamp to MAX_PUMP_FRAMES so the scratch buffers
        // are bounded.
        uint32_t freeBytes = IntelHda::GetWriteSpace(g_hdaHandle);
        uint32_t frames = freeBytes / 4;
        if (frames == 0) return;
        if (frames > MAX_PUMP_FRAMES) frames = MAX_PUMP_FRAMES;

        uint32_t streamFrames = 0;
        bool hasUnpaused = false;
        bool hasAudible = false;
        for (int i = 0; i < MAX_STREAMS; i++) {
            VirtualStream& s = g_streams[i];
            if (!s.active || s.paused) continue;
            hasUnpaused = true;
            if (!s.muted && s.volume > 0) hasAudible = true;
            uint32_t n = ProducibleFrames(s, frames);
            if (n > streamFrames) streamFrames = n;
        }

        // Write silence intentionally when there is no stream data to mix:
        // either the mixer is idle, every stream is paused/muted/silent, or an
        // audible producer has genuinely under-run. When data exists, write no
        // more than the longest available stream window so we don't append a
        // large zero tail after a small real-time write.
        bool silenceOnly = (streamFrames == 0);
        if (!silenceOnly) {
            frames = streamFrames;
        } else if (g_activeCount == 0 || !hasAudible || !hasUnpaused) {
            memset(g_outScratch, 0, frames * 2 * sizeof(int16_t));
            IntelHda::Write(g_hdaHandle, (const uint8_t*)g_outScratch, frames * 4);
            return;
        }

        // Zero the scratch accumulator.
        memset(g_mixScratch, 0, frames * 2 * sizeof(int32_t));

        for (int i = 0; i < MAX_STREAMS; i++) {
            VirtualStream& s = g_streams[i];
            if (!s.active || s.paused) continue;

            // posQ32 is the fractional read cursor; integer part is the index
            // of the *next* input frame to consume.
            uint64_t pos = s.posQ32;
            uint64_t step = s.stepQ32;
            uint32_t available = s.writeFrame - s.readFrame;
            if (available < 2) continue;

            // Per-stream gain in Q15 (post-master mix).
            int32_t gain = (int32_t)s.volume;  // 0..100
            bool emit = !s.muted && gain > 0;

            for (uint32_t f = 0; f < frames; f++) {
                uint64_t intPart = pos >> 32;
                if (intPart + 1 >= available) {
                    break;  // would read past end of ring; stop early
                }
                uint32_t i0 = (s.readFrame + (uint32_t)intPart) % INPUT_RING_FRAMES;
                uint32_t i1 = (s.readFrame + (uint32_t)intPart + 1) % INPUT_RING_FRAMES;

                int32_t l0 = s.ring[i0 * 2 + 0];
                int32_t r0 = s.ring[i0 * 2 + 1];
                int32_t l1 = s.ring[i1 * 2 + 0];
                int32_t r1 = s.ring[i1 * 2 + 1];

                // 16-bit fraction is the high half of the Q32 fractional part.
                int32_t frac = (int32_t)((pos & 0xFFFFFFFFull) >> 16);  // 0..65535
                int32_t inv  = 65536 - frac;
                int32_t l = (l0 * inv + l1 * frac) >> 16;
                int32_t r = (r0 * inv + r1 * frac) >> 16;

                if (emit) {
                    // Apply per-stream volume.
                    l = (l * gain) / 100;
                    r = (r * gain) / 100;

                    g_mixScratch[f * 2 + 0] += l;
                    g_mixScratch[f * 2 + 1] += r;
                }

                pos += step;
            }

            // Advance the ring read cursor by the integer part of `pos`, and
            // keep only the leftover fractional bit for next pump.
            uint32_t consumedFrames = (uint32_t)(pos >> 32);
            if (consumedFrames > available - 1) consumedFrames = available - 1;
            s.readFrame += consumedFrames;
            s.posQ32 = pos - ((uint64_t)consumedFrames << 32);
            s.consumedInputBytes += (uint64_t)consumedFrames *
                                    (uint64_t)((s.inputChannels == 1) ? 2 : 4);
        }

        // Saturate, apply master volume + mute, and emit s16 stereo.
        int32_t masterGain = g_masterMute ? 0 : g_masterVolume;  // 0..100
        for (uint32_t f = 0; f < frames; f++) {
            int32_t l = (g_mixScratch[f * 2 + 0] * masterGain) / 100;
            int32_t r = (g_mixScratch[f * 2 + 1] * masterGain) / 100;
            g_outScratch[f * 2 + 0] = SatI16(l);
            g_outScratch[f * 2 + 1] = SatI16(r);
        }

        // Hand off to HDA. IntelHda::Write returns the number of bytes
        // actually accepted (limited by free space in the DMA ring).
        IntelHda::Write(g_hdaHandle, (const uint8_t*)g_outScratch, frames * 4);
    }

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

    int Open(uint32_t sampleRate, uint8_t channels, uint8_t bitsPerSample,
             int ownerPid, const char* ownerName) {
        if (channels < 1 || channels > 2) return -1;
        if (bitsPerSample != 16) return -1;
        if (sampleRate < 8000 || sampleRate > 192000) return -1;

        // Allocate the input ring before taking g_lock. Walking the PFA free
        // list for 16 contiguous pages and zeroing 64 KiB is long enough that
        // doing it under the lock (with IRQs disabled) lets HDA BCIS IRQs
        // spin on g_lock on other CPUs — the same hazard that wedged Close().
        int16_t* ring = AllocRing();
        if (!ring) return -1;

        g_lock.Acquire();

        if (!EnsureHdaOpen()) {
            g_lock.Release();
            FreeRing(ring);
            return -1;
        }

        int slot = -1;
        for (int i = 0; i < MAX_STREAMS; i++) {
            if (!g_streams[i].active) { slot = i; break; }
        }
        if (slot < 0) {
            g_lock.Release();
            FreeRing(ring);
            return -1;
        }

        VirtualStream& s = g_streams[slot];
        s.active = true;
        s.ownerPid = ownerPid;
        int n = 0;
        if (ownerName) {
            for (; n < 63 && ownerName[n]; n++) s.name[n] = ownerName[n];
        }
        s.name[n] = '\0';
        s.inputRate = sampleRate;
        s.inputChannels = channels;
        s.inputBits = bitsPerSample;
        s.volume = 100;
        s.muted = false;
        s.paused = false;
        s.ring = ring;
        s.writeFrame = 0;
        s.readFrame = 0;
        s.consumedInputBytes = 0;
        s.posQ32 = 0;
        // step = inputRate / MIX_RATE in Q32.
        s.stepQ32 = ((uint64_t)sampleRate << 32) / MIX_RATE;
        g_activeCount++;
        BumpSerialLocked();

        g_lock.Release();
        Sched::WakeObjectWaiters((void*)&g_serial);
        return slot;
    }

    void Close(int handle) {
        if (handle < 0 || handle >= MAX_STREAMS) return;
        g_lock.Acquire();
        VirtualStream& s = g_streams[handle];
        bool changed = false;
        // Pull the ring pointer out of the slot before freeing it. Once
        // s.active=false and s.ring=nullptr are visible under the lock no
        // other path can reach this pointer, so the actual PFA free can run
        // with the lock released — keeping the HDA buffer-complete IRQ from
        // spinning on g_lock while ReallocConsecutive walks the free list.
        int16_t* ringToFree = nullptr;
        bool lastStream = false;
        if (s.active) {
            ringToFree = s.ring;
            s.ring = nullptr;
            s.active = false;
            s.ownerPid = 0;
            s.name[0] = '\0';
            if (g_activeCount > 0) g_activeCount--;
            lastStream = (g_activeCount == 0);
            BumpSerialLocked();
            changed = true;
        }
        // Quiesce the hardware when the final virtual stream goes away so
        // the HDA DMA ring stops looping its last 32 KiB of samples. The
        // next Open() reopens the HDA stream via EnsureHdaOpen().
        bool closeHda = lastStream && g_hdaOpened;
        int hdaHandle = g_hdaHandle;
        if (closeHda) {
            g_hdaOpened = false;
            g_hdaHandle = -1;
        }
        g_lock.Release();
        if (closeHda) IntelHda::Close(hdaHandle);
        if (ringToFree) FreeRing(ringToFree);
        if (changed) Sched::WakeObjectWaiters((void*)&g_serial);
    }

    int Write(int handle, const uint8_t* data, uint32_t size) {
        if (handle < 0 || handle >= MAX_STREAMS || !data || size == 0) return -1;
        g_lock.Acquire();
        VirtualStream& s = g_streams[handle];
        if (!s.active) { g_lock.Release(); return -1; }

        uint32_t written = ConvertAndPush(s, data, size);

        // Run a pump cycle so the HDA buffer stays fed.
        Pump();

        g_lock.Release();
        return (int)written;
    }

    int Control(int handle, int cmd, int value) {
        // Master-scope commands ignore the handle.
        switch (cmd) {
            case Montauk::AUDIO_CTL_SET_MASTER_VOLUME:
                SetMasterVolume(value);
                return 0;
            case Montauk::AUDIO_CTL_GET_MASTER_VOLUME:
                return GetMasterVolume();
            case Montauk::AUDIO_CTL_SET_MASTER_MUTE:
                SetMasterMute(value != 0);
                return 0;
            case Montauk::AUDIO_CTL_GET_MASTER_MUTE:
                return GetMasterMute() ? 1 : 0;
        }

        if (handle < 0 || handle >= MAX_STREAMS) return -1;

        g_lock.Acquire();
        VirtualStream& s = g_streams[handle];
        if (!s.active) { g_lock.Release(); return -1; }

        int rv = -1;
        bool changed = false;
        switch (cmd) {
            case Montauk::AUDIO_CTL_SET_VOLUME: {
                int v = value; if (v < 0) v = 0; if (v > 100) v = 100;
                if (s.volume != (uint8_t)v) { s.volume = (uint8_t)v; changed = true; }
                rv = 0;
                break;
            }
            case Montauk::AUDIO_CTL_GET_VOLUME:
                rv = s.volume;
                break;
            case Montauk::AUDIO_CTL_GET_POS:
                rv = (int)(s.consumedInputBytes & 0x7FFFFFFF);
                break;
            case Montauk::AUDIO_CTL_PAUSE:
                if (s.paused != (value != 0)) { s.paused = (value != 0); changed = true; }
                rv = 0;
                break;
            case Montauk::AUDIO_CTL_SET_MUTE:
                if (s.muted != (value != 0)) { s.muted = (value != 0); changed = true; }
                rv = 0;
                break;
            case Montauk::AUDIO_CTL_GET_MUTE:
                rv = s.muted ? 1 : 0;
                break;
            default:
                rv = -1;
                break;
        }
        if (changed) BumpSerialLocked();
        g_lock.Release();
        if (changed) Sched::WakeObjectWaiters((void*)&g_serial);
        return rv;
    }

    int List(Montauk::AudioStreamInfo* buf, int maxCount) {
        if (!buf || maxCount <= 0) return 0;
        g_lock.Acquire();
        int count = 0;
        for (int i = 0; i < MAX_STREAMS && count < maxCount; i++) {
            VirtualStream& s = g_streams[i];
            if (!s.active) continue;
            buf[count].handle = i;
            buf[count].ownerPid = s.ownerPid;
            int j = 0;
            for (; j < 63 && s.name[j]; j++) buf[count].name[j] = s.name[j];
            buf[count].name[j] = '\0';
            buf[count].sampleRate = s.inputRate;
            buf[count].channels = s.inputChannels;
            buf[count].bitsPerSample = s.inputBits;
            buf[count].volume = s.volume;
            buf[count].muted = s.muted ? 1 : 0;
            buf[count].paused = s.paused ? 1 : 0;
            for (int k = 0; k < 7; k++) buf[count]._pad[k] = 0;
            count++;
        }
        g_lock.Release();
        return count;
    }

    void CleanupProcess(int pid) {
        int16_t* ringsToFree[MAX_STREAMS] = {};
        int ringCount = 0;
        g_lock.Acquire();
        bool changed = false;
        for (int i = 0; i < MAX_STREAMS; i++) {
            VirtualStream& s = g_streams[i];
            if (s.active && s.ownerPid == pid) {
                ringsToFree[ringCount++] = s.ring;
                s.ring = nullptr;
                s.active = false;
                s.ownerPid = 0;
                s.name[0] = '\0';
                if (g_activeCount > 0) g_activeCount--;
                changed = true;
            }
        }
        bool closeHda = changed && (g_activeCount == 0) && g_hdaOpened;
        int hdaHandle = g_hdaHandle;
        if (closeHda) {
            g_hdaOpened = false;
            g_hdaHandle = -1;
        }
        if (changed) BumpSerialLocked();
        g_lock.Release();
        if (closeHda) IntelHda::Close(hdaHandle);
        for (int i = 0; i < ringCount; i++) FreeRing(ringsToFree[i]);
        if (changed) Sched::WakeObjectWaiters((void*)&g_serial);
    }

    void SetMasterVolume(int percent) {
        if (percent < 0) percent = 0;
        if (percent > 100) percent = 100;

        g_lock.Acquire();
        bool changed = (g_masterVolume != percent);
        g_masterVolume = percent;
        if (changed) BumpSerialLocked();
        g_lock.Release();

        if (changed) Sched::WakeObjectWaiters((void*)&g_serial);
    }

    int GetMasterVolume() {
        return g_masterVolume;
    }

    void SetMasterMute(bool muted) {
        g_lock.Acquire();
        bool changed = (g_masterMute != muted);
        g_masterMute = muted;
        if (changed) {
            g_masterHwSeq++;
            BumpSerialLocked();
        }
        g_lock.Release();

        if (changed) SyncHdaMasterMute();
        if (changed) Sched::WakeObjectWaiters((void*)&g_serial);
    }

    bool GetMasterMute() {
        return g_masterMute;
    }

    void OnHdaBufferComplete() {
        // Called from the HDA buffer-completion interrupt. Mixer state and
        // HDA register access are both serialized through g_lock (which
        // disables interrupts on acquire), so this is safe to call from IRQ.
        g_lock.Acquire();
        Pump();
        g_lock.Release();
    }

    uint64_t GetSerial() {
        return g_serial;
    }

    // BlockOnObjectIf callback: returns true (i.e. "do block") only if the
    // current serial still matches the caller's snapshot. Avoids the classic
    // wakeup-before-block race: state could change between the caller's first
    // read of g_serial and the scheduler dropping the process to Blocked.
    struct WaitCtx { uint64_t expected; };
    static bool WaitShouldBlock(void* ctx) {
        return g_serial == ((WaitCtx*)ctx)->expected;
    }

    uint64_t Wait(uint64_t prevSerial, uint64_t timeoutMs) {
        // Fast path: state already moved on, no need to enter the scheduler.
        if (g_serial != prevSerial) return g_serial;
        if (timeoutMs == 0) return g_serial;

        WaitCtx ctx{prevSerial};
        Sched::BlockOnObjectIf((void*)&g_serial, timeoutMs,
                               WaitShouldBlock, &ctx);
        return g_serial;
    }

};
