/*
 *  image.c - Image utilities for the tests
 *
 *  Copyright (C) 2010-2011 Splitted-Desktop Systems
 *
 *  This library is free software; you can redistribute it and/or
 *  modify it under the terms of the GNU Lesser General Public License
 *  as published by the Free Software Foundation; either version 2.1
 *  of the License, or (at your option) any later version.
 *
 *  This library is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *  Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public
 *  License along with this library; if not, write to the Free
 *  Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
 *  Boston, MA 02110-1301 USA
 */

#include "image.h"

static gboolean
image_draw_rectangle(
    GstVaapiImage *image,
    gint           x,
    gint           y,
    guint          width,
    guint          height,
    guint32        color,
    guint32        flags
);

static gboolean
image_draw_color_rectangles(
    GstVaapiImage      *image,
    guint               width,
    guint               height,
    const guint32       colors[4],
    guint32             flags
    )
{
    const guint w = width / 2;
    const guint h = height / 2;

    if (!image_draw_rectangle(image, 0, 0, w, h, colors[0], flags))
        return FALSE;
    if (!image_draw_rectangle(image, w, 0, w, h, colors[1], flags))
        return FALSE;
    if (!image_draw_rectangle(image, 0, h, w, h, colors[2], flags))
        return FALSE;
    if (!image_draw_rectangle(image, w, h, w, h, colors[3], flags))
        return FALSE;
    return TRUE;
}

GstVaapiImage *
image_generate(
    GstVaapiDisplay    *display,
    GstVideoFormat      format,
    guint               width,
    guint               height
)
{
    return image_generate_full(display, format, width, height, 0);
}

GstVaapiImage *
image_generate_full(
    GstVaapiDisplay    *display,
    GstVideoFormat      format,
    guint               width,
    guint               height,
    guint32             flags
)
{
    GstVaapiImage *image;

    static const guint32 rgb_colors[4] =
        { 0xffff0000, 0xff00ff00, 0xff0000ff, 0xff000000 };
    static const guint32 bgr_colors[4] =
        { 0xff000000, 0xff0000ff, 0xff00ff00, 0xffff0000 };
    static const guint32 inv_colors[4] =
        { 0xffdeadc0, 0xffdeadc0, 0xffdeadc0, 0xffdeadc0 };

    image = gst_vaapi_image_new(display, format, width, height);
    if (!image)
        return NULL;

    if (flags) {
        if (!image_draw_color_rectangles(image, width, height,
                ((flags & GST_VAAPI_PICTURE_STRUCTURE_TOP_FIELD) ?
                 rgb_colors : inv_colors),
                GST_VAAPI_PICTURE_STRUCTURE_TOP_FIELD))
            goto error;

        if (!image_draw_color_rectangles(image, width, height,
                ((flags & GST_VAAPI_PICTURE_STRUCTURE_BOTTOM_FIELD) ?
                 bgr_colors : inv_colors),
                GST_VAAPI_PICTURE_STRUCTURE_BOTTOM_FIELD))
            goto error;
    }
    else if (!image_draw_color_rectangles(image, width, height, rgb_colors, 0))
        goto error;
    return image;

error:
    gst_vaapi_object_unref(image);
    return NULL;
}

typedef void (*DrawRectFunc)(
    guchar *pixels[3],
    guint   stride[3],
    gint    x,
    gint    y,
    guint   width,
    guint   height,
    guint32 color
);

static void draw_rect_ARGB(
    guchar *pixels[3],
    guint   stride[3],
    gint    x,
    gint    y,
    guint   width,
    guint   height,
    guint32 color
)
{
    guint i, j;

    color = GUINT32_TO_BE(color);

    for (j = 0; j < height; j++) {
        guint32 *p = (guint32 *)(pixels[0] + (y + j) * stride[0] + x * 4);
        for (i = 0; i < width; i++)
            p[i] = color;
    }
}

static void draw_rect_BGRA(
    guchar *pixels[3],
    guint   stride[3],
    gint    x,
    gint    y,
    guint   width,
    guint   height,
    guint32 color
)
{
    // Converts ARGB color to BGRA
    color = GUINT32_SWAP_LE_BE(color);

    draw_rect_ARGB(pixels, stride, x, y, width, height, color);
}

static void draw_rect_RGBA(
    guchar *pixels[3],
    guint   stride[3],
    gint    x,
    gint    y,
    guint   width,
    guint   height,
    guint32 color
)
{
    // Converts ARGB color to RGBA
    color = ((color >> 24) & 0xff) | ((color & 0xffffff) << 8);

    draw_rect_ARGB(pixels, stride, x, y, width, height, color);
}

static void draw_rect_ABGR(
    guchar *pixels[3],
    guint   stride[3],
    gint    x,
    gint    y,
    guint   width,
    guint   height,
    guint32 color
)
{
    // Converts ARGB color to ABGR
    color = ((color & 0xff00ff00)   |
             ((color >> 16) & 0xff) |
             ((color & 0xff) << 16));

    draw_rect_ARGB(pixels, stride, x, y, width, height, color);
}

static void draw_rect_NV12( // Y, UV planes
    guchar *pixels[3],
    guint   stride[3],
    gint    x,
    gint    y,
    guint   width,
    guint   height,
    guint32 color
)
{
    const guchar Y  = color >> 16;
    const guchar Cb = color >> 8;
    const guchar Cr = color;
    guchar *dst;
    guint i, j;

    dst = pixels[0] + y * stride[0] + x;
    for (j = 0; j < height; j++, dst += stride[0])
        for (i = 0; i < width; i++)
            dst[i] = Y;

    x      /= 2;
    y      /= 2;
    width  /= 2;
    height /= 2;

    dst = pixels[1] + y * stride[1] + x * 2;
    for (j = 0; j < height; j++, dst += stride[1])
        for (i = 0; i < width; i++) {
            dst[2*i + 0] = Cb;
            dst[2*i + 1] = Cr;
        }
}

static void draw_rect_YV12( // Y, V, U planes
    guchar *pixels[3],
    guint   stride[3],
    gint    x,
    gint    y,
    guint   width,
    guint   height,
    guint32 color
)
{
    const guchar Y  = color >> 16;
    const guchar Cb = color >> 8;
    const guchar Cr = color;
    guchar *pY, *pU, *pV;
    guint i, j;

    pY = pixels[0] + y * stride[0] + x;
    for (j = 0; j < height; j++, pY += stride[0])
        for (i = 0; i < width; i++)
            pY[i] = Y;

    x      /= 2;
    y      /= 2;
    width  /= 2;
    height /= 2;

    pV = pixels[1] + y * stride[1] + x;
    pU = pixels[2] + y * stride[2] + x;
    for (j = 0; j < height; j++, pU += stride[1], pV += stride[2])
        for (i = 0; i < width; i++) {
            pU[i] = Cb;
            pV[i] = Cr;
        }
}

static void draw_rect_I420( // Y, U, V planes
    guchar *pixels[3],
    guint   stride[3],
    gint    x,
    gint    y,
    guint   width,
    guint   height,
    guint32 color
)
{
    guchar *new_pixels[3] = { pixels[0], pixels[2], pixels[1] };
    guint   new_stride[3] = { stride[0], stride[2], stride[1] };

    draw_rect_YV12(new_pixels, new_stride, x, y, width, height, color);
}

static void draw_rect_YUV422(guchar *pixels[3], guint stride[3],
    gint x, gint y, guint width, guint height, guint32 color)
{
    guint i, j;

    width /= 2;
    for (j = 0; j < height; j++) {
        guint32 * const p = (guint32 *)
            (pixels[0] + (y + j) * stride[0] + x * 2);
        for (i = 0; i < width; i++)
            p[i] = color;
    }
}

static void draw_rect_YUY2(guchar *pixels[3], guint stride[3],
    gint x, gint y, guint width, guint height, guint32 color)
{
    const guchar Y  = color >> 16;
    const guchar Cb = color >> 8;
    const guchar Cr = color;

    color = (Y << 24) | (Cb << 16) | (Y << 8) | Cr;
    draw_rect_YUV422(pixels, stride, x, y, width, height, GUINT32_TO_BE(color));
}

static void draw_rect_UYVY(guchar *pixels[3], guint stride[3],
    gint x, gint y, guint width, guint height, guint32 color)
{
    const guchar Y  = color >> 16;
    const guchar Cb = color >> 8;
    const guchar Cr = color;

    color = (Cb << 24) | (Y << 16) | (Cr << 8) | Y;
    draw_rect_YUV422(pixels, stride, x, y, width, height, GUINT32_TO_BE(color));
}

static void draw_rect_AYUV(
    guchar *pixels[3],
    guint   stride[3],
    gint    x,
    gint    y,
    guint   width,
    guint   height,
    guint32 color
)
{
    guint i, j;

    color = color | 0xff000000;

    for (j = 0; j < height; j++) {
        guint32 *p = (guint32 *)(pixels[0] + (y + j) * stride[0] + x * 4);
        for (i = 0; i < width; i++)
            p[i] = color;
    }
}

static inline guint32 argb2yuv(guint32 color)
{
    const gint32 r = (color >> 16) & 0xff;
    const gint32 g = (color >>  8) & 0xff;
    const gint32 b = (color      ) & 0xff;

    const guint32 y = (( 306 * r + 601 * g + 116 * b) >> 10);
    const guint32 u = ((-172 * r - 339 * g + 512 * b) >> 10) + 128;
    const guint32 v = (( 512 * r - 428 * g -  83 * b) >> 10) + 128;

    return (y << 16) | (u << 8) | v;
}

gboolean
image_draw_rectangle(
    GstVaapiImage *image,
    gint           x,
    gint           y,
    guint          width,
    guint          height,
    guint32        color,
    guint32        flags
)
{
    const GstVideoFormat      image_format = gst_vaapi_image_get_format(image);
    const guint               image_width  = gst_vaapi_image_get_width(image);
    const guint               image_height = gst_vaapi_image_get_height(image);
    GstVaapiDisplay          *display;
    guchar                   *pixels[3];
    guint                     stride[3];
    DrawRectFunc              draw_rect = NULL;
    guint                     i;

    static const struct {
        GstVideoFormat        format;
        DrawRectFunc          draw_rect;
    }
    map[] = {
#define _(FORMAT) { GST_VIDEO_FORMAT_##FORMAT, draw_rect_##FORMAT }
        _(ARGB),
        _(BGRA),
        _(RGBA),
        _(ABGR),
        _(NV12),
        _(YV12),
        _(I420),
        _(YUY2),
        _(UYVY),
        _(AYUV),
#undef  _
        { 0, }
    };

    for (i = 0; !draw_rect && map[i].format; i++)
        if (map[i].format == image_format)
            draw_rect = map[i].draw_rect;
    if (!draw_rect)
        return FALSE;

    if (x < 0)
        x = 0;
    if (y < 0)
        y = 0;
    if (width > image_width - x)
        width = image_width - x;
    if (height > image_height - y)
        height = image_height - y;

    display = gst_vaapi_object_get_display(GST_VAAPI_OBJECT(image));
    if (!display)
        return FALSE;

    if (!gst_vaapi_image_map(image))
        return FALSE;

    for (i = 0; i < gst_vaapi_image_get_plane_count(image); i++) {
        pixels[i] = gst_vaapi_image_get_plane(image, i);
        stride[i] = gst_vaapi_image_get_pitch(image, i);
        switch (flags) {
        case GST_VAAPI_PICTURE_STRUCTURE_BOTTOM_FIELD:
            pixels[i] += stride[i];
            // fall-through
        case GST_VAAPI_PICTURE_STRUCTURE_TOP_FIELD:
            stride[i] *= 2;
            break;
        }
    }

    if (flags)
        y /= 2, height /= 2;

    if (gst_vaapi_video_format_is_yuv(image_format))
        color = argb2yuv(color);

    draw_rect(pixels, stride, x, y, width, height, color);
    return gst_vaapi_image_unmap(image);
}

gboolean
image_upload(GstVaapiImage *image, GstVaapiSurface *surface)
{
    GstVaapiDisplay    *display;
    GstVideoFormat      format;
    GstVaapiImage      *surface_image;
    GstVaapiSubpicture *subpicture;
    gboolean            success;

    display = gst_vaapi_object_get_display(GST_VAAPI_OBJECT(surface));
    if (!display)
        return FALSE;

    format = gst_vaapi_image_get_format(image);
    if (!format)
        return FALSE;

    if (gst_vaapi_surface_put_image(surface, image))
        return TRUE;

    surface_image = gst_vaapi_surface_derive_image(surface);
    if (surface_image) {
        success = gst_vaapi_image_copy(surface_image, image);
        gst_vaapi_object_unref(surface_image);
        if (success)
            return TRUE;
    }

    g_print("could not upload %s image to surface\n",
            gst_vaapi_video_format_to_string(format));

    if (!gst_vaapi_display_has_subpicture_format(display, format, NULL))
        return FALSE;

    g_print("trying as a subpicture\n");

    subpicture = gst_vaapi_subpicture_new(image, 0);
    if (!subpicture)
        g_error("could not create VA subpicture");

    if (!gst_vaapi_surface_associate_subpicture(surface, subpicture,
                                                NULL, NULL))
        g_error("could not associate subpicture to surface");

    /* The surface holds a reference to the subpicture. This is safe */
    gst_vaapi_object_unref(subpicture);
    return TRUE;
}