#include "pipeline.h"

/**
 * @brief Construct a new PipeLine object.
 * 
 * This constructor initializes the image pipeline, including:
 * - Setting up the camera sensor.
 * - Configuring AI and display frame sizes.
 * - Initializing the display device.
 * - Creating the OSD (On-Screen Display) buffer.
 * - Binding the sensor output channel to a display layer.
 * 
 * @param sensor_id          Sensor device ID.
 * @param display_type       Display type (e.g., "ST7701", "HDMI").
 * @param ai_frame_width     Width of AI input frame.
 * @param ai_frame_height    Height of AI input frame.
 * @param display_width      Display output width.
 * @param display_height     Display output height.
 */
PipeLine::PipeLine(int sensor_id, std::string& display_type, size_t ai_frame_width,
                   size_t ai_frame_height, size_t display_width, size_t display_height)
{
    // Initialize basic parameters
    sensor_id_       = sensor_id;
    display_type_    = display_type;
    ai_frame_width_  = ai_frame_width;
    ai_frame_height_ = ai_frame_height;
    display_width_   = display_width;
    display_height_  = display_height;

    // Define input tensor shape [N, C, H, W]
    tensor_shape = {1, 3, ai_frame_height_, ai_frame_width_};

    // Clear frame structure
    memset(&frame_info, 0, sizeof(k_video_frame_info));

    // ---------------- Sensor Initialization ----------------
    sensor = new Sensor(sensor_id_);
    if (!sensor) {
        std::cout << "Too many sensors" << std::endl;
        return;
    }

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

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

    // Configure AI channel (CHN 1)
    if (0x00 != sensor->set_framesize(ai_frame_width_, ai_frame_height_, false, VICAP_CHN_ID_1)) {
        std::cout << "Failed to set sensor framesize" << std::endl;
        return;
    }
    if (0x00 != sensor->set_pixformat(SensorPixelFormat::RGB888P, VICAP_CHN_ID_1)) {
        std::cout << "Failed to set sensor pixelformat" << std::endl;
        return;
    }

    // ---------------- Display Initialization ----------------
    display = &Display::instance();
    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 Sensor to Display ----------------
    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},
        Display::LAYER_VIDEO1,
        {0, 0, std::get<0>(sensor_frame_size), std::get<1>(sensor_frame_size)},
        sensor_frame_pix_format,
        Display::FLAG_ROTATION_90,
        0);

    // ---------------- OSD Frame Buffer Allocation ----------------
    k_u32 buffer_size = display_width_ * display_height_ * 4;
    osd_frame_pool_id = kd_mpi_vb_create_pool_ex(buffer_size, 1, VB_REMAP_MODE_NOCACHE);

    // Check allocation success
    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);

    // Fill OSD frame info
    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        = display_width_;
    osd_frame_info.v_frame.height       = display_height_;
    osd_frame_info.v_frame.pixel_format = PIXEL_FORMAT_BGRA_8888;
    osd_frame_info.v_frame.stride[0]    = display_width_ * 4;  // Bytes per row
    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 OpenCV OSD image wrapper
    osd_frame = cv::Mat(display_height_, display_width_, CV_8UC4,
                        (void *)osd_frame_buffer.virt_addr,
                        osd_frame_info.v_frame.stride[0]);

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

/**
 * @brief Destructor for PipeLine.
 * 
 * Cleans up resources such as sensor and display instances.
 */
PipeLine::~PipeLine()
{
    if (sensor) {
        delete sensor;
        sensor = nullptr;
    }
    if (display) {
        display = nullptr;
    }
}

/**
 * @brief Capture one frame from the sensor (AI channel).
 * 
 * @return k_video_frame_info Captured frame information.
 */
k_video_frame_info PipeLine::get_frame()
{
    if (sensor) {
        if (0x00 != sensor->snapshot(frame_info, VICAP_CHN_ID_1)) {
            std::cout << "Failed snapshot sensor" << std::endl;
        }
    }
    return frame_info;
}

/**
 * @brief Release the current video frame (placeholder function).
 */
void PipeLine::release_frame()
{
    // Intentionally left blank (can be extended for frame pool management)
}

/**
 * @brief Convert current sensor frame to runtime tensor for AI inference.
 * 
 * This method maps the physical frame memory to virtual address space and
 * constructs a tensor structure compatible with nncase runtime.
 * 
 * @return runtime_tensor The generated input tensor.
 */
runtime_tensor PipeLine::get_tensor()
{
    if (sensor) {
        if (0x00 != sensor->snapshot(frame_info, VICAP_CHN_ID_1)) {
            std::cout << "Failed snapshot sensor" << std::endl;
        }
    }

    // Map frame physical memory for tensor access
    frame_virt_addr = reinterpret_cast<uintptr_t>(
        kd_mpi_sys_mmap(frame_info.v_frame.phys_addr[0],
                        tensor_shape[0] * tensor_shape[1] * tensor_shape[2] * tensor_shape[3]));
    frame_phy_addr = reinterpret_cast<uintptr_t>(frame_info.v_frame.phys_addr[0]);

    // Create tensor with shared physical buffer
    frame_tensor = host_runtime_tensor::create(
                       typecode_t::dt_uint8,
                       tensor_shape,
                       {(gsl::byte *)frame_virt_addr,
                        tensor_shape[0] * tensor_shape[1] * tensor_shape[2] * tensor_shape[3]},
                       false, hrt::pool_shared, frame_phy_addr)
                       .expect("cannot create input tensor");

    // Synchronize memory
    hrt::sync(frame_tensor, sync_op_t::sync_write_back, true).unwrap();

    return frame_tensor;
}

/**
 * @brief Release the mapped tensor memory.
 */
void PipeLine::release_tensor()
{
    kd_mpi_sys_munmap(reinterpret_cast<void *>(frame_virt_addr),
                      tensor_shape[0] * tensor_shape[1] * tensor_shape[2] * tensor_shape[3]);
}

/**
 * @brief Display the current OSD frame on screen.
 * 
 * The rotation mode depends on the display type:
 * - ST7701: Rotate 90 degrees.
 * - Others: No rotation.
 */
void PipeLine::show_image()
{
    if (display_type_ == "ST7701") {
        display->show(osd_frame_info, std::make_tuple(0, 0),
                      Display::LAYER_OSD0, 255, Display::FLAG_ROTATION_90);
    } else {
        display->show(osd_frame_info, std::make_tuple(0, 0),
                      Display::LAYER_OSD0, 255, Display::FLAG_ROTATION_0);
    }
}
