/*
	Copyright (C) 2025  Chud

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program 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 General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

/*
gcc -Wall -Wextra -Werror -Wstrict-prototypes
-pedantic -pedantic-errors -fstack-protector-all
-fsanitize=address -std=c99 -Ofast -pipe -lm
file2png.c -o file2png
*/

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <math.h>
#define STB_IMAGE_IMPLEMENTATION
#include <stb/stb_image.h>
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <stb/stb_image_write.h>

int file_exists(const char *filename) {
	FILE *f = fopen(filename, "rb");
	if (!f) return 0;
	fclose(f);
	return 1;
}

uint8_t *read_file(const char *filename, size_t *out_size) {
	FILE *f = fopen(filename, "rb");
	if (!f) return NULL;
	fseek(f, 0, SEEK_END);
	*out_size = ftell(f);
	rewind(f);

	uint8_t *data = malloc(*out_size);
	fread(data, 1, *out_size, f);
	fclose(f);
	return data;
}

int main(int argc, char *argv[]) {
	if (argc < 4) {
		printf("Usage: %s [data file] [overlay png] [output png]\n"
	"Example: %s input.mp4 input.png output.png\n", argv[0], argv[0]);
		return 1;
	}
	if (!file_exists(argv[1])) {
		printf("%s doesn't exist.\n", argv[1]);
		return 1;
	}
	if (!file_exists(argv[2])) {
		printf("%s doesn't exist.\n", argv[2]);
		return 1;
	}
	size_t data_size;
	uint8_t *data = read_file(argv[1], &data_size);
	if (!data) {
		fprintf(stderr, "Failed to read data file.\n");
		return 1;
	}
	int overlay_w, overlay_h, overlay_channels;
	uint8_t *overlay = stbi_load(argv[2], &overlay_w, &overlay_h, &overlay_channels, 1);
	if (!overlay) {
		fprintf(stderr, "Failed to load overlay image.\n");
		free(data);
		return 1;
	}

	// calculate pixel count and res
	int pixel_count = (int)ceil(data_size / 3.0);
	float pixel_root = sqrt((float)pixel_count);
	float aspect = (float)overlay_w / overlay_h;
	float aspect_root = sqrtf(aspect);
	int width = (int)ceil(pixel_root * aspect_root);
	int height = (int)ceil(pixel_root / aspect_root);
	int padded_size = width * height * 3;
	uint8_t *rgb_data = calloc(padded_size, 1);
	memcpy(rgb_data, data, data_size);
	free(data);

	// resize overlay png to match res
	uint8_t *resized_overlay = malloc(width * height);
	for (int y = 0; y < height; ++y) {
		int src_y = y * overlay_h / height;
		for (int x = 0; x < width; ++x) {
			int src_x = x * overlay_w / width;
			resized_overlay[y * width + x] = overlay[src_y * overlay_w + src_x];
		}
	}
	free(overlay);

	for (int i = 0; i < width * height; ++i) {
		if (resized_overlay[i] == 0)
			resized_overlay[i] = 1;
	}

	// RGB+A
	uint8_t *rgba_data = malloc(width * height * 4);
	for (int i = 0; i < width * height; ++i) {
		rgba_data[4 * i + 0] = rgb_data[3 * i + 0];
		rgba_data[4 * i + 1] = rgb_data[3 * i + 1];
		rgba_data[4 * i + 2] = rgb_data[3 * i + 2];
		rgba_data[4 * i + 3] = resized_overlay[i];
	}
	free(rgb_data);
	free(resized_overlay);

	if (!stbi_write_png(argv[3], width, height, 4, rgba_data, width * 4)) {
		fprintf(stderr, "Failed to write PNG.\n");
		free(rgba_data);
		return 1;
	}
	free(rgba_data);

	printf("ffmpeg -i %s -pix_fmt rgb24 -f rawvideo - | mpv -\n", argv[3]);
	return 0;
}