#include "DeviceDiscovery.h"
#include "BoostLog.h"
#include <boost/scope/scope_exit.hpp>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#ifdef WIN32
#include <comutil.h>
#include <dshow.h>
#include <mfapi.h>
#include <mfcaptureengine.h>
#include <strmif.h>
#else
#include <fcntl.h>
#include <linux/videodev2.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#endif

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

DeviceDiscovery::DeviceDiscovery() {
}

#ifdef WIN32

void DeleteMediaType(AM_MEDIA_TYPE *pmt) {
    if (pmt != nullptr) {
        if (pmt->cbFormat != 0) {
            CoTaskMemFree((PVOID)pmt->pbFormat);
            pmt->cbFormat = 0;
            pmt->pbFormat = nullptr;
        }
        if (pmt->pUnk != nullptr) {
            // Unecessary because pUnk should not be used, but safest.
            pmt->pUnk->Release();
            pmt->pUnk = nullptr;
        }
        CoTaskMemFree((PVOID)pmt);
    }
}

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 = _com_util::ConvertBSTRToString(friendlyName);
    }
    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<DeviceDiscovery::Device> DeviceDiscovery::devices() {
    std::vector<Device> devices;
    ICreateDevEnum *pDevEnum = nullptr;
    IEnumMoniker *pEnum = nullptr;

    HRESULT hr =
        CoCreateInstance(CLSID_SystemDeviceEnum, nullptr, CLSCTX_INPROC_SERVER, IID_ICreateDevEnum, (void **)&pDevEnum);

    if (FAILED(hr)) {
        return devices;
    }

    hr = pDevEnum->CreateClassEnumerator(CLSID_VideoInputDeviceCategory, &pEnum, 0);
    if (hr != S_OK) {
        pDevEnum->Release();
        return devices;
    }

    IMoniker *pMoniker = nullptr;
    while (pEnum->Next(1, &pMoniker, nullptr) == S_OK) {
        Device info;
        GetDeviceNames(pMoniker, info);

        IBaseFilter *pFilter;
        hr = pMoniker->BindToObject(nullptr, nullptr, IID_IBaseFilter, (void **)&pFilter);
        if (SUCCEEDED(hr)) {
            IEnumPins *pEnumPins;
            hr = pFilter->EnumPins(&pEnumPins);
            if (SUCCEEDED(hr)) {
                IPin *pPin;
                while (pEnumPins->Next(1, &pPin, nullptr) == S_OK) {
                    PIN_DIRECTION pinDir;
                    pPin->QueryDirection(&pinDir);
                    if (pinDir == PINDIR_OUTPUT) {
                        GetSupportedResolutions(pPin, info.resolutions);
                    }
                    pPin->Release();
                }
                pEnumPins->Release();
            }
            pFilter->Release();
        }
        devices.push_back(info);
        pMoniker->Release();
    }
    pEnum->Release();
    pDevEnum->Release();
    return devices;
}

bool DeviceDiscovery::SetResolution(const std::string &deviceName, int width, int height) {
    ICreateDevEnum *pDevEnum = nullptr;
    IEnumMoniker *pEnum = nullptr;

    HRESULT hr =
        CoCreateInstance(CLSID_SystemDeviceEnum, nullptr, CLSCTX_INPROC_SERVER, IID_ICreateDevEnum, (void **)&pDevEnum);

    if (FAILED(hr)) {
        return false;
    }

    hr = pDevEnum->CreateClassEnumerator(CLSID_VideoInputDeviceCategory, &pEnum, 0);
    if (hr != S_OK) {
        pDevEnum->Release();
        return false;
    }

    IMoniker *pMoniker = nullptr;
    while (pEnum->Next(1, &pMoniker, nullptr) == S_OK) {
        IPropertyBag *pPropBag;
        hr = pMoniker->BindToStorage(nullptr, nullptr, IID_IPropertyBag, (void **)&pPropBag);
        if (SUCCEEDED(hr)) {
            VARIANT varName;
            VariantInit(&varName);
            hr = pPropBag->Read(L"FriendlyName", &varName, nullptr);
            if (SUCCEEDED(hr)) {
                std::string currentName = _com_util::ConvertBSTRToString(varName.bstrVal);
                if (currentName == deviceName) {
                    IBaseFilter *pFilter;
                    hr = pMoniker->BindToObject(nullptr, nullptr, IID_IBaseFilter, (void **)&pFilter);
                    if (SUCCEEDED(hr)) {
                        if (SetPinResolution(pFilter, width, height)) {
                            pFilter->Release();
                            VariantClear(&varName);
                            pPropBag->Release();
                            pMoniker->Release();
                            pEnum->Release();
                            pDevEnum->Release();
                            return true;
                        }
                        pFilter->Release();
                    }
                }
            }
            VariantClear(&varName);
            pPropBag->Release();
        }
        pMoniker->Release();
    }
    pEnum->Release();
    pDevEnum->Release();
    return false;
}

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;
}

void DeviceDiscovery::GetDeviceNames(IMoniker *pMoniker, Device &info) {
    IPropertyBag *pPropBag;
    HRESULT hr = pMoniker->BindToStorage(nullptr, nullptr, IID_IPropertyBag, (void **)&pPropBag);
    if (SUCCEEDED(hr)) {
        VARIANT varName;
        VariantInit(&varName);
        hr = pPropBag->Read(L"FriendlyName", &varName, nullptr);
        if (SUCCEEDED(hr)) {

            info.friendlyName = _com_util::ConvertBSTRToString(varName.bstrVal);
        }
        VariantClear(&varName);

        hr = pPropBag->Read(L"DevicePath", &varName, nullptr);
        if (SUCCEEDED(hr)) {
            info.alternativeName = _com_util::ConvertBSTRToString(varName.bstrVal);
        }
        VariantClear(&varName);

        pPropBag->Release();
    }
}

void DeviceDiscovery::GetSupportedResolutions(IPin *pPin, std::vector<std::pair<int, int>> &resolutions) {
    IEnumMediaTypes *pEnumMediaTypes = nullptr;
    HRESULT hr = pPin->EnumMediaTypes(&pEnumMediaTypes);
    if (FAILED(hr)) {
        return;
    }

    AM_MEDIA_TYPE *pMediaType = nullptr;
    while (pEnumMediaTypes->Next(1, &pMediaType, nullptr) == S_OK) {
        if (pMediaType->formattype == FORMAT_VideoInfo) {
            VIDEOINFOHEADER *pVih = reinterpret_cast<VIDEOINFOHEADER *>(pMediaType->pbFormat);
            int width = pVih->bmiHeader.biWidth;
            int height = pVih->bmiHeader.biHeight;
            resolutions.emplace_back(width, height);
        }
        DeleteMediaType(pMediaType);
    }
    pEnumMediaTypes->Release();
}

bool DeviceDiscovery::SetPinResolution(IBaseFilter *pFilter, int width, int height) {
    IEnumPins *pEnumPins;
    HRESULT hr = pFilter->EnumPins(&pEnumPins);
    if (FAILED(hr)) {
        return false;
    }

    IPin *pPin;
    bool success = false;
    while (pEnumPins->Next(1, &pPin, nullptr) == S_OK) {
        PIN_DIRECTION pinDir;
        pPin->QueryDirection(&pinDir);
        if (pinDir == PINDIR_OUTPUT) {
            IEnumMediaTypes *pEnumMediaTypes = nullptr;
            hr = pPin->EnumMediaTypes(&pEnumMediaTypes);
            if (SUCCEEDED(hr)) {
                AM_MEDIA_TYPE *pMediaType = nullptr;
                while (pEnumMediaTypes->Next(1, &pMediaType, nullptr) == S_OK) {
                    if (pMediaType->formattype == FORMAT_VideoInfo) {
                        VIDEOINFOHEADER *pVih = reinterpret_cast<VIDEOINFOHEADER *>(pMediaType->pbFormat);
                        if (pVih->bmiHeader.biWidth == width && pVih->bmiHeader.biHeight == height) {
                            hr = pPin->QueryAccept(pMediaType);
                            if (hr == S_OK) {
                                success = true;
                            }
                        }
                    }
                    DeleteMediaType(pMediaType);
                    if (success) break;
                }
                pEnumMediaTypes->Release();
            }
        }
        pPin->Release();
        if (success) break;
    }
    pEnumPins->Release();
    return success;
}

DeviceDiscovery::Device::Device(IMFMediaSource *source) : source(source) {
    source->AddRef();
    auto result = MFCreateSourceReaderFromMediaSource(source, nullptr, &reader);
    if (FAILED(result)) {
        LOG(error) << "MFCreateSourceReaderFromMediaSource() failed, result: " << result;
    }
}
#else
std::vector<DeviceDiscovery::Device> DeviceDiscovery::devices() {
    std::vector<Device> ret;
    for (const auto &entry : std::filesystem::directory_iterator("/dev")) {
        if (entry.is_character_file() && entry.path().string().find("video") != std::string::npos) {
            int fd = open(entry.path().c_str(), O_RDWR);
            if (fd < 0) {
                continue;
            }

            struct v4l2_capability cap;
            if (ioctl(fd, VIDIOC_QUERYCAP, &cap) == 0) {
                if ((strstr(reinterpret_cast<const char *>(cap.card), DeviceName) != nullptr) &&
                    (cap.capabilities & V4L2_CAP_VIDEO_CAPTURE)) {
                    Device device;
                    device.friendlyName = entry.path().string();
                    device.alternativeName = entry.path().string();
                    ret.push_back(device);
                }
            }
            close(fd);
        }
    }
    return ret;
}

static std::string find_video_device_by_name(const std::string &targetName) {
    std::string ret;
    const std::string base_path = "/sys/class/video4linux";
    for (const auto &entry : std::filesystem::directory_iterator(base_path)) {
        if (!std::filesystem::is_directory(entry.path())) continue;

        std::string interface;
        std::ifstream ifs(entry.path().string() + "/device/interface");
        if (ifs.is_open()) {
            std::getline(ifs, interface);
        }
        if (interface != targetName) continue;

        ifs = std::ifstream(entry.path().string() + "/index");
        std::string index;
        if (ifs.is_open()) {
            std::getline(ifs, index);
        }
        if (index != "0") continue;

        ret = "/dev/" + entry.path().filename().string();
        break;
    }
    return ret;
}

std::shared_ptr<DeviceDiscovery::Device> DeviceDiscovery::find(const std::string &deviceName, std::error_code &error) {
    auto ret = std::make_shared<Device>();
    ret->friendlyName = find_video_device_by_name(deviceName);
    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;
    }
    int fd = open(device->friendlyName.c_str(), O_RDWR);
    if (fd <= 0) {
        LOG(error) << "Failed to open device " << device->friendlyName;
    } else {
        struct v4l2_format format;
        memset(&format, 0, sizeof(format));
        format.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
        format.fmt.pix.width = OtaSpecificWidth;
        format.fmt.pix.height = otaSpecificHeight;
        format.fmt.pix.pixelformat = V4L2_PIX_FMT_MJPEG;
        format.fmt.pix.field = V4L2_FIELD_INTERLACED;

        if (ioctl(fd, VIDIOC_S_FMT, &format) == -1) {
            perror("Setting Pixel Format");
            return;
        }
#define CLEAR(x) memset(&(x), 0, sizeof(x))
        struct v4l2_requestbuffers req;
        CLEAR(req);
        req.count = 1;
        req.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
        req.memory = V4L2_MEMORY_MMAP;

        if (ioctl(fd, VIDIOC_REQBUFS, &req) == -1) {
            perror("Requesting Buffer");
            return;
        }

        struct v4l2_buffer buf;
        CLEAR(buf);
        buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
        buf.memory = V4L2_MEMORY_MMAP;
        buf.index = 0;

        if (ioctl(fd, VIDIOC_QUERYBUF, &buf) == -1) {
            perror("Querying Buffer");
            return;
        }

        struct buffer {
            void *start;
            size_t length;
        };

        struct buffer buffer;
        buffer.length = buf.length;
        buffer.start = mmap(NULL, buf.length, PROT_READ | PROT_WRITE, MAP_SHARED, fd, buf.m.offset);

        if (buffer.start == MAP_FAILED) {
            perror("Mapping Buffer");
            return;
        }

        if (ioctl(fd, VIDIOC_QBUF, &buf) == -1) {
            perror("Queue Buffer");
            return;
        }

        enum v4l2_buf_type type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
        if (ioctl(fd, VIDIOC_STREAMON, &type) == -1) {
            perror("Start Capture");
            return;
        }

        if (ioctl(fd, VIDIOC_STREAMOFF, &type) == -1) {
            perror("Stop Capture");
            return;
        }

        if (munmap(buffer.start, buffer.length) == -1) {
            perror("Unmapping Buffer");
            return;
        }

        close(fd);
    }
}

DeviceDiscovery::Resolutions DeviceDiscovery::deviceResolutions(const std::shared_ptr<Device> &source) {
    Resolutions ret;
    int fd = open(source->friendlyName.c_str(), O_RDWR);
    if (fd <= 0) {
        LOG(error) << "Failed to open device " << source->friendlyName;
    } else {
        struct v4l2_fmtdesc fmt;
        fmt.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
        fmt.index = 0;

        while (ioctl(fd, VIDIOC_ENUM_FMT, &fmt) == 0) {
            std::cout << "Pixel Format: " << fmt.description << std::endl;
            struct v4l2_frmsizeenum frmsize;
            frmsize.pixel_format = fmt.pixelformat;
            frmsize.index = 0;

            while (ioctl(fd, VIDIOC_ENUM_FRAMESIZES, &frmsize) == 0) {
                if (frmsize.type == V4L2_FRMSIZE_TYPE_DISCRETE) {
                    ret.emplace_back(frmsize.discrete.width, frmsize.discrete.height);
                } else if (frmsize.type == V4L2_FRMSIZE_TYPE_STEPWISE) {
                    for (int width = frmsize.stepwise.min_width; width <= frmsize.stepwise.max_width;
                         width += frmsize.stepwise.step_width) {
                        for (int height = frmsize.stepwise.min_height; height <= frmsize.stepwise.max_height;
                             height += frmsize.stepwise.step_height) {
                            ret.emplace_back(width, height);
                        }
                    }
                }
                frmsize.index++;
            }
            fmt.index++;
        }
        close(fd);
    }
    return ret;
}
#endif

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