/* See LICENSE file for copyright and license details. */

#include <stdint.h>
#include <stdio.h>
#include <string.h>

#include <libavutil/frame.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
#include <png.h>
#include <jpeglib.h>

#include "diff.h"
#include "util.h"

/*
 * Compute 64-bit average perceptual hash.
 *
 * 1. Scale frame to 8x8 grayscale.
 * 2. Compute mean pixel value.
 * 3. Each bit = 1 if pixel > mean, 0 otherwise.
 *
 * Fast and effective for detecting duplicate/near-duplicate frames.
 */
uint64_t
diff_phash(AVFrame *frame, struct SwsContext *sws,
           int src_w, int src_h)
{
	uint8_t gray[64];
	uint8_t *dst_data[1];
	int dst_linesize[1];
	uint64_t hash;
	unsigned sum;
	uint8_t mean;
	int i;

	dst_data[0] = gray;
	dst_linesize[0] = 8;

	sws_scale(sws,
	          (const uint8_t *const *)frame->data,
	          frame->linesize,
	          0, src_h,
	          dst_data, dst_linesize);

	/* compute mean */
	sum = 0;
	for (i = 0; i < 64; i++)
		sum += gray[i];
	mean = (uint8_t)(sum / 64);

	/* build hash */
	hash = 0;
	for (i = 0; i < 64; i++) {
		if (gray[i] > mean)
			hash |= ((uint64_t)1 << i);
	}

	return hash;
}

int
diff_hamming(uint64_t a, uint64_t b)
{
	uint64_t x;
	int count;

	x = a ^ b;
	count = 0;
	while (x) {
		count++;
		x &= x - 1;
	}
	return count;
}

static int
save_png(const uint8_t *rgb, int w, int h, int stride,
         const char *path)
{
	FILE *fp;
	png_structp png;
	png_infop info;
	int y;

	fp = fopen(path, "wb");
	if (!fp) {
		warn("cannot open '%s' for writing:", path);
		return -1;
	}

	png = png_create_write_struct(PNG_LIBPNG_VER_STRING,
	                              NULL, NULL, NULL);
	if (!png) {
		fclose(fp);
		return -1;
	}

	info = png_create_info_struct(png);
	if (!info) {
		png_destroy_write_struct(&png, NULL);
		fclose(fp);
		return -1;
	}

	if (setjmp(png_jmpbuf(png))) {
		png_destroy_write_struct(&png, &info);
		fclose(fp);
		return -1;
	}

	png_init_io(png, fp);
	png_set_IHDR(png, info, w, h, 8,
	             PNG_COLOR_TYPE_RGB,
	             PNG_INTERLACE_NONE,
	             PNG_COMPRESSION_TYPE_DEFAULT,
	             PNG_FILTER_TYPE_DEFAULT);
	png_write_info(png, info);

	for (y = 0; y < h; y++)
		png_write_row(png, rgb + y * stride);

	png_write_end(png, NULL);
	png_destroy_write_struct(&png, &info);
	fclose(fp);
	return 0;
}

static int
save_jpg(const uint8_t *rgb, int w, int h, int stride,
         const char *path)
{
	FILE *fp;
	struct jpeg_compress_struct cinfo;
	struct jpeg_error_mgr jerr;
	JSAMPROW row;
	int y;

	fp = fopen(path, "wb");
	if (!fp) {
		warn("cannot open '%s' for writing:", path);
		return -1;
	}

	cinfo.err = jpeg_std_error(&jerr);
	jpeg_create_compress(&cinfo);
	jpeg_stdio_dest(&cinfo, fp);

	cinfo.image_width = w;
	cinfo.image_height = h;
	cinfo.input_components = 3;
	cinfo.in_color_space = JCS_RGB;
	jpeg_set_defaults(&cinfo);
	jpeg_set_quality(&cinfo, 95, 1);
	jpeg_start_compress(&cinfo, 1);

	for (y = 0; y < h; y++) {
		row = (JSAMPROW)(rgb + y * stride);
		jpeg_write_scanlines(&cinfo, &row, 1);
	}

	jpeg_finish_compress(&cinfo);
	jpeg_destroy_compress(&cinfo);
	fclose(fp);
	return 0;
}

int
diff_save_frame(AVFrame *frame, struct SwsContext *sws_rgb,
                int src_w, int src_h,
                const char *path, const char *fmt)
{
	AVFrame *rgb_frame;
	int ret;

	rgb_frame = av_frame_alloc();
	if (!rgb_frame)
		return -1;

	rgb_frame->format = AV_PIX_FMT_RGB24;
	rgb_frame->width = src_w;
	rgb_frame->height = src_h;

	ret = av_image_alloc(rgb_frame->data,
	                     rgb_frame->linesize,
	                     src_w, src_h,
	                     AV_PIX_FMT_RGB24, 32);
	if (ret < 0) {
		av_frame_free(&rgb_frame);
		return -1;
	}

	sws_scale(sws_rgb,
	          (const uint8_t *const *)frame->data,
	          frame->linesize,
	          0, src_h,
	          rgb_frame->data, rgb_frame->linesize);

	if (strcmp(fmt, "jpg") == 0 || strcmp(fmt, "jpeg") == 0) {
		ret = save_jpg(rgb_frame->data[0],
		               src_w, src_h,
		               rgb_frame->linesize[0], path);
	} else {
		ret = save_png(rgb_frame->data[0],
		               src_w, src_h,
		               rgb_frame->linesize[0], path);
	}

	av_freep(&rgb_frame->data[0]);
	av_frame_free(&rgb_frame);
	return ret;
}