#include <chrono>
#include <iostream>
#include <vector>
#include <utility>
#include <fsdk/FaceEngine.h>

#import <UIKit/UIKit.h>
#import "FaceRecognitionManager.h"

#import "KYC_Demo-Swift.h"

#define OUT(x) std::cout << "[" << __PRETTY_FUNCTION__ << "]: " << x << '\n';

struct Constant {

};

struct VectorArchive: fsdk::IArchive
{
    std::vector<uint8_t>& dataOut;
    int index = 0;

    bool write(const void* data, size_t size) noexcept override {
        const uint8_t* p = reinterpret_cast<const uint8_t*>(data);
        dataOut.insert(dataOut.end(), p, p+size);
        return true;
    }

    bool read(void* data, size_t size) noexcept override {
        assert(size <= dataOut.size()-index);
        memcpy(data, (void*)&dataOut[0+index], size);
        index += size;
        return true;
    }

    void setSizeHint(size_t /*hint*/) noexcept override {}

    VectorArchive(std::vector<uint8_t>& dataref):
    dataOut(dataref)
    {}
};

@implementation FaceRecognitionManager {
    fsdk::IFaceEngineMobilePtr engine;                  // Root SDK object
    fsdk::IDescriptorExtractorPtr descriptorExtractor;  // Face descriptor extractor
    fsdk::IDescriptorMatcherPtr descriptorMatcher;      // Face descriptor matcher
    fsdk::IDetectorPtr detector;
    fsdk::IWarperPtr warper;      // Face descriptor warper

    std::mutex mutex;

    bool loaded;
    bool activated;
    bool completeEdition;
}

static NSString * const kModelsDirectory = @"data";

// returns vlf image in rgb format (the one extractor accepts)
// uiimage is expected to be in bgra
fsdk::Image uiimageToVlfImage(UIImage* uiimage) {
    CGImage* cgimage = uiimage.CGImage;
    NSDictionary *options = @{
                              (NSString*)kCVPixelBufferCGImageCompatibilityKey : @YES,
                              (NSString*)kCVPixelBufferCGBitmapContextCompatibilityKey : @YES,
                              };

    CVPixelBufferRef pxbuffer = NULL;
    CVReturn status = CVPixelBufferCreate(kCFAllocatorDefault,
                                            CGImageGetWidth(cgimage),
                                            CGImageGetHeight(cgimage),
                                            kCVPixelFormatType_32BGRA,
                                            (__bridge CFDictionaryRef) options,
                                            &pxbuffer);
    fsdk::Image resultingImage;

    if (status != kCVReturnSuccess || !pxbuffer) {
        OUT("Operation failed");
        return resultingImage;
    }

    CVPixelBufferLockBaseAddress(pxbuffer, 0);

    CGColorSpaceRef rgbColorSpace = CGColorSpaceCreateDeviceRGB();
    CGContextRef context = CGBitmapContextCreate(CVPixelBufferGetBaseAddress(pxbuffer),
                                                    CGImageGetWidth(cgimage),
                                                    CGImageGetHeight(cgimage),
                                                    8,
                                                    4 * CGImageGetWidth(cgimage),
                                                    rgbColorSpace,
                                                    kCGImageAlphaNoneSkipLast);

    CGContextDrawImage(context, CGRectMake(0, 0, CGImageGetWidth(cgimage), CGImageGetHeight(cgimage)), cgimage);
    CGColorSpaceRelease(rgbColorSpace);
    CGContextRelease(context);

    fsdk::Image vlfimage = fsdk::Image((int32_t)CVPixelBufferGetWidth(pxbuffer),
                                (int32_t)CVPixelBufferGetHeight(pxbuffer),
                                fsdk::Format::R8G8B8X8,
                                CVPixelBufferGetBaseAddress(pxbuffer));

    if(!vlfimage.convert(resultingImage, fsdk::Format::R8G8B8))
        OUT("Failed to convert image to rgb");

    CVPixelBufferUnlockBaseAddress(pxbuffer, 0);
    CVPixelBufferRelease(pxbuffer);

    return resultingImage;
    }

/********************************************************/
UIImage* vlfImagetoUiimage(const fsdk::Image& vlfimage) {
    assert(vlfimage.getFormat() == fsdk::Format::R8G8B8);
    const auto channelNum = 3;
    NSData *data = [NSData dataWithBytes: vlfimage.getData() length: vlfimage.getWidth() * vlfimage.getHeight() * channelNum];
    CGColorSpaceRef colorSpace;
    CGBitmapInfo bitmapInfo;

    colorSpace = CGColorSpaceCreateDeviceRGB();
    bitmapInfo = kCGBitmapByteOrder32Big | kCGImageAlphaNone;

    CGDataProviderRef provider = CGDataProviderCreateWithCFData((__bridge CFDataRef)data);

    // Creating CGImage from cv::Mat
    CGImageRef imageRef = CGImageCreate(
                                        vlfimage.getWidth(),                 //width
                                        vlfimage.getHeight(),                 //height
                                        8,                          //bits per component
                                        8 * channelNum,       //bits per pixel
                                        vlfimage.getWidth() * channelNum,              //bytesPerRow
                                        colorSpace,                 //colorspace
                                        bitmapInfo,                 // bitmap info
                                        provider,                   //CGDataProviderRef
                                        NULL,                       //decode
                                        false,                      //should interpolate
                                        kCGRenderingIntentDefault   //intent
                                        );

    // Getting UIImage from CGImage
    UIImage *finalImage = [UIImage imageWithCGImage:imageRef scale:1 orientation: UIImageOrientationUp];
    CGImageRelease(imageRef);
    CGDataProviderRelease(provider);
    CGColorSpaceRelease(colorSpace);

    return finalImage;
}

#pragma mark - Public

+ (FaceRecognitionManager *)shared {
    static FaceRecognitionManager *sharedManager = nil;
    static dispatch_once_t onceToken;

    dispatch_once(&onceToken, ^{
        sharedManager = [[self alloc] init];
    });

    return sharedManager;
}

- (FaceEngineLoadError)load {
    std::lock_guard<decltype(mutex)> guard(mutex);
    NSString *path = [[[NSBundle bundleWithIdentifier:@"ru.visionlabs.FaceEngine"] resourcePath] stringByAppendingPathComponent:kModelsDirectory];
    fsdk::ISettingsProviderPtr config;
    fsdk::ISettingsProviderPtr runtimeConfig;
    std::string configPath = std::string([path UTF8String]) + "/faceengine.conf";
    std::string licenseConfPath = std::string([path UTF8String]) + "/license.conf";
    std::string runtimeConfPath = std::string([path UTF8String]) + "/runtime.conf";

    loaded = false;
    activated = false;

    OUT("face engine data path: " << [path UTF8String]);
    auto resEngine = fsdk::createFaceEngineMobile([path UTF8String]);
    if (resEngine.isError()) {
        OUT("Failed to create face engine. " << resEngine.what())
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }
    engine = resEngine.getValue();

    if(!engine) {
        OUT("Failed to create face engine.")
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }

//    engine->setDataDirectory([path UTF8String]);

    completeEdition = (engine->getFaceEngineEdition() == fsdk::FaceEngineEdition::CompleteEdition);

    auto resRuntimeConfig = fsdk::createSettingsProvider(runtimeConfPath.c_str());
    if (resRuntimeConfig.isError()) {
        OUT("[ISettingsProvider creating] Failed to create Runtime ISettingsProvider. " << resRuntimeConfig.what());
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }
    runtimeConfig = resRuntimeConfig.getValue();

    if (!runtimeConfig) {
        OUT("[ISettingsProvider creating] Failed to create Runtime ISettingsProvider. ");
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }

    auto resConfig = fsdk::createSettingsProvider(configPath.c_str());
    if (resConfig.isError()) {
        OUT("[ISettingsProvider creating] Failed to create ISettingsProvider. " << resConfig.what());
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }
    config = resConfig.getValue();
    if (!config) {
        OUT("[ISettingsProvider creating] Failed to create ISettingsProvider");
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }

    config->setValue("DescriptorFactory::Settings", "useMobileNet", 1);

    config->setValue("FaceDetV2::Settings", "minFaceSize", (int)[Luna5Settings minIndent]);

    engine->setSettingsProvider(config);
    runtimeConfig->setValue("Runtime", "numThreads", -1);
    engine->setRuntimeSettingsProvider(runtimeConfig);
    fsdk::ILicense* license = engine->getLicense();
    if (!license) {
        OUT("[License creating] Failed to get license");
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }

    //    Make activation
    bool activationResult = fsdk::activateLicense(license, licenseConfPath.c_str());
    if (!(activationResult)) {
        OUT("[Descriptor creating] Failed to activate license");
        return FaceEngineLoadError::LicenseActivationFailed;
    } else {
        activated = true;
        OUT("[Descriptor creating] license is activated");
    }
    auto resDescriptorExtractor = engine->createExtractor();
    if (resDescriptorExtractor.isError()) {
        OUT("Failed to create descriptor extractor. " << resDescriptorExtractor.what());
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }
    descriptorExtractor = resDescriptorExtractor.getValue();
    if (!descriptorExtractor) {
        OUT("Failed to create descriptor extractor");
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }
    auto resDescriptorMatcher = engine->createMatcher();

    if (resDescriptorMatcher.isError()) {
        OUT("Failed to create descriptor matcher. " << resDescriptorMatcher.what());
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }
    descriptorMatcher = resDescriptorMatcher.getValue();
    if (!descriptorMatcher) {
        OUT("Failed to create descriptor matcher");
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }

    // Create default detector, see faceengine.conf - "defaultDetectorType"
    auto resFaceDetector = engine->createDetector();
    if (!resFaceDetector) {
        OUT("Failed to create face detector instance.");
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }
    detector = resFaceDetector.getValue();

    auto resWarper = engine->createWarper();
    if(!resWarper) {
        OUT("Failed to create warper instance.");
        return FaceEngineLoadError::FaceEngineNotLoaded;
    }
    warper = resWarper.getValue();

    loaded = true;
    activated = true;

    return FaceEngineLoadError::FaceEngineLoadSuccess;
}

- (NSData*)extractDescriptorFrom: (UIImage*)image {
    std::lock_guard<decltype(mutex)> guard(mutex);

    if(!loaded) {
        OUT("Not loaded properly");
        return nil;
    }

    // convert UIImage to fsdk::Image
    fsdk::Image rgbimage = uiimageToVlfImage(image);

    auto resDescriptor = engine->createDescriptor();
    if(resDescriptor.isError()) {
        OUT("[extractDescriptorFrom] Failed to create descriptor instance. " << resDescriptor.what());
        return nil;
    }
    fsdk::IDescriptorPtr descriptor = resDescriptor.getValue();
    if(!descriptor) {
        OUT("[extractDescriptorFrom] Failed to create descriptor instance.");
        return nil;
    }
    if(!rgbimage.isValid()) {
        OUT("[extractDescriptorFrom] Image is not valid. No extraction");
        return nil;
    }

    const auto result = descriptorExtractor->extractFromWarpedImage(rgbimage, descriptor);

    if(result) {
        OUT("Descriptor extracted");
        std::vector<uint8_t> descriptorData;
        VectorArchive archiveDescriptor(descriptorData);

        descriptor->save(&archiveDescriptor);
        NSData* nsdata = [[NSData alloc] initWithBytes: descriptorData.data() length:descriptorData.size()];

        if(!nsdata) {
            OUT("Couldnt convert descriptor to nsdata");
            return nil;
        }

        return nsdata;
    } else {
        OUT("Failed to extract descriptor: " << result.what());
        return nil;
    }
}

- (UIImage*)warp: (UIImage*)image {
    std::lock_guard<decltype(mutex)> guard(mutex);

    if(!loaded) {
        OUT("Not loaded properly");
        return nil;
    }

    // convert UIImage to fsdk::Image
    fsdk::Image rgbimage = uiimageToVlfImage(image);


    // Facial feature detection confidence threshold.
    const float confidenceThreshold = 0.25f;

    if (!rgbimage) {
        OUT("Request image is invalid.");
        return nil;
    }

    // Stage 1. Detect a face.
    OUT("Detecting faces.");

    // Detect no more than 10 faces in the image.
    const uint32_t maxDetections = 1;

    // Data used for detection.
    uint32_t detectionsCount(maxDetections);
    fsdk::Rect imageRect = rgbimage.getRect();

    fsdk::ResultValue<fsdk::FSDKError, fsdk::Face> detectorResult = detector->detectOne(rgbimage, imageRect, fsdk::DT_LANDMARKS5);

//    // Detect faces in the image.
//    fsdk::ResultValue<fsdk::FSDKError, fsdk::Ref<fsdk::IFaceDetectionBatch>> detectorResult =
//        detector->detect(
//            fsdk::Span<const fsdk::Image>(&rgbimage, 1),
//            fsdk::Span<const fsdk::Rect>(&imageRect, 1),
//            detectionsCount,
//            fsdk::DT_ALL
//        );

    if (detectorResult.isError()) {
        OUT("Failed to detect face detection. Reason: " << detectorResult.what());
        return nullptr;
    }

    fsdk::Face face = detectorResult.getValue();
//    fsdk::Ref<fsdk::IFaceDetectionBatch> detectionBatch = detectorResult.getValue();
//    detectionsCount = static_cast<uint32_t>(detectionBatch->getSize(0));
//    if (detectionsCount == 0) {
//        OUT("Faces are not found.");
//        return nullptr;
//    }
//
//    OUT("Found " << detectionsCount << " face(s).");

//    const fsdk::Span<const fsdk::Detection>& detections = detectionBatch->getDetections(0);
//    const fsdk::Span<const fsdk::Landmarks5>& landmarks5 = detectionBatch->getLandmarks5(0);

//    int bestDetectionIndex(-1);
//    float bestScore(-1.f);

    // Loop through all the faces and find one with the best score.
//    for (size_t i = 0; i < detectionsCount; ++i) {
//        const fsdk::Detection &detection = detections[i];
//        OUT("Detecting facial features (" << i + 1 << "/" << detectionsCount << ")");
//
//        // Choose the best detection.
//        float score = detection.getScore();
//        if (score > bestScore) {
//            bestScore = score;
//            bestDetectionIndex = static_cast<int>(i);
//        }
//    }

//    // If detection confidence score is too low, abort.
//    if (bestScore < confidenceThreshold) {
//        OUT("Face detection succeeded, but no faces with good confidence found.");
//        return nullptr;
//    }
//    OUT("Best face confidence is " << bestScore);

    // Stage 2. Create face descriptor.
    OUT("Extracting descriptor.");

    // Create face descriptor.
    auto resDescriptor = engine->createDescriptor();
    if (!resDescriptor) {
        OUT("Failed to create face descrtiptor instance.");
        return nullptr;
    }
    fsdk::IDescriptorPtr descriptor = resDescriptor.getValue();

    // Extract face descriptor.
    // This is typically the most time-consuming task.
//    fsdk::Transformation transform = warper->createTransformation(
//        detections[bestDetectionIndex],
//        landmarks5[bestDetectionIndex]
//    );

    if (face.landmarks5.has_value() == false) {
        OUT("Failed to get face landmark.");
        return nullptr;
    }

    fsdk::Transformation transform = warper->createTransformation(
        face.detection,
        face.landmarks5.value()
    );

    fsdk::Image warpedImage = {};
    fsdk::Result<fsdk::FSDKError> warpResult = warper->warp(rgbimage, transform, warpedImage);
    if(warpResult.isError()) {
        OUT("Failed to warp image. Reason: " << warpResult.what());
        return nullptr;
    }

    UIImage *result = vlfImagetoUiimage(warpedImage);

    if(result) {
        return result;
    } else {
        return nil;
    }
}

fsdk::IDescriptorPtr descriptorFromNSData(const fsdk::IFaceEngineMobilePtr engine, NSData* data) {
    const uint8_t* buffer = static_cast<const uint8_t*>([data bytes]);
    std::vector<uint8_t> descriptorData(buffer, buffer + [data length]);
    auto resDescriptor = engine->createDescriptor();
    if(resDescriptor.isError()) {
        OUT("[descriptorFromNSData] Failed to create descriptor instance. " << resDescriptor.what());
        return nil;
    }
    fsdk::IDescriptorPtr descriptor = resDescriptor.getValue();
    VectorArchive archiveDescriptor(descriptorData);
    descriptor->load(&archiveDescriptor);

    if(!descriptor)
        OUT("Failed to get descriptor from NSData");

    return descriptor;
}

// MARK: Estimators
// quality estimation
// only pass WARPED image - you can get it from warpImage method
/********************************************************/
//- (bool)estimateQualityOf: (UIImage*)warpedImage estimation: (float*) estimation {
//    std::lock_guard<decltype(mutex)> guard(mutex);
//
//    if(![self isValid]) {
//        OUT("Object is not valid (probably you didnt call load())");
//        return nil;
//    }
//
//    // check args correctness
//    if(warpedImage == nullptr || estimation == nullptr) {
//        OUT("Invalid arguments");
//        return false;
//    }
//
//    fsdk::Image rgbimage = uiimageToVlfImage(warpedImage);
//    fsdk::SubjectiveQuality quality;
//
//    const auto result = qualityEstimator->estimate(rgbimage, quality);
//
//    *estimation = quality.isGood();
//
//    if(result.isError()) {
//        OUT("Failed to estimate quality. Reason: " << result.what());
//        return false;
//    }
//
//    return true;
//}


- (float)qualityOf:(UIImage *)warpedImage {
    // TODO: 2022!!!
    return 1.0;
}

- (float)matchDescriptor: (NSData*) firstDescriptorData and: (NSData*) secondDescriptorData {
    std::lock_guard<decltype(mutex)> guard(mutex);
    float similarity = 0.f;

    if(!loaded) {
        OUT("Not loaded properly");
        return similarity;
    }

    const auto firstDescriptor = descriptorFromNSData(engine, firstDescriptorData);
    const auto secondDescriptor = descriptorFromNSData(engine, secondDescriptorData);
    const auto result = descriptorMatcher->match(firstDescriptor, secondDescriptor);

    if(result) {
        similarity = result.getValue().similarity;
        OUT("Descriptors matched with similarity=" << similarity);
    }
    else {
        OUT("Failed to match descriptors: " << result.what());
    }

    return similarity;
}

@end