#pragma once

#include <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstring>
#include <mutex>
#include <numeric>
#include <stdexcept>
#include <string>
#include <vector>
#include <DeckLinkAPI.h>
#include "Signal.hpp"

namespace dmf {

class DeckLinkReceiver {
public:
    struct VideoInfo {
        int width   = 0;
        int height  = 0;
        int fps_num = 0;
        int fps_den = 0;
    };

    struct AudioInfo {
        int sample_rate = 48000;  // DeckLink always delivers 48 kHz
        int channels    = 0;
    };

    struct DeviceInfo {
        uint32_t    index;
        std::string name;
    };

    std::vector<DeviceInfo> devices;
    VideoInfo               video_info{};
    AudioInfo               audio_info{};
    bool                    has_audio = false;

    DeckLinkReceiver() { enumerate_devices(); }

    ~DeckLinkReceiver() {
        if (decklink_input) {
            decklink_input->StopStreams();
            decklink_input->DisableVideoInput();
            if (has_audio) decklink_input->DisableAudioInput();
            decklink_input->SetCallback(nullptr);
            decklink_input->Release();
        }
        delete input_callback;
        for (auto* d : raw_devices) d->Release();
        if (selected_device) selected_device->Release();
    }

    // audio_channels > 0 enables audio capture at 48 kHz / 32-bit int.
    void start_capture(uint32_t device_index, int audio_channels = 0) {
        if (device_index >= raw_devices.size())
            throw std::runtime_error("Device index out of range");

        selected_device = raw_devices[device_index];
        selected_device->AddRef();

        HRESULT r = selected_device->QueryInterface(IID_IDeckLinkInput, (void**)&decklink_input);
        if (r != S_OK) throw std::runtime_error("Could not obtain IDeckLinkInput");

        input_callback = new InputCallback(*this);
        r = decklink_input->SetCallback(input_callback);
        if (r != S_OK) throw std::runtime_error("Could not set input callback");

        r = decklink_input->EnableVideoInput(bmdModeNTSC, bmdFormat10BitYUV,
                                              bmdVideoInputEnableFormatDetection);
        if (r != S_OK) throw std::runtime_error("Could not enable video input");

        if (audio_channels > 0) {
            r = decklink_input->EnableAudioInput(bmdAudioSampleRate48kHz,
                                                  bmdAudioSampleType32bitInteger,
                                                  static_cast<uint32_t>(audio_channels));
            if (r != S_OK) throw std::runtime_error("Could not enable audio input");
            audio_info.channels = audio_channels;
            has_audio = true;
        }

        r = decklink_input->StartStreams();
        if (r != S_OK) throw std::runtime_error("Could not start streams");
    }

    bool wait_for_format(int timeout_ms = 5000) {
        std::unique_lock<std::mutex> lk(mutex);
        return format_cv.wait_for(lk, std::chrono::milliseconds(timeout_ms),
            [this] { return format_detected; });
    }

    // Blocks until a fresh frame arrives or g_running goes false.
    // Copies video into video_dst and (if audio_dst != nullptr) deinterleaved
    // float32 audio into audio_dst[channel * max_samples + sample].
    // Returns false on shutdown.
    bool wait_for_frame(uint8_t* video_dst, uint32_t dst_stride, int width, int height,
                        float* audio_dst = nullptr, int max_samples = 0, int* samples_written = nullptr) {
        std::unique_lock<std::mutex> lk(mutex);
        frame_cv.wait(lk, [this] {
            return frame_ready || !dmf::g_running.load(std::memory_order_relaxed);
        });
        if (!dmf::g_running.load(std::memory_order_relaxed)) return false;

        // Video copy
        const uint32_t src_stride = frame_row_bytes;
        const int      rows       = std::min(frame_height, height);
        const uint32_t copy_bytes = std::min(src_stride, dst_stride);
        std::memset(video_dst, 0, static_cast<size_t>(dst_stride) * static_cast<size_t>(height));
        const uint8_t* src = frame_buffer.data();
        uint8_t*       d   = video_dst;
        for (int y = 0; y < rows; ++y, src += src_stride, d += dst_stride)
            std::memcpy(d, src, copy_bytes);

        // Audio copy — planar float32: channel c starts at audio_dst + c * max_samples
        if (audio_dst && max_samples > 0 && samples_written) {
            const int n  = std::min(audio_samples, max_samples);
            const int ch = audio_info.channels;
            *samples_written = n;
            for (int c = 0; c < ch; ++c)
                std::memcpy(audio_dst + c * max_samples,
                            audio_buffer.data() + c * audio_samples,
                            static_cast<size_t>(n) * sizeof(float));
        } else if (samples_written) {
            *samples_written = 0;
        }

        frame_ready = false;
        return true;
    }

private:
    class InputCallback : public IDeckLinkInputCallback {
    public:
        explicit InputCallback(DeckLinkReceiver& owner) : owner(owner) {}

        HRESULT STDMETHODCALLTYPE VideoInputFormatChanged(
            BMDVideoInputFormatChangedEvents events,
            IDeckLinkDisplayMode*            mode,
            BMDDetectedVideoInputFormatFlags /*flags*/) override
        {
            if (!(events & bmdVideoInputDisplayModeChanged)) return S_OK;

            {
                std::lock_guard<std::mutex> lk(owner.mutex);
                owner.video_info.width  = static_cast<int>(mode->GetWidth());
                owner.video_info.height = static_cast<int>(mode->GetHeight());
                BMDTimeValue dur = 0; BMDTimeScale ts = 0;
                mode->GetFrameRate(&dur, &ts);
                if (dur > 0 && ts > 0) {
                    owner.video_info.fps_num = static_cast<int>(ts);
                    owner.video_info.fps_den = static_cast<int>(dur);
                    const int g = std::gcd(owner.video_info.fps_num, owner.video_info.fps_den);
                    if (g > 1) { owner.video_info.fps_num /= g; owner.video_info.fps_den /= g; }
                }
                owner.frame_ready     = false;
                owner.frame_buffer.clear();
                owner.format_detected = true;
            }
            owner.format_cv.notify_one();

            HRESULT r = owner.decklink_input->StopStreams();
            if (r == S_OK) {
                r = owner.decklink_input->EnableVideoInput(
                    mode->GetDisplayMode(), bmdFormat10BitYUV,
                    bmdVideoInputEnableFormatDetection);
                if (r == S_OK) owner.decklink_input->StartStreams();
            }
            return S_OK;
        }

        HRESULT STDMETHODCALLTYPE VideoInputFrameArrived(
            IDeckLinkVideoInputFrame*  video_frame,
            IDeckLinkAudioInputPacket* audio_packet) override
        {
            if (!video_frame) return S_OK;

            IDeckLinkVideoBuffer* buf = nullptr;
            if (video_frame->QueryInterface(IID_IDeckLinkVideoBuffer, (void**)&buf) != S_OK)
                return S_OK;
            buf->StartAccess(bmdBufferAccessRead);
            void* src = nullptr;
            buf->GetBytes(&src);

            if (src) {
                const uint32_t row_bytes = static_cast<uint32_t>(video_frame->GetRowBytes());
                const int      fw        = static_cast<int>(video_frame->GetWidth());
                const int      fh        = static_cast<int>(video_frame->GetHeight());
                const size_t   sz        = static_cast<size_t>(row_bytes) * static_cast<size_t>(fh);

                std::lock_guard<std::mutex> lk(owner.mutex);

                // Video
                if (owner.frame_buffer.size() < sz) owner.frame_buffer.resize(sz);
                std::memcpy(owner.frame_buffer.data(), src, sz);
                owner.frame_row_bytes = row_bytes;
                owner.frame_width     = fw;
                owner.frame_height    = fh;

                // Audio — deinterleave int32 → float32 planar under the same lock
                if (owner.has_audio && audio_packet) {
                    void* asrc = nullptr;
                    audio_packet->GetBytes(&asrc);
                    const long nb = audio_packet->GetSampleFrameCount();
                    const int  ch = owner.audio_info.channels;
                    if (asrc && nb > 0 && ch > 0) {
                        const auto*  in   = static_cast<const int32_t*>(asrc);
                        const size_t need = static_cast<size_t>(ch) * static_cast<size_t>(nb);
                        if (owner.audio_buffer.size() < need) owner.audio_buffer.resize(need);
                        for (long s = 0; s < nb; ++s)
                            for (int c = 0; c < ch; ++c)
                                owner.audio_buffer[static_cast<size_t>(c) * static_cast<size_t>(nb) + static_cast<size_t>(s)]
                                    = static_cast<float>(in[s * ch + c]) / 2147483648.0f;
                        owner.audio_samples = static_cast<int>(nb);
                    }
                }

                owner.frame_ready = true;
            }

            buf->EndAccess(bmdBufferAccessRead);
            buf->Release();
            if (src) owner.frame_cv.notify_one();
            return S_OK;
        }

        HRESULT STDMETHODCALLTYPE QueryInterface(REFIID, LPVOID*) override { return E_NOINTERFACE; }
        ULONG   STDMETHODCALLTYPE AddRef()  override { return ++ref_count; }
        ULONG   STDMETHODCALLTYPE Release() override { return --ref_count; }

    private:
        DeckLinkReceiver&    owner;
        std::atomic<int32_t> ref_count{1};
    };

    // DeckLink SDK objects
    std::vector<IDeckLink*>  raw_devices;
    IDeckLink*               selected_device = nullptr;
    IDeckLinkInput*          decklink_input  = nullptr;
    InputCallback*           input_callback  = nullptr;

    // Synchronisation — one mutex guards all shared state below
    std::mutex              mutex;
    std::condition_variable format_cv;   // signalled when format is detected
    std::condition_variable frame_cv;    // signalled when a frame arrives
    bool                    format_detected = false;

    // Frame state — written by InputCallback, read by wait_for_frame (both under mutex)
    std::vector<uint8_t>    frame_buffer;
    uint32_t                frame_row_bytes = 0;
    int                     frame_width     = 0;
    int                     frame_height    = 0;
    bool                    frame_ready     = false;

    // Audio state — planar float32: channel c at audio_buffer[c * audio_samples + s]
    std::vector<float>      audio_buffer;
    int                     audio_samples = 0;

    void enumerate_devices() {
        IDeckLinkIterator* it = CreateDeckLinkIteratorInstance();
        if (!it) throw std::runtime_error("DeckLink drivers not installed");

        IDeckLink* device = nullptr;
        uint32_t   index  = 0;
        while (it->Next(&device) == S_OK) {
            IDeckLinkInput* inp = nullptr;
            if (device->QueryInterface(IID_IDeckLinkInput, (void**)&inp) == S_OK) {
                inp->Release();
                const char* name = nullptr;
                device->GetDisplayName(&name);
                devices.push_back({index, name ? name : "?"});
                raw_devices.push_back(device);
                index++;
            } else {
                device->Release();
            }
        }
        it->Release();
        if (raw_devices.empty())
            throw std::runtime_error("No DeckLink input devices found");
    }
};

} // namespace dmf