using System.Runtime.InteropServices;

// Single-process WASAPI loopback capture via AUDIOCLIENT_ACTIVATION_PARAMS
// (Windows 10 2004+ / Build 19041+).
//
// Threading: ALL WASAPI init runs on the capture thread (MTA). If called from the
// WinForms UI thread (STA), ActivateAudioInterfaceAsync fires ActivateCompleted on
// an MTA pool thread; COM marshals that back to the STA pump — but the STA thread is
// blocked on CompletionEvent.Wait → deadlock. MTA capture thread avoids this.
//
// COM QI policy: the COM objects returned by the process-loopback activation path
// reject QueryInterface for their own IIDs under .NET's RCW mechanism. Every call
// to IAudioClient and IAudioCaptureClient is therefore dispatched via raw vtable
// pointer arithmetic, bypassing .NET COM interop entirely.
namespace VoiceCat.App.Audio;

public sealed class ProcessLoopbackCapture : IDisposable
{
    public enum Mode { Include, Exclude }

    // Fired on the capture thread every 20 ms (960 samples/channel @ 48 kHz, interleaved s16).
    public event Action<short[], int /*samplesPerChannel*/, int /*channels*/>? PcmFrameReady;

    private const int SampleRate   = 48000;
    private const int FrameSamples = 960;  // 20 ms
    private const string LoopbackDevicePath = "VAD\\Process_Loopback";

    private readonly int  _pid;
    private readonly Mode _mode;

    // Raw COM pointers — managed via explicit AddRef/Release, no RCW wrapping.
    private IntPtr _audioClientPtr;   // IAudioClient*
    private IntPtr _captureClientPtr; // IAudioCaptureClient*

    private AutoResetEvent? _bufferEvent;
    private Thread?         _captureThread;
    private volatile bool   _running;
    private int             _channels;

    // Accumulator: assembles driver-callback-sized fragments into FrameSamples chunks.
    private short[] _accumBuf = [];
    private int     _accumCount;

    // Init-done signal: Set() by the capture thread after ActivateClient() completes.
    private readonly ManualResetEventSlim _initDone = new(false);
    private bool _initOk;

    public ProcessLoopbackCapture(int pid, Mode mode)
    {
        _pid  = pid;
        _mode = mode;
    }

    /// <summary>Starts capture. Blocks until WASAPI activation completes (typically &lt;100 ms).
    /// Returns false if the process cannot be captured.</summary>
    public bool Start()
    {
        if (_running) return false;
        _running       = true;
        _captureThread = new Thread(CaptureThreadProc)
        {
            IsBackground = true,
            Name         = $"ProcLoopback:{_pid}",
        };
        _captureThread.Start();

        bool ok = _initDone.Wait(5000) && _initOk;
        if (!ok) _running = false;
        return ok;
    }

    public void Stop()
    {
        _running = false;
        _bufferEvent?.Set();
        _captureThread?.Join(500);
        if (_audioClientPtr != IntPtr.Zero) AC_Stop(_audioClientPtr);
    }

    public void Dispose()
    {
        Stop();
        ComRelease(ref _captureClientPtr);
        ComRelease(ref _audioClientPtr);
        _bufferEvent?.Dispose();
        _initDone.Dispose();
    }

    // ── Capture thread (MTA) ──────────────────────────────────────────────────

    private void CaptureThreadProc()
    {
        _initOk = ActivateAndStart();
        _initDone.Set();
        if (!_initOk) return;
        CaptureLoop();
    }

    private bool ActivateAndStart()
    {
        if (!ActivateClient()) return false;

        _bufferEvent = new AutoResetEvent(false);
        if (AC_SetEventHandle(_audioClientPtr, _bufferEvent.SafeWaitHandle.DangerousGetHandle()) < 0)
            return false;

        return AC_Start(_audioClientPtr) >= 0;
    }

    // ── Activation ───────────────────────────────────────────────────────────

    private unsafe bool ActivateClient()
    {
        var activationParams = new AudioClientActivationParams
        {
            ActivationType      = 1,   // AUDCLNT_ACTIVATIONTYPE_PROCESS_LOOPBACK
            TargetProcessId     = (uint)_pid,
            ProcessLoopbackMode = _mode == Mode.Include ? 0u : 1u,
        };

        var handler = new ActivationCompletionHandler();
        var audioClientIid = new Guid("1CB9AD4C-DBFA-4c32-B178-C2F568A703B2");

        IntPtr opPtr;
        {
            AudioClientActivationParams* pParams = &activationParams;
            // PROPVARIANT (VT_BLOB) x64: vt(2)+res(6)+cbSize(4)+pad(4)+pBlobData(8) = 24 B.
            var pv = stackalloc byte[24];
            *(ushort*)(pv + 0)  = 65;
            *(uint*)  (pv + 8)  = (uint)sizeof(AudioClientActivationParams);
            *(nint*)  (pv + 16) = (nint)pParams;

            int hr = ActivateAudioInterfaceAsync(
                LoopbackDevicePath, ref audioClientIid,
                (IntPtr)pv, handler, out opPtr);
            if (hr < 0) return false;
        }

        if (!handler.CompletionEvent.Wait(3000))
        {
            ComRelease(ref opPtr);
            return false;
        }

        // Call GetActivateResult via vtable (slot 3) — avoids QI on the async-op object.
        if (!Vtable_GetActivateResult(handler.OperationPtr, out int activateHr, out IntPtr activatedPtr))
        {
            handler.ReleaseOp();
            ComRelease(ref opPtr);
            return false;
        }
        handler.ReleaseOp();
        ComRelease(ref opPtr);

        if (activateHr < 0 || activatedPtr == IntPtr.Zero) return false;

        // Store the raw IAudioClient* — do NOT create an RCW; use vtable dispatch instead.
        _audioClientPtr = activatedPtr;
        // activatedPtr already has ref count from GetActivateResult; don't double-release.

        return InitializeStream();
    }

    // IActivateAudioInterfaceAsyncOperation vtable slot 3: GetActivateResult(HRESULT*, IUnknown**)
    private static unsafe bool Vtable_GetActivateResult(IntPtr op,
        out int activateHr, out IntPtr activatedPtr)
    {
        activateHr   = unchecked((int)0x80004005);
        activatedPtr = IntPtr.Zero;
        if (op == IntPtr.Zero) return false;
        void** vtable = *(void***)op.ToPointer();
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, int*, IntPtr*, int>)vtable[3];
        fixed (int*    pHr  = &activateHr)
        fixed (IntPtr* pPtr = &activatedPtr)
            return fn(op, pHr, pPtr) >= 0;
    }

    private unsafe bool InitializeStream()
    {
        // Try s16 stereo first; fall back to s16 mono.
        foreach (int ch in new[] { 2, 1 })
        {
            var fmt = new WaveFormatEx
            {
                wFormatTag      = 1,           // WAVE_FORMAT_PCM
                nChannels       = (ushort)ch,
                nSamplesPerSec  = SampleRate,
                wBitsPerSample  = 16,
                nBlockAlign     = (ushort)(ch * 2),
                nAvgBytesPerSec = (uint)(SampleRate * ch * 2),
                cbSize          = 0,
            };

            // Process-loopback requires LOOPBACK (deliver rendered audio) + EVENTCALLBACK +
            // AUTOCONVERTPCM (resample the app's native format to our requested s16/48k). Without
            // LOOPBACK every buffer comes back AUDCLNT_BUFFERFLAGS_SILENT; without AUTOCONVERTPCM
            // the requested format is rejected. Matches the MS ApplicationLoopback sample.
            //   AUDCLNT_SHAREMODE_SHARED      = 0
            //   AUDCLNT_STREAMFLAGS_LOOPBACK  = 0x00020000
            //   AUDCLNT_STREAMFLAGS_EVENTCALLBACK = 0x00040000
            //   AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM = 0x80000000
            const uint streamFlags = 0x00020000u | 0x00040000u | 0x80000000u;
            int hr = AC_Initialize(_audioClientPtr, 0, streamFlags,
                                   2_000_000 /*200ms hns*/, 0, &fmt, null);
            if (hr >= 0)
            {
                _channels   = ch;
                _accumBuf   = new short[FrameSamples * ch];
                _accumCount = 0;
                break;
            }
            if (ch == 1) return false;
        }

        var captureIid = new Guid("C8ADBD64-E71E-48a0-A4DE-185C395CD317");
        int getHr = AC_GetService(_audioClientPtr, ref captureIid, out _captureClientPtr);
        return getHr >= 0 && _captureClientPtr != IntPtr.Zero;
    }

    // ── Capture loop ─────────────────────────────────────────────────────────

    private void CaptureLoop()
    {
        while (_running)
        {
            _bufferEvent!.WaitOne(100);
            if (!_running) break;

            while (_running)
            {
                int hr = CC_GetNextPacketSize(_captureClientPtr, out uint packetSize);
                if (hr < 0 || packetSize == 0) break;

                hr = CC_GetBuffer(_captureClientPtr, out IntPtr dataPtr, out uint framesAvailable,
                    out uint flags);
                if (hr < 0) break;

                bool silent = (flags & 2) != 0; // AUDCLNT_BUFFERFLAGS_SILENT
                if (framesAvailable > 0)
                {
                    if (silent) AccumulateSilence((int)framesAvailable);
                    else        AccumulatePcm(dataPtr, (int)framesAvailable);
                }

                CC_ReleaseBuffer(_captureClientPtr, framesAvailable);
            }
        }
    }

    private unsafe void AccumulatePcm(IntPtr data, int frames)
    {
        var src   = (short*)data.ToPointer();
        int total = frames * _channels;
        int idx   = 0;
        while (idx < total)
        {
            int space = _accumBuf.Length - _accumCount;
            int copy  = Math.Min(total - idx, space);
            fixed (short* dst = _accumBuf)
                Buffer.MemoryCopy(src + idx, dst + _accumCount, copy * 2L, copy * 2L);
            _accumCount += copy;
            idx         += copy;
            if (_accumCount == _accumBuf.Length)
                FlushFrame();
        }
    }

    private void AccumulateSilence(int frames)
    {
        int total = frames * _channels;
        int idx   = 0;
        while (idx < total)
        {
            int space = _accumBuf.Length - _accumCount;
            int fill  = Math.Min(total - idx, space);
            Array.Clear(_accumBuf, _accumCount, fill);
            _accumCount += fill;
            idx         += fill;
            if (_accumCount == _accumBuf.Length)
                FlushFrame();
        }
    }

    private void FlushFrame()
    {
        var copy = new short[_accumBuf.Length];
        _accumBuf.AsSpan().CopyTo(copy);
        PcmFrameReady?.Invoke(copy, FrameSamples, _channels);
        _accumCount = 0;
    }

    // ── IAudioClient vtable helpers (raw dispatch, no RCW / no QI) ───────────
    //
    // Vtable layout (IUnknown base: 0=QI 1=AddRef 2=Release; then IAudioClient methods):
    //   3=Initialize  4=GetBufferSize  5=GetStreamLatency  6=GetCurrentPadding
    //   7=IsFormatSupported  8=GetMixFormat  9=GetDevicePeriod
    //   10=Start  11=Stop  12=Reset  13=SetEventHandle  14=GetService

    private static unsafe int AC_Initialize(IntPtr ac, int shareMode, uint streamFlags,
        long hnsBufferDuration, long hnsPeriodicity, WaveFormatEx* pFormat, Guid* pSession)
    {
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, int, uint, long, long, WaveFormatEx*, Guid*, int>)
            (*(void***)ac)[3];
        return fn(ac, shareMode, streamFlags, hnsBufferDuration, hnsPeriodicity, pFormat, pSession);
    }

    private static unsafe int AC_Start(IntPtr ac)
    {
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, int>)(*(void***)ac)[10];
        return fn(ac);
    }

    private static unsafe int AC_Stop(IntPtr ac)
    {
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, int>)(*(void***)ac)[11];
        return fn(ac);
    }

    private static unsafe int AC_SetEventHandle(IntPtr ac, IntPtr eventHandle)
    {
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, IntPtr, int>)(*(void***)ac)[13];
        return fn(ac, eventHandle);
    }

    private static unsafe int AC_GetService(IntPtr ac, ref Guid riid, out IntPtr ppv)
    {
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, Guid*, IntPtr*, int>)(*(void***)ac)[14];
        fixed (Guid* pIid = &riid)
        fixed (IntPtr* pPpv = &ppv)
            return fn(ac, pIid, pPpv);
    }

    // ── IAudioCaptureClient vtable helpers ────────────────────────────────────
    //
    // Vtable (IUnknown: 0-2; then): 3=GetBuffer  4=ReleaseBuffer  5=GetNextPacketSize

    private static unsafe int CC_GetBuffer(IntPtr cc, out IntPtr ppData,
        out uint pNumFrames, out uint pdwFlags)
    {
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, IntPtr*, uint*, uint*, ulong*, ulong*, int>)
            (*(void***)cc)[3];
        ulong devPos = 0, qpcPos = 0;
        fixed (IntPtr* p0 = &ppData)
        fixed (uint*   p1 = &pNumFrames)
        fixed (uint*   p2 = &pdwFlags)
            return fn(cc, p0, p1, p2, &devPos, &qpcPos);
    }

    private static unsafe int CC_ReleaseBuffer(IntPtr cc, uint numFrames)
    {
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, uint, int>)(*(void***)cc)[4];
        return fn(cc, numFrames);
    }

    private static unsafe int CC_GetNextPacketSize(IntPtr cc, out uint pNumFrames)
    {
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, uint*, int>)(*(void***)cc)[5];
        fixed (uint* p = &pNumFrames)
            return fn(cc, p);
    }

    // ── COM utilities ─────────────────────────────────────────────────────────

    private static unsafe void ComRelease(ref IntPtr ptr)
    {
        if (ptr == IntPtr.Zero) return;
        var fn = (delegate* unmanaged[Stdcall]<IntPtr, uint>)(*(void***)ptr)[2]; // IUnknown::Release
        fn(ptr);
        ptr = IntPtr.Zero;
    }

    // ── P/Invoke & structs ────────────────────────────────────────────────────

    [DllImport("Mmdevapi.dll", CharSet = CharSet.Unicode)]
    private static extern int ActivateAudioInterfaceAsync(
        string deviceInterfacePath,
        ref Guid riid,
        IntPtr activationParams,
        [MarshalAs(UnmanagedType.Interface)] IActivateAudioInterfaceCompletionHandler completionHandler,
        out IntPtr activationOperation);

    [StructLayout(LayoutKind.Sequential)]
    private struct AudioClientActivationParams
    {
        public int  ActivationType;       // AUDCLNT_ACTIVATIONTYPE_PROCESS_LOOPBACK = 1
        public uint TargetProcessId;
        public uint ProcessLoopbackMode;  // INCLUDE=0, EXCLUDE=1
    }

    [StructLayout(LayoutKind.Sequential, Pack = 2)]
    private struct WaveFormatEx
    {
        public ushort wFormatTag;
        public ushort nChannels;
        public uint   nSamplesPerSec;
        public uint   nAvgBytesPerSec;
        public ushort nBlockAlign;
        public ushort wBitsPerSample;
        public ushort cbSize;
    }

    // ── Completion handler CCW ────────────────────────────────────────────────
    //
    // Only this object still uses .NET COM interop (as a CCW). The activateOperation
    // parameter is IntPtr to avoid QI on the incoming async-op pointer.

    [ComVisible(true), ClassInterface(ClassInterfaceType.None)]
    private sealed class ActivationCompletionHandler : IActivateAudioInterfaceCompletionHandler
    {
        public readonly ManualResetEventSlim CompletionEvent = new(false);
        public IntPtr OperationPtr { get; private set; }

        public void ActivateCompleted(IntPtr activateOperation)
        {
            OperationPtr = activateOperation;
            if (OperationPtr != IntPtr.Zero) Marshal.AddRef(OperationPtr);
            CompletionEvent.Set();
        }

        public void ReleaseOp()
        {
            if (OperationPtr == IntPtr.Zero) return;
            Marshal.Release(OperationPtr);
            OperationPtr = IntPtr.Zero;
        }
    }

    [ComImport, Guid("41D949AB-9862-444A-80F6-C261334DA5EB"),
     InterfaceType(ComInterfaceType.InterfaceIsIUnknown)]
    private interface IActivateAudioInterfaceCompletionHandler
    {
        void ActivateCompleted(IntPtr activateOperation);
    }
}