#include "buffer_converter.h"

#ifdef __aarch64__
#include <arm_neon.h>
#endif

#include <cstring>
#include <algorithm>

namespace buffer_converter {

void copy_plane_neon(
    const uint8_t* src, int src_stride,
    uint8_t* dst, int dst_stride,
    int width, int height) {

#ifdef __aarch64__
    for (int row = 0; row < height; row++) {
        const uint8_t* src_row = src + row * src_stride;
        uint8_t* dst_row = dst + row * dst_stride;
        int col = 0;

        for (; col + 15 < width; col += 16) {
            uint8x16_t data = vld1q_u8(src_row + col);
            vst1q_u8(dst_row + col, data);
        }
        for (; col < width; col++) {
            dst_row[col] = src_row[col];
        }
    }
#else
    for (int row = 0; row < height; row++) {
        std::memcpy(dst + row * dst_stride, src + row * src_stride, width);
    }
#endif
}

void interleave_uv_to_nv21_neon(
    const uint8_t* src_u, int src_u_stride,
    const uint8_t* src_v, int src_v_stride,
    uint8_t* dst_vu, int dst_vu_stride,
    int width, int height) {

    int uv_width = width / 2;
    int uv_height = height / 2;

#ifdef __aarch64__
    for (int row = 0; row < uv_height; row++) {
        const uint8_t* u_row = src_u + row * src_u_stride;
        const uint8_t* v_row = src_v + row * src_v_stride;
        uint8_t* dst_row = dst_vu + row * dst_vu_stride;
        int col = 0;

        for (; col + 15 < uv_width; col += 16) {
            uint8x16_t u_data = vld1q_u8(u_row + col);
            uint8x16_t v_data = vld1q_u8(v_row + col);
            uint8x16x2_t vu_interleaved;
            vu_interleaved.val[0] = v_data;
            vu_interleaved.val[1] = u_data;
            vst2q_u8(dst_row + col * 2, vu_interleaved);
        }
        for (; col < uv_width; col++) {
            dst_row[col * 2]     = v_row[col];
            dst_row[col * 2 + 1] = u_row[col];
        }
    }
#else
    for (int row = 0; row < uv_height; row++) {
        const uint8_t* u_row = src_u + row * src_u_stride;
        const uint8_t* v_row = src_v + row * src_v_stride;
        uint8_t* dst_row = dst_vu + row * dst_vu_stride;
        for (int col = 0; col < uv_width; col++) {
            dst_row[col * 2]     = v_row[col];
            dst_row[col * 2 + 1] = u_row[col];
        }
    }
#endif
}

bool yuv420_planar_to_nv21(
    const uint8_t* src_y, int src_y_stride,
    const uint8_t* src_u, int src_u_stride,
    const uint8_t* src_v, int src_v_stride,
    uint8_t* dst_nv21, int dst_y_stride, int dst_uv_stride,
    int width, int height) {

    if (!src_y || !src_u || !src_v || !dst_nv21 || width <= 0 || height <= 0) {
        return false;
    }

    copy_plane_neon(src_y, src_y_stride, dst_nv21, dst_y_stride, width, height);

    uint8_t* dst_vu = dst_nv21 + (dst_y_stride * height);
    interleave_uv_to_nv21_neon(src_u, src_u_stride, src_v, src_v_stride,
                                dst_vu, dst_uv_stride, width, height);

    return true;
}

bool yuv420_planar_copy_with_stride(
    const uint8_t* src_y, int src_y_stride,
    const uint8_t* src_u, int src_u_stride,
    const uint8_t* src_v, int src_v_stride,
    uint8_t* dst_y, int dst_y_stride,
    uint8_t* dst_u, int dst_u_stride,
    uint8_t* dst_v, int dst_v_stride,
    int width, int height) {

    if (!src_y || !src_u || !src_v || !dst_y || !dst_u || !dst_v ||
        width <= 0 || height <= 0) {
        return false;
    }

    copy_plane_neon(src_y, src_y_stride, dst_y, dst_y_stride, width, height);

    int uv_width = width / 2;
    int uv_height = height / 2;
    copy_plane_neon(src_u, src_u_stride, dst_u, dst_u_stride, uv_width, uv_height);
    copy_plane_neon(src_v, src_v_stride, dst_v, dst_v_stride, uv_width, uv_height);

    return true;
}

bool nv21_to_yuv420_planar(
    const uint8_t* src_nv21, int src_y_stride, int src_uv_stride,
    uint8_t* dst_y, int dst_y_stride,
    uint8_t* dst_u, int dst_u_stride,
    uint8_t* dst_v, int dst_v_stride,
    int width, int height) {

    if (!src_nv21 || !dst_y || !dst_u || !dst_v || width <= 0 || height <= 0) {
        return false;
    }

    copy_plane_neon(src_nv21, src_y_stride, dst_y, dst_y_stride, width, height);

    int uv_width = width / 2;
    int uv_height = height / 2;
    const uint8_t* src_vu = src_nv21 + (src_y_stride * height);

#ifdef __aarch64__
    for (int row = 0; row < uv_height; row++) {
        const uint8_t* vu_row = src_vu + row * src_uv_stride;
        uint8_t* u_row = dst_u + row * dst_u_stride;
        uint8_t* v_row = dst_v + row * dst_v_stride;
        int col = 0;

        for (; col + 15 < uv_width; col += 16) {
            uint8x16x2_t vu = vld2q_u8(vu_row + col * 2);
            vst1q_u8(v_row + col, vu.val[0]);
            vst1q_u8(u_row + col, vu.val[1]);
        }
        for (; col < uv_width; col++) {
            v_row[col] = vu_row[col * 2];
            u_row[col] = vu_row[col * 2 + 1];
        }
    }
#else
    for (int row = 0; row < uv_height; row++) {
        const uint8_t* vu_row = src_vu + row * src_uv_stride;
        uint8_t* u_row = dst_u + row * dst_u_stride;
        uint8_t* v_row = dst_v + row * dst_v_stride;
        for (int col = 0; col < uv_width; col++) {
            v_row[col] = vu_row[col * 2];
            u_row[col] = vu_row[col * 2 + 1];
        }
    }
#endif

    return true;
}

} // namespace buffer_converter
