#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>
#include "k230_display.h"
#include "k230_sensor.h"
#include "ai_utils.h"
#include "mpi_sys_api.h"

// ==========================================
// Global Configuration Parameters
// ==========================================
size_t ai_frame_width = 1280;
size_t ai_frame_height = 720;
size_t display_width = 800;
size_t display_height = 480;
size_t osd_frame_width = 800;
size_t osd_frame_height = 480;

// ==========================================
// Global Objects and Buffers
// ==========================================
Sensor* sensor = nullptr;          // Pointer to the camera sensor
VbBuffer osd_frame_buffer;         // Video buffer for OSD (On-Screen Display)
k_video_frame_info osd_frame_info; // Frame metadata for OSD buffer
k_s32 osd_frame_pool_id = VB_INVALID_POOLID; // Video buffer pool ID
cv::Mat osd_frame;                 // OpenCV Mat mapped to the OSD buffer

// ==========================================
// setup()
// Called once to initialize sensor, display, and OSD buffer
// ==========================================
void setup() {
    // Initialize sensor instance (Sensor ID = 2)
    sensor = new Sensor(2);
    if (!sensor) {
        std::cout << "Too many sensors" << std::endl;
        return;
    }

    // Reset and initialize the sensor
    if (sensor->reset() != 0x00) {
        std::cout << "Failed to probe sensor" << std::endl;
        return;
    }

    // Configure channel 0 for display output (YUV420SP format)
    if (sensor->set_framesize(display_width, display_height, false, VICAP_CHN_ID_0) != 0x00) {
        std::cout << "Failed to set sensor framesize" << std::endl;
        return;
    }
    if (sensor->set_pixformat(SensorPixelFormat::YUV420SP, VICAP_CHN_ID_0) != 0x00) {
        std::cout << "Failed to set sensor pixelformat" << std::endl;
        return;
    }

    // Configure channel 1 for AI processing (RGB888 planar format)
    if (sensor->set_framesize(ai_frame_width, ai_frame_height, false, VICAP_CHN_ID_1) != 0x00) {
        std::cout << "Failed to set sensor framesize" << std::endl;
        return;
    }
    if (sensor->set_pixformat(SensorPixelFormat::RGB888P, VICAP_CHN_ID_1) != 0x00) {
        std::cout << "Failed to set sensor pixelformat" << std::endl;
        return;
    }

    // Start video streaming
    if (sensor->run() != 0x00) {
        std::cout << "Failed to run sensor" << std::endl;
        return;
    }

    // ==========================================
    // Initialize Display
    // ==========================================
    Display& display = Display::instance();

    // Initialize ST7701 display with 2-lane MIPI, 480x800 resolution, 30 FPS
    int ret = display.init(ST7701_V1_MIPI_2LAN_480X800_30FPS, 1);
    if (ret != 0) {
        printf("Display init failed. Halting.");
        while (1)
            delay(1000);
    }
    printf("Display initialized: %dx%d\n", display.width(), display.height());

    // ==========================================
    // Bind the sensor channel to the display video layer
    // ==========================================
    int sensor_mod = K_ID_VI;
    int sensor_dev = sensor->_dev_id();
    int sensor_chn = VICAP_CHN_ID_0;

    std::tuple<int, int> sensor_frame_size = sensor->framesize(VICAP_CHN_ID_0);
    k_pixel_format sensor_frame_pix_format =
        Sensor::to_k_pixel_format(sensor->pixformat(VICAP_CHN_ID_0));

    display.bind_layer(
        {sensor_mod, sensor_dev, sensor_chn},  // Source (sensor module/device/channel)
        Display::LAYER_VIDEO1,                 // Target display layer
        {0, 0, std::get<0>(sensor_frame_size), std::get<1>(sensor_frame_size)}, // ROI
        sensor_frame_pix_format,               // Pixel format
        Display::FLAG_ROTATION_90,             // Rotate display 90 degrees
        0                                      // Z-order (optional)
    );

    // ==========================================
    // Create and Map OSD Frame Buffer
    // ==========================================
    k_u32 buffer_size = osd_frame_width * osd_frame_height * 4; // BGRA8888 (4 bytes per pixel)
    osd_frame_pool_id = kd_mpi_vb_create_pool_ex(buffer_size, 1, VB_REMAP_MODE_NOCACHE);

    if ((VB_INVALID_POOLID == osd_frame_pool_id) ||
        (!osd_frame_buffer.get(buffer_size, osd_frame_pool_id) || (!osd_frame_buffer.mmap()))) {
        printf("Failed to allocate source buffer. Halting.");
        while (1)
            delay(1000);
    }
    printf("Source buffer allocated: %d bytes\n", buffer_size);

    // Initialize OSD frame information
    memset(&osd_frame_info, 0, sizeof(osd_frame_info));
    osd_frame_info.pool_id = osd_frame_buffer.pool_id;
    osd_frame_info.v_frame.width = osd_frame_width;
    osd_frame_info.v_frame.height = osd_frame_height;
    osd_frame_info.v_frame.pixel_format = PIXEL_FORMAT_BGRA_8888;
    osd_frame_info.v_frame.stride[0] = osd_frame_width * 4;
    osd_frame_info.v_frame.phys_addr[0] = osd_frame_buffer.phys_addr;
    osd_frame_info.v_frame.virt_addr[0] = (k_u64)osd_frame_buffer.virt_addr;

    // Create an OpenCV Mat object that wraps the OSD buffer memory
    osd_frame = cv::Mat(osd_frame_height, osd_frame_width, CV_8UC4,
                        (void*)osd_frame_buffer.virt_addr,
                        osd_frame_info.v_frame.stride[0]);
}

// ==========================================
// loop()
// Main loop function — captures a frame, creates tensor, and displays OSD
// ==========================================
void loop() {
    PROFILE_SCOPE("total:", 1);  // Timing macro for profiling (if enabled)

    k_video_frame_info vf_info;

    // Capture one frame from the sensor’s AI channel
    if (sensor) {
        if (sensor->snapshot(vf_info, VICAP_CHN_ID_1) != 0x00) {
            std::cout << "Failed snapshot sensor" << std::endl;
            return;
        }
    }

    // Define tensor input shape: (N, C, H, W)
    dims_t input_shape{1, 3, ai_frame_height, ai_frame_width};

    // Map physical address to virtual address for CPU access
    uintptr_t virt_addr = reinterpret_cast<uintptr_t>(
        kd_mpi_sys_mmap(vf_info.v_frame.phys_addr[0],
                        input_shape[0] * input_shape[1] * input_shape[2] * input_shape[3]));

    uintptr_t phy_addr = reinterpret_cast<uintptr_t>(vf_info.v_frame.phys_addr[0]);

    // Create runtime tensor (non-copy shared memory)
    runtime_tensor input_tensor = host_runtime_tensor::create(
        typecode_t::dt_uint8,
        input_shape,
        {(gsl::byte*)virt_addr, input_shape[0] * input_shape[1] * input_shape[2] * input_shape[3]},
        false,                   // Do not copy data
        hrt::pool_shared,        // Shared memory pool
        phy_addr                 // Physical address for DMA access
    ).expect("cannot create input tensor");

    // Synchronize tensor data (flush cache)
    hrt::sync(input_tensor, sync_op_t::sync_write_back, true).unwrap();

    // ==========================================
    // Draw OSD (example text rendering)
    // ==========================================
    osd_frame.setTo(cv::Scalar(0, 0, 0, 0)); // Clear OSD buffer (transparent)
    printf("Got tensor! You can add kmodel inference part!\n");
    cv::putText(osd_frame,
                "Got tensor! You can add kmodel inference part!",
                {20, 20},
                cv::FONT_HERSHEY_COMPLEX,
                0.8,
                cv::Scalar(255, 0, 0, 255),
                2,
                8,
                0);

    // Unmap system memory
    kd_mpi_sys_munmap(reinterpret_cast<void*>(virt_addr),
                      input_shape[0] * input_shape[1] * input_shape[2] * input_shape[3]);

    // ==========================================
    // Show OSD layer on the display
    // ==========================================
    Display& display = Display::instance();
    display.show(osd_frame_info,
                 std::make_tuple(0, 0),
                 Display::LAYER_OSD0,
                 255,
                 Display::FLAG_ROTATION_90);
}
