/*
 * Support functions for the IMAGE datatype.
 *
 * Written by Andrew Tipton <andrew@kiwidrew.com>
 *
 * pg_image/pg_image.c
 */

#include <string.h>


#include "postgres.h"
#include "fmgr.h"
#include "access/hash.h"
#include "utils/builtins.h"


#include "gd.h"


#include "pg_image.h"
#include "base64.h"


/* PG_MODULE_MAGIC must appear exactly once in the entire module. */
PG_MODULE_MAGIC;


static int
mimetype_to_imgtype(char *mimetype) {
    if (strcmp(mimetype, "image/png") == 0) {
        return IMAGE_PNG;
    } else if (strcmp(mimetype, "image/jpeg") == 0) {
        return IMAGE_JPG;
    } else {
        return IMAGE_UNKNOWN;
    }
}


static char *
imgtype_to_mimetype(int imgtype) {
    switch (imgtype) {
        case IMAGE_PNG:  return "image/png";
        case IMAGE_JPG:  return "image/jpeg";
        default:         return "unknown";
    }
}


static int
add_metadata(struct pg_image *img)
{
    gdImage *gd_img = NULL;

    /* Load the image metadata and initialize the pg_image datum. */
    switch (img->imgtype) {
        case IMAGE_JPG:  gd_img = gdImageCreateFromJpegPtr(img->datasz, img->data); break;
        case IMAGE_PNG:  gd_img = gdImageCreateFromPngPtr(img->datasz, img->data); break;
        default:         elog(ERROR, "pg_image_from_data: unknown imgtype %d", img->imgtype); return 0;
    }
    if (!gd_img) {
        return 0;
    }
    img->width = gd_img->sx;
    img->height = gd_img->sy;
    return 1;
}


static struct pg_image *
pg_image_decode(const char *input)
{
    int input_length = strlen(input);
    const char *cur_input = input;
    struct pg_image *img;
    uint8 imgtype;

    char uri_mimetype[64], uri_encoding[64];
    int uri_mimetype_len, uri_encoding_len;
    int recognized;

    /*
     * Input URI must have the "data:" scheme.
     */
    if (strncmp(input, "data:", 5) != 0) {
        ereport(ERROR, (errmsg("image input must be a \"data:\"-style URI")));
    }
    cur_input += 5;

    /* 
     * Mimetype immediately follows, and is terminated by a ";" or ","
     * character.
     */
    uri_mimetype_len = strcspn(cur_input, ";,");
    if (uri_mimetype_len > 63) {
        ereport(ERROR, (errmsg("image input has a mimetype longer than 63 bytes")));
    }
    strncpy(uri_mimetype, cur_input, uri_mimetype_len);
    uri_mimetype[uri_mimetype_len] = '\0';
    imgtype = mimetype_to_imgtype(uri_mimetype);
    if (imgtype == IMAGE_UNKNOWN) {
        ereport(ERROR, (errmsg("image input has an unrecognized mimetype"),
                        errdetail("The input has a mimetype of \"%s\".", uri_mimetype)));
        return NULL;    /* Keep the compiler happy. */
    }
    cur_input += uri_mimetype_len;

    /*
     * The data encoding, if different from the default "urlescaped", is
     * specified by a ";" character followed by the encoding name.  We only
     * support the "base64" encoding.
     */
    if (*cur_input != ';') {
        strcpy(uri_encoding, "urlescaped");
        uri_encoding_len = strlen(uri_encoding);
    } else {
        cur_input += 1;
        uri_encoding_len = strcspn(cur_input, ",");
        if (uri_encoding_len > 63) {
            ereport(ERROR, (errmsg("image input has an encoding name longer than 63 bytes")));
        }
        strncpy(uri_encoding, cur_input, uri_encoding_len);
        uri_encoding[uri_encoding_len] = '\0';
        cur_input += uri_encoding_len;
    }
    if (strcmp(uri_encoding, "base64") != 0) {
        ereport(ERROR, (errmsg("image input has an unsupported encoding"),
                        errdetail("The input has an encoding of \"%s\".", uri_encoding)));
    }


    /*
     * Decode the binary data as base64.  We have already made one pass through
     * the data to estimate the buffer size, followed by the decoding pass.  The
     * actual size may be smaller due to whitespace padding in the base64 input.
     */
    if (*cur_input != ',') {
        ereport(ERROR, (errmsg("image input does not contain any data")));
    }
    cur_input += 1;
    input_length -= (cur_input - input);
    img = palloc(sizeof(*img) + b64_decode_len(input_length));
    img->datasz = b64_decode(cur_input, input_length, img->data);
    img->imgtype = imgtype;
    SET_VARSIZE(img, sizeof(struct pg_image) + img->datasz);

    /*
     * Validate the image format and fill in the metadata.
     */
    recognized = add_metadata(img);
    if (!recognized) {
        ereport(ERROR, (errmsg("image input cannot be handled"), errhint("Is the input a valid JPEG or PNG image?")));
    }
    return img;
}


static char *
pg_image_encode(struct pg_image *img)
{
    const char *prefix = "data:";
    const char *suffix = ";base64,\n"; 
    char *result, *ptr;
    int headersz;

    /* Determine the image's mimetype string. */
    char *mimetype = imgtype_to_mimetype(img->imgtype);
    if (!mimetype) {
        ereport(ERROR, (errmsg("unknown imgtype %d in image datum", img->imgtype)));
    }

    /* Calculate the output length and allocate a buffer. */
    headersz = strlen(prefix) + strlen(mimetype) + strlen(suffix);
    result = palloc(headersz + b64_encode_len(img->datasz));
    
    /* Generate the data: URI and return it. */
    ptr = result;
    ptr += sprintf(ptr, "%s%s%s", prefix, mimetype, suffix);
    ptr += b64_encode(img->data, img->datasz, ptr);
    *ptr = '\0';
    return result;
}


PG_FUNCTION_INFO_V1(pg_image_in);
Datum
pg_image_in(PG_FUNCTION_ARGS)
{
    char *input = PG_GETARG_CSTRING(0);
    struct pg_image *img;
    
    img = pg_image_decode(input);
    if (img == NULL) {
        char detail[128];
        if (strlen(input) < 64) {
            sprintf(detail, "\"%63s\"", input);
        } else {
            sprintf(detail, "\"%63s\" (...truncated)", input); 
        }
        ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION),
                        errmsg("invalid pg_image input"),
                        errdetail("input was: %s", detail)
                       ));
    }

    PG_RETURN_IMAGE_P(img);
}


PG_FUNCTION_INFO_V1(pg_image_out);
Datum
pg_image_out(PG_FUNCTION_ARGS)
{
    struct pg_image *img = PG_GETARG_IMAGE_P(0);

    const char *result = pg_image_encode(img);

    PG_RETURN_CSTRING(result);
}


PG_FUNCTION_INFO_V1(pg_image_in_text);
Datum
pg_image_in_text(PG_FUNCTION_ARGS)
{
    text *input = PG_GETARG_TEXT_PP(0);
    char *str = text_to_cstring(input);
    struct pg_image *img = NULL;

    img = pg_image_decode(str);
    if (img == NULL) {
        ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION),
                        errmsg("invalid pg_image input"),
                        errdetail("input was: \"%63s\"%s", str,
                                    strlen(str) <= 63 ? "" : " (...truncated)")
                       ));
    }

    PG_FREE_IF_COPY(input, 0);
    PG_RETURN_IMAGE_P(img);
}


PG_FUNCTION_INFO_V1(pg_image_out_text);
Datum
pg_image_out_text(PG_FUNCTION_ARGS)
{
    struct pg_image *img = PG_GETARG_IMAGE_P(0);

    const char *buf = pg_image_encode(img);

    text *result = cstring_to_text(buf);
    PG_RETURN_TEXT_P(result);
}


PG_FUNCTION_INFO_V1(pg_image_from_bytea);
Datum
pg_image_from_bytea(PG_FUNCTION_ARGS)
{
    bytea *input = PG_GETARG_BYTEA_P(0);
    text *mimetype_text = PG_GETARG_TEXT_P(1);
    char *mimetype = text_to_cstring(mimetype_text);
    int input_len = VARSIZE(input) - VARHDRSZ;
    struct pg_image *img;
    int recognized;

    img = palloc(sizeof(struct pg_image) + input_len);
    memcpy(img->data, VARDATA(input), input_len);
    img->datasz = input_len;
    img->imgtype = mimetype_to_imgtype(mimetype);
    if (img->imgtype == IMAGE_UNKNOWN) {
        ereport(ERROR, (errmsg("image input has an unrecognized mimetype"),
                        errdetail("The input has a mimetype of \"%s\".", mimetype)));
    }
    SET_VARSIZE(img, sizeof(struct pg_image) + img->datasz);

    recognized = add_metadata(img);
    if (!recognized) {
        ereport(ERROR, (errmsg("image input cannot be handled"), errhint("Is the input a valid JPEG or PNG image?")));
    }

    PG_FREE_IF_COPY(input, 0);
    PG_RETURN_IMAGE_P(img);
}


PG_FUNCTION_INFO_V1(pg_image_width);
Datum
pg_image_width(PG_FUNCTION_ARGS)
{
    struct pg_image *img = PG_GETARG_IMAGE_P(0);

    PG_RETURN_INT32(img->width);
}


PG_FUNCTION_INFO_V1(pg_image_height);
Datum
pg_image_height(PG_FUNCTION_ARGS)
{
    struct pg_image *img = PG_GETARG_IMAGE_P(0);

    PG_RETURN_INT32(img->height);
}


PG_FUNCTION_INFO_V1(pg_image_filesize);
Datum
pg_image_filesize(PG_FUNCTION_ARGS)
{
    struct pg_image *img = PG_GETARG_IMAGE_P(0);

    PG_RETURN_INT32(img->datasz);
}


PG_FUNCTION_INFO_V1(pg_image_mimetype);
Datum
pg_image_mimetype(PG_FUNCTION_ARGS)
{
    struct pg_image *img = PG_GETARG_IMAGE_P(0);
    char *mimetype = "unknown";
    text *result;

    switch (img->imgtype) {
        case IMAGE_PNG:  mimetype = "image/png";  break;
        case IMAGE_JPG:  mimetype = "image/jpeg"; break;
    }

    result = cstring_to_text(mimetype);
    PG_RETURN_TEXT_P(result);
}


PG_FUNCTION_INFO_V1(pg_image_imagedata);
Datum
pg_image_imagedata(PG_FUNCTION_ARGS)
{
    struct pg_image *img = PG_GETARG_IMAGE_P(0);

    char *data = palloc(VARHDRSZ + img->datasz);
    SET_VARSIZE(data, VARHDRSZ + img->datasz);
    memcpy(data + VARHDRSZ, img->data, img->datasz);

    PG_RETURN_BYTEA_P(data);
}


PG_FUNCTION_INFO_V1(pg_image_eq);
Datum
pg_image_eq(PG_FUNCTION_ARGS)
{
    struct pg_image *img1 = PG_GETARG_IMAGE_P(0);
    struct pg_image *img2 = PG_GETARG_IMAGE_P(1);
    
    if (img1->imgtype != img2->imgtype) {
        PG_RETURN_BOOL(0);
    }
    if (img1->datasz != img2->datasz) {
        PG_RETURN_BOOL(0);
    }
    if (memcmp(img1->data, img2->data, img1->datasz) != 0) {
        PG_RETURN_BOOL(0);
    }
    PG_RETURN_BOOL(1);
}


PG_FUNCTION_INFO_V1(pg_image_ne);
Datum
pg_image_ne(PG_FUNCTION_ARGS)
{
    struct pg_image *img1 = PG_GETARG_IMAGE_P(0);
    struct pg_image *img2 = PG_GETARG_IMAGE_P(1);
    
    if (img1->imgtype != img2->imgtype) {
        PG_RETURN_BOOL(1);
    }
    if (img1->datasz != img2->datasz) {
        PG_RETURN_BOOL(1);
    }
    if (memcmp(img1->data, img2->data, img1->datasz) != 0) {
        PG_RETURN_BOOL(1);
    }
    PG_RETURN_BOOL(0);
}


PG_FUNCTION_INFO_V1(pg_image_hash);
Datum
pg_image_hash(PG_FUNCTION_ARGS)
{
    struct pg_image *img = PG_GETARG_IMAGE_P(0);
    int32 hash;

    hash = hash_any((uint8*)img->data, img->datasz);

    PG_RETURN_INT32(hash);
}