// Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#include "ndkcamera.h"

#include <string>

#include <android/log.h>

#include <opencv2/core/core.hpp>

#include "mat.h"

static void onDisconnected(void *context, ACameraDevice *device) {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onDisconnected %p", device);
}

static void onError(void *context, ACameraDevice *device, int error) {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onError %p %d", device, error);
}

static void onImageAvailable(void *context, AImageReader *reader) {
//     __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onImageAvailable %p", reader);

    AImage *image = nullptr;
    media_status_t status = AImageReader_acquireLatestImage(reader, &image);

    if (status != AMEDIA_OK) {
        // error
        return;
    }

    int32_t format;
    AImage_getFormat(image, &format);

    // assert format == AIMAGE_FORMAT_YUV_420_888

    int32_t width = 0;
    int32_t height = 0;
    AImage_getWidth(image, &width);
    AImage_getHeight(image, &height);

    int32_t y_pixelStride = 0;
    int32_t u_pixelStride = 0;
    int32_t v_pixelStride = 0;
    AImage_getPlanePixelStride(image, 0, &y_pixelStride);
    AImage_getPlanePixelStride(image, 1, &u_pixelStride);
    AImage_getPlanePixelStride(image, 2, &v_pixelStride);

    int32_t y_rowStride = 0;
    int32_t u_rowStride = 0;
    int32_t v_rowStride = 0;
    AImage_getPlaneRowStride(image, 0, &y_rowStride);
    AImage_getPlaneRowStride(image, 1, &u_rowStride);
    AImage_getPlaneRowStride(image, 2, &v_rowStride);

    uint8_t *y_data = nullptr;
    uint8_t *u_data = nullptr;
    uint8_t *v_data = nullptr;
    int y_len = 0;
    int u_len = 0;
    int v_len = 0;
    AImage_getPlaneData(image, 0, &y_data, &y_len);
    AImage_getPlaneData(image, 1, &u_data, &u_len);
    AImage_getPlaneData(image, 2, &v_data, &v_len);

    if (u_data == v_data + 1 && v_data == y_data + width * height && y_pixelStride == 1 &&
        u_pixelStride == 2 && v_pixelStride == 2 && y_rowStride == width && u_rowStride == width &&
        v_rowStride == width) {
        // already nv21  :)
        ((NdkCamera *) context)->on_image((unsigned char *) y_data, (int) width, (int) height);
    } else {
        // construct nv21
        auto *nv21 = new unsigned char[width * height + width * height / 2];
        {
            // Y
            unsigned char *y_ptr = nv21;
            for (int y = 0; y < height; y++) {
                const unsigned char *y_data_ptr = y_data + y_rowStride * y;
                for (int x = 0; x < width; x++) {
                    y_ptr[0] = y_data_ptr[0];
                    y_ptr++;
                    y_data_ptr += y_pixelStride;
                }
            }

            // UV
            unsigned char *uv_ptr = nv21 + width * height;
            for (int y = 0; y < height / 2; y++) {
                const unsigned char *v_data_ptr = v_data + v_rowStride * y;
                const unsigned char *u_data_ptr = u_data + u_rowStride * y;
                for (int x = 0; x < width / 2; x++) {
                    uv_ptr[0] = v_data_ptr[0];
                    uv_ptr[1] = u_data_ptr[0];
                    uv_ptr += 2;
                    v_data_ptr += v_pixelStride;
                    u_data_ptr += u_pixelStride;
                }
            }
        }

        ((NdkCamera *) context)->on_image((unsigned char *) nv21, (int) width, (int) height);

        delete[] nv21;
    }

    AImage_delete(image);
}

static void onSessionActive(void *context, ACameraCaptureSession *session) {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onSessionActive %p", session);
}

static void onSessionReady(void *context, ACameraCaptureSession *session) {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onSessionReady %p", session);
}

static void onSessionClosed(void *context, ACameraCaptureSession *session) {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onSessionClosed %p", session);
}

void onCaptureFailed(void *context, ACameraCaptureSession *session, ACaptureRequest *request,
                     ACameraCaptureFailure *failure) {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onCaptureFailed %p %p %p", session, request,
                        failure);
}

void onCaptureSequenceCompleted(void *context, ACameraCaptureSession *session, int sequenceId,
                                int64_t frameNumber) {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onCaptureSequenceCompleted %p %d %ld",
                        session, sequenceId, frameNumber);
}

void onCaptureSequenceAborted(void *context, ACameraCaptureSession *session, int sequenceId) {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onCaptureSequenceAborted %p %d", session,
                        sequenceId);
}

void onCaptureCompleted(void *context, ACameraCaptureSession *session, ACaptureRequest *request,
                        const ACameraMetadata *result) {
//     __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "onCaptureCompleted %p %p %p", session, request, result);
}

NdkCamera::NdkCamera() {
    camera_facing = 0;
    camera_orientation = 0;

    camera_manager = nullptr;
    camera_device = nullptr;
    image_reader = nullptr;
    image_reader_surface = nullptr;
    image_reader_target = nullptr;
    capture_request = nullptr;
    capture_session_output_container = nullptr;
    capture_session_output = nullptr;
    capture_session = nullptr;


    // setup imagereader and its surface
    {
        AImageReader_new(640, 480, AIMAGE_FORMAT_YUV_420_888, /*maxImages*/2, &image_reader);

        AImageReader_ImageListener listener;
        listener.context = this;
        listener.onImageAvailable = onImageAvailable;

        AImageReader_setImageListener(image_reader, &listener);

        AImageReader_getWindow(image_reader, &image_reader_surface);

        ANativeWindow_acquire(image_reader_surface);
    }
}

NdkCamera::~NdkCamera() {
    close();

    if (image_reader) {
        AImageReader_delete(image_reader);
        image_reader = nullptr;
    }

    if (image_reader_surface) {
        ANativeWindow_release(image_reader_surface);
        image_reader_surface = nullptr;
    }
}

int NdkCamera::open(int _camera_facing) {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "open");

    camera_facing = _camera_facing;

    camera_manager = ACameraManager_create();

    // find front camera
    std::string camera_id;
    {
        ACameraIdList *camera_id_list = nullptr;
        ACameraManager_getCameraIdList(camera_manager, &camera_id_list);

        for (int i = 0; i < camera_id_list->numCameras; ++i) {
            const char *id = camera_id_list->cameraIds[i];
            ACameraMetadata *camera_metadata = nullptr;
            ACameraManager_getCameraCharacteristics(camera_manager, id, &camera_metadata);

            // query faceing
            acamera_metadata_enum_android_lens_facing_t facing = ACAMERA_LENS_FACING_FRONT;
            {
                ACameraMetadata_const_entry e = {0};
                ACameraMetadata_getConstEntry(camera_metadata, ACAMERA_LENS_FACING, &e);
                facing = (acamera_metadata_enum_android_lens_facing_t) e.data.u8[0];
            }

            if (camera_facing == 0 && facing != ACAMERA_LENS_FACING_FRONT) {
                ACameraMetadata_free(camera_metadata);
                continue;
            }

            if (camera_facing == 1 && facing != ACAMERA_LENS_FACING_BACK) {
                ACameraMetadata_free(camera_metadata);
                continue;
            }

            camera_id = id;

            // query orientation
            int orientation = 0;
            {
                ACameraMetadata_const_entry e = {0};
                ACameraMetadata_getConstEntry(camera_metadata, ACAMERA_SENSOR_ORIENTATION, &e);

                orientation = (int) e.data.i32[0];
            }

            camera_orientation = orientation;

            ACameraMetadata_free(camera_metadata);

            break;
        }

        ACameraManager_deleteCameraIdList(camera_id_list);
    }

    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "open %s %d", camera_id.c_str(),
                        camera_orientation);

    // open camera
    {
        ACameraDevice_StateCallbacks camera_device_state_callbacks;
        camera_device_state_callbacks.context = this;
        camera_device_state_callbacks.onDisconnected = onDisconnected;
        camera_device_state_callbacks.onError = onError;

        ACameraManager_openCamera(camera_manager, camera_id.c_str(), &camera_device_state_callbacks,
                                  &camera_device);
    }

    // capture request
    {
        ACameraDevice_createCaptureRequest(camera_device, TEMPLATE_PREVIEW, &capture_request);

        ACameraOutputTarget_create(image_reader_surface, &image_reader_target);
        ACaptureRequest_addTarget(capture_request, image_reader_target);
    }

    // capture session
    {
        ACameraCaptureSession_stateCallbacks camera_capture_session_state_callbacks;
        camera_capture_session_state_callbacks.context = this;
        camera_capture_session_state_callbacks.onActive = onSessionActive;
        camera_capture_session_state_callbacks.onReady = onSessionReady;
        camera_capture_session_state_callbacks.onClosed = onSessionClosed;

        ACaptureSessionOutputContainer_create(&capture_session_output_container);

        ACaptureSessionOutput_create(image_reader_surface, &capture_session_output);

        ACaptureSessionOutputContainer_add(capture_session_output_container,
                                           capture_session_output);

        ACameraDevice_createCaptureSession(camera_device, capture_session_output_container,
                                           &camera_capture_session_state_callbacks,
                                           &capture_session);

        ACameraCaptureSession_captureCallbacks camera_capture_session_capture_callbacks;
        camera_capture_session_capture_callbacks.context = this;
        camera_capture_session_capture_callbacks.onCaptureStarted = nullptr;
        camera_capture_session_capture_callbacks.onCaptureProgressed = nullptr;
        camera_capture_session_capture_callbacks.onCaptureCompleted = onCaptureCompleted;
        camera_capture_session_capture_callbacks.onCaptureFailed = onCaptureFailed;
        camera_capture_session_capture_callbacks.onCaptureSequenceCompleted = onCaptureSequenceCompleted;
        camera_capture_session_capture_callbacks.onCaptureSequenceAborted = onCaptureSequenceAborted;
        camera_capture_session_capture_callbacks.onCaptureBufferLost = nullptr;

        ACameraCaptureSession_setRepeatingRequest(capture_session,
                                                  &camera_capture_session_capture_callbacks, 1,
                                                  &capture_request, nullptr);
    }

    return 0;
}

void NdkCamera::close() {
    __android_log_print(ANDROID_LOG_WARN, "NdkCamera", "close");

    if (capture_session) {
        ACameraCaptureSession_stopRepeating(capture_session);
        ACameraCaptureSession_close(capture_session);
        capture_session = nullptr;
    }

    if (camera_device) {
        ACameraDevice_close(camera_device);
        camera_device = nullptr;
    }

    if (capture_session_output_container) {
        ACaptureSessionOutputContainer_free(capture_session_output_container);
        capture_session_output_container = nullptr;
    }

    if (capture_session_output) {
        ACaptureSessionOutput_free(capture_session_output);
        capture_session_output = nullptr;
    }

    if (capture_request) {
        ACaptureRequest_free(capture_request);
        capture_request = nullptr;
    }

    if (image_reader_target) {
        ACameraOutputTarget_free(image_reader_target);
        image_reader_target = nullptr;
    }

    if (camera_manager) {
        ACameraManager_delete(camera_manager);
        camera_manager = nullptr;
    }
}

void NdkCamera::on_image(const cv::Mat &rgb) const {
}

void NdkCamera::on_image(const unsigned char *nv21, int nv21_width, int nv21_height) const {
    // rotate nv21
    int w = 0;
    int h = 0;
    int rotate_type = 0;
    {
        if (camera_orientation == 0) {
            w = nv21_width;
            h = nv21_height;
            rotate_type = camera_facing == 0 ? 2 : 1;
        }
        if (camera_orientation == 90) {
            w = nv21_height;
            h = nv21_width;
            rotate_type = camera_facing == 0 ? 5 : 6;
        }
        if (camera_orientation == 180) {
            w = nv21_width;
            h = nv21_height;
            rotate_type = camera_facing == 0 ? 4 : 3;
        }
        if (camera_orientation == 270) {
            w = nv21_height;
            h = nv21_width;
            rotate_type = camera_facing == 0 ? 7 : 8;
        }
    }

    cv::Mat nv21_rotated(h + h / 2, w, CV_8UC1);
    ncnn::kanna_rotate_yuv420sp(nv21, nv21_width, nv21_height, nv21_rotated.data, w, h,
                                rotate_type);

    // nv21_rotated to rgb
    cv::Mat rgb(h, w, CV_8UC3);
    ncnn::yuv420sp2rgb(nv21_rotated.data, w, h, rgb.data);

    on_image(rgb);
}

static const int NDKCAMERAWINDOW_ID = 233;

NdkCameraWindow::NdkCameraWindow() : NdkCamera() {
    sensor_manager = nullptr;
    sensor_event_queue = nullptr;
    accelerometer_sensor = nullptr;
    win = nullptr;

    accelerometer_orientation = 0;

    // sensor
    sensor_manager = ASensorManager_getInstance();

    accelerometer_sensor = ASensorManager_getDefaultSensor(sensor_manager,
                                                           ASENSOR_TYPE_ACCELEROMETER);
}

NdkCameraWindow::~NdkCameraWindow() {
    if (accelerometer_sensor) {
        ASensorEventQueue_disableSensor(sensor_event_queue, accelerometer_sensor);
        accelerometer_sensor = nullptr;
    }

    if (sensor_event_queue) {
        ASensorManager_destroyEventQueue(sensor_manager, sensor_event_queue);
        sensor_event_queue = nullptr;
    }

    if (win) {
        ANativeWindow_release(win);
    }
}

void NdkCameraWindow::set_window(ANativeWindow *_win) {
    if (win) {
        ANativeWindow_release(win);
    }

    win = _win;
    ANativeWindow_acquire(win);
}

void NdkCameraWindow::on_image_render(cv::Mat &rgb) const {
}

void NdkCameraWindow::on_image(const unsigned char *nv21, int nv21_width, int nv21_height) const {
    // resolve orientation from camera_orientation and accelerometer_sensor
    {
        if (!sensor_event_queue) {
            sensor_event_queue = ASensorManager_createEventQueue(sensor_manager, ALooper_prepare(
                    ALOOPER_PREPARE_ALLOW_NON_CALLBACKS), NDKCAMERAWINDOW_ID, 0, 0);

            ASensorEventQueue_enableSensor(sensor_event_queue, accelerometer_sensor);
        }

        int id = ALooper_pollAll(0, nullptr, nullptr, nullptr);
        if (id == NDKCAMERAWINDOW_ID) {
            ASensorEvent e[8];
            ssize_t num_event = 0;
            while (ASensorEventQueue_hasEvents(sensor_event_queue) == 1) {
                num_event = ASensorEventQueue_getEvents(sensor_event_queue, e, 8);
                if (num_event < 0)
                    break;
            }

            if (num_event > 0) {
                float acceleration_x = e[num_event - 1].acceleration.x;
                float acceleration_y = e[num_event - 1].acceleration.y;
                float acceleration_z = e[num_event - 1].acceleration.z;
//                 __android_log_print(ANDROID_LOG_WARN, "NdkCameraWindow", "x = %f, y = %f, z = %f", x, y, z);

                if (acceleration_y > 7) {
                    accelerometer_orientation = 0;
                }
                if (acceleration_x < -7) {
                    accelerometer_orientation = 90;
                }
                if (acceleration_y < -7) {
                    accelerometer_orientation = 180;
                }
                if (acceleration_x > 7) {
                    accelerometer_orientation = 270;
                }
            }
        }
    }

    // roi crop and rotate nv21
    int nv21_roi_x = 0;
    int nv21_roi_y = 0;
    int nv21_roi_w = 0;
    int nv21_roi_h = 0;
    int roi_x = 0;
    int roi_y = 0;
    int roi_w = 0;
    int roi_h = 0;
    int rotate_type = 0;
    int render_w = 0;
    int render_h = 0;
    int render_rotate_type = 0;
    {
        int win_w = ANativeWindow_getWidth(win);
        int win_h = ANativeWindow_getHeight(win);

        if (accelerometer_orientation == 90 || accelerometer_orientation == 270) {
            std::swap(win_w, win_h);
        }

        const int final_orientation = (camera_orientation + accelerometer_orientation) % 360;

        if (final_orientation == 0 || final_orientation == 180) {
            if (win_w * nv21_height > win_h * nv21_width) {
                roi_w = nv21_width;
                roi_h = (nv21_width * win_h / win_w) / 2 * 2;
                roi_x = 0;
                roi_y = ((nv21_height - roi_h) / 2) / 2 * 2;
            } else {
                roi_h = nv21_height;
                roi_w = (nv21_height * win_w / win_h) / 2 * 2;
                roi_x = ((nv21_width - roi_w) / 2) / 2 * 2;
                roi_y = 0;
            }

            nv21_roi_x = roi_x;
            nv21_roi_y = roi_y;
            nv21_roi_w = roi_w;
            nv21_roi_h = roi_h;
        }
        if (final_orientation == 90 || final_orientation == 270) {
            if (win_w * nv21_width > win_h * nv21_height) {
                roi_w = nv21_height;
                roi_h = (nv21_height * win_h / win_w) / 2 * 2;
                roi_x = 0;
                roi_y = ((nv21_width - roi_h) / 2) / 2 * 2;
            } else {
                roi_h = nv21_width;
                roi_w = (nv21_width * win_w / win_h) / 2 * 2;
                roi_x = ((nv21_height - roi_w) / 2) / 2 * 2;
                roi_y = 0;
            }

            nv21_roi_x = roi_y;
            nv21_roi_y = roi_x;
            nv21_roi_w = roi_h;
            nv21_roi_h = roi_w;
        }

        if (camera_facing == 0) {
            if (camera_orientation == 0 && accelerometer_orientation == 0) {
                rotate_type = 2;
            }
            if (camera_orientation == 0 && accelerometer_orientation == 90) {
                rotate_type = 7;
            }
            if (camera_orientation == 0 && accelerometer_orientation == 180) {
                rotate_type = 4;
            }
            if (camera_orientation == 0 && accelerometer_orientation == 270) {
                rotate_type = 5;
            }
            if (camera_orientation == 90 && accelerometer_orientation == 0) {
                rotate_type = 5;
            }
            if (camera_orientation == 90 && accelerometer_orientation == 90) {
                rotate_type = 2;
            }
            if (camera_orientation == 90 && accelerometer_orientation == 180) {
                rotate_type = 7;
            }
            if (camera_orientation == 90 && accelerometer_orientation == 270) {
                rotate_type = 4;
            }
            if (camera_orientation == 180 && accelerometer_orientation == 0) {
                rotate_type = 4;
            }
            if (camera_orientation == 180 && accelerometer_orientation == 90) {
                rotate_type = 5;
            }
            if (camera_orientation == 180 && accelerometer_orientation == 180) {
                rotate_type = 2;
            }
            if (camera_orientation == 180 && accelerometer_orientation == 270) {
                rotate_type = 7;
            }
            if (camera_orientation == 270 && accelerometer_orientation == 0) {
                rotate_type = 7;
            }
            if (camera_orientation == 270 && accelerometer_orientation == 90) {
                rotate_type = 4;
            }
            if (camera_orientation == 270 && accelerometer_orientation == 180) {
                rotate_type = 5;
            }
            if (camera_orientation == 270 && accelerometer_orientation == 270) {
                rotate_type = 2;
            }
        } else {
            if (final_orientation == 0) {
                rotate_type = 1;
            }
            if (final_orientation == 90) {
                rotate_type = 6;
            }
            if (final_orientation == 180) {
                rotate_type = 3;
            }
            if (final_orientation == 270) {
                rotate_type = 8;
            }
        }

        if (accelerometer_orientation == 0) {
            render_w = roi_w;
            render_h = roi_h;
            render_rotate_type = 1;
        }
        if (accelerometer_orientation == 90) {
            render_w = roi_h;
            render_h = roi_w;
            render_rotate_type = 8;
        }
        if (accelerometer_orientation == 180) {
            render_w = roi_w;
            render_h = roi_h;
            render_rotate_type = 3;
        }
        if (accelerometer_orientation == 270) {
            render_w = roi_h;
            render_h = roi_w;
            render_rotate_type = 6;
        }
    }

    // crop and rotate nv21
    cv::Mat nv21_croprotated(roi_h + roi_h / 2, roi_w, CV_8UC1);
    {
        const unsigned char *srcY = nv21 + nv21_roi_y * nv21_width + nv21_roi_x;
        unsigned char *dstY = nv21_croprotated.data;
        ncnn::kanna_rotate_c1(srcY, nv21_roi_w, nv21_roi_h, nv21_width, dstY, roi_w, roi_h, roi_w,
                              rotate_type);

        const unsigned char *srcUV =
                nv21 + nv21_width * nv21_height + nv21_roi_y * nv21_width / 2 + nv21_roi_x;
        unsigned char *dstUV = nv21_croprotated.data + roi_w * roi_h;
        ncnn::kanna_rotate_c2(srcUV, nv21_roi_w / 2, nv21_roi_h / 2, nv21_width, dstUV, roi_w / 2,
                              roi_h / 2, roi_w, rotate_type);
    }

    // nv21_croprotated to rgb
    cv::Mat rgb(roi_h, roi_w, CV_8UC3);
    ncnn::yuv420sp2rgb(nv21_croprotated.data, roi_w, roi_h, rgb.data);

    on_image_render(rgb);

    // rotate to native window orientation
    cv::Mat rgb_render(render_h, render_w, CV_8UC3);
    ncnn::kanna_rotate_c3(rgb.data, roi_w, roi_h, rgb_render.data, render_w, render_h,
                          render_rotate_type);

    ANativeWindow_setBuffersGeometry(win, render_w, render_h,
                                     AHARDWAREBUFFER_FORMAT_R8G8B8A8_UNORM);

    ANativeWindow_Buffer buf;
    ANativeWindow_lock(win, &buf, NULL);

    // scale to target size
    if (buf.format == AHARDWAREBUFFER_FORMAT_R8G8B8A8_UNORM ||
        buf.format == AHARDWAREBUFFER_FORMAT_R8G8B8X8_UNORM) {
        for (int y = 0; y < render_h; y++) {
            const auto *ptr = rgb_render.ptr<const unsigned char>(y);
            unsigned char *outptr = (unsigned char *) buf.bits + buf.stride * 4 * y;

            int x = 0;
#if __ARM_NEON
            for (; x + 7 < render_w; x += 8)
            {
                uint8x8x3_t _rgb = vld3_u8(ptr);
                uint8x8x4_t _rgba;
                _rgba.val[0] = _rgb.val[0];
                _rgba.val[1] = _rgb.val[1];
                _rgba.val[2] = _rgb.val[2];
                _rgba.val[3] = vdup_n_u8(255);
                vst4_u8(outptr, _rgba);

                ptr += 24;
                outptr += 32;
            }
#endif // __ARM_NEON
            for (; x < render_w; x++) {
                outptr[0] = ptr[0];
                outptr[1] = ptr[1];
                outptr[2] = ptr[2];
                outptr[3] = 255;

                ptr += 3;
                outptr += 4;
            }
        }
    }

    ANativeWindow_unlockAndPost(win);
}