#include "DeviceDiscovery.h"
#include "BoostLog.h"
#include "StringUtility.h"
#include <boost/scope/scope_exit.hpp>
#include <cstdlib>
#include <mfapi.h>
#include <mfcaptureengine.h>

struct DeviceDiscovery::Device {
    Device(IMFMediaSource *source);
    ~Device();
    IMFMediaSource *source = nullptr;
    IMFSourceReader *reader = nullptr;
};

template <class T>
void SafeRelease(T **ppT) {
    if (*ppT) {
        (*ppT)->Release();
        *ppT = NULL;
    }
}

DeviceDiscovery::DeviceDiscovery() {
}

static std::string deviceName(IMFActivate *device) {
    std::string ret;
    WCHAR *friendlyName = nullptr;
    uint32_t nameLength = 0;
    auto result = device->GetAllocatedString(MF_DEVSOURCE_ATTRIBUTE_FRIENDLY_NAME, &friendlyName, &nameLength);

    if (SUCCEEDED(result)) {
        ret = Amass::StringUtility::wstringToString(std::wstring(friendlyName, nameLength));
    }
    if (friendlyName != nullptr) {
        CoTaskMemFree(friendlyName);
    }
    return ret;
}

std::shared_ptr<DeviceDiscovery::Device> DeviceDiscovery::find(const std::string &deviceName, std::error_code &error) {
    std::shared_ptr<Device> ret;

    IMFAttributes *attributes = nullptr;
    boost::scope::scope_exit guard([&attributes] { SafeRelease(&attributes); });
    auto result = MFCreateAttributes(&attributes, 1);
    if (FAILED(result)) {
        return ret;
    }

    result = attributes->SetGUID(MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE, MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_VIDCAP_GUID);
    if (FAILED(result)) {
        return ret;
    }

    UINT32 count;
    IMFActivate **devices = nullptr;
    result = MFEnumDeviceSources(attributes, &devices, &count);
    if (FAILED(result)) {
        return ret;
    }

    if (count == 0) {
        return ret;
    }

    int index = -1;
    for (int i = 0; i < count; i++) {
        auto name = ::deviceName(devices[i]);
        LOG(info) << "device[" << i << "]: " << name;
        if (name == deviceName) {
            index = i;
            break;
        }
    }
    if (index >= 0) {
        IMFMediaSource *source = nullptr;
        result = devices[index]->ActivateObject(IID_PPV_ARGS(&source));
        if (FAILED(result)) {
        } else {
            ret = std::make_shared<Device>(source);
        }
    }
    return ret;
}

void DeviceDiscovery::enterOtaMode(const std::shared_ptr<Device> &device, std::error_code &error) {
    auto resolutions = deviceResolutions(device);
    LOG(info) << "device resolutions:";
    for (auto &[w, h] : resolutions) {
        LOG(info) << "\t" << w << "*" << h;
    }

    int32_t otaSpecificHeight = -1;
    for (auto &[width, height] : resolutions) {
        if (width == OtaSpecificWidth) {
            otaSpecificHeight = height;
            break;
        }
    }
    if (otaSpecificHeight <= 0) {
        LOG(error) << "cannot find ota specific resolution.";
        return;
    } else {
        LOG(info) << "found ota specific resolution: " << OtaSpecificWidth << "x" << otaSpecificHeight;
    }

    IMFMediaType *type = nullptr;
    auto result = MFCreateMediaType(&type);
    if (SUCCEEDED(result)) {
        result = type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video);
        result = type->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_MJPG);
        result = MFSetAttributeSize(type, MF_MT_FRAME_SIZE, OtaSpecificWidth, otaSpecificHeight);
        result = type->SetUINT32(MF_MT_ALL_SAMPLES_INDEPENDENT, TRUE);
        if (SUCCEEDED(result)) {
            result = device->reader->SetCurrentMediaType((DWORD)MF_SOURCE_READER_FIRST_VIDEO_STREAM, nullptr, type);
        }
        type->Release();
    }
    DWORD streamIndex, flags;
    LONGLONG llTimeStamp;
    IMFSample *pSample = NULL;
    result = device->reader->ReadSample((DWORD)MF_SOURCE_READER_FIRST_VIDEO_STREAM, 0, &streamIndex, &flags,
                                        &llTimeStamp, &pSample);
}

std::vector<std::string> DeviceDiscovery::devices() {
    std::vector<std::string> ret;
    IMFAttributes *attributes = nullptr;
    boost::scope::scope_exit guard([&attributes] { SafeRelease(&attributes); });
    auto result = MFCreateAttributes(&attributes, 1);
    if (FAILED(result)) {
        return ret;
    }

    result = attributes->SetGUID(MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE, MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_VIDCAP_GUID);
    if (FAILED(result)) {
        return ret;
    }

    UINT32 count;
    IMFActivate **devices = nullptr;
    result = MFEnumDeviceSources(attributes, &devices, &count);
    if (FAILED(result)) {
        return ret;
    }

    if (count == 0) {
        return ret;
    }

    for (int i = 0; i < count; i++) {
        auto name = ::deviceName(devices[i]);
        ret.push_back(name);
    }
    return ret;
}

DeviceDiscovery::Resolutions DeviceDiscovery::deviceResolutions(const std::shared_ptr<Device> &source) {
    DeviceDiscovery::Resolutions ret;
    HRESULT result = S_OK;
    DWORD index = 0;
    while (SUCCEEDED(result)) {
        IMFMediaType *mediaType = nullptr;
        result = source->reader->GetNativeMediaType(MF_SOURCE_READER_FIRST_VIDEO_STREAM, index++, &mediaType);
        if (SUCCEEDED(result)) {
            UINT32 width, height;
            MFGetAttributeSize(mediaType, MF_MT_FRAME_SIZE, &width, &height);
            mediaType->Release();
            auto iterator = std::find_if(ret.begin(), ret.end(), [&width, &height](const Resolution &resolution) {
                return (resolution.first == width) && (resolution.second == height);
            });
            if (iterator == ret.end()) {
                ret.push_back({width, height});
            }
        }
    }
    return ret;
}

DeviceDiscovery::Device::Device(IMFMediaSource *source) : source(source) {
    source->AddRef();
    auto result = MFCreateSourceReaderFromMediaSource(source, nullptr, &reader);
    if (FAILED(result)) {
        LOG(error) << "MFCreateSourceReaderFromMediaSource() failed, result: " << result;
    }
}

DeviceDiscovery::Device::~Device() {
    if (source != nullptr) {
        source->Release();
    }
    if (reader != nullptr) {
        reader->Release();
    }
}