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

#include <ftw.h>
#include <sys/stat.h>

#include <libbluray/bluray.h>
#include <libaacs/aacs.h>

#ifndef BUFFER_SIZE
#define BUFFER_SIZE 1024
#endif

const char* aacs_error_string(int error_code)
{
	switch(error_code)
	{
		case AACS_ERROR_CORRUPTED_DISC:
			return "disc corrupted";
			break;
		case AACS_ERROR_NO_CONFIG:
			return "AACS config file not found";
			break;
		case AACS_ERROR_NO_PK:
			return "no processing key found for this disc";
			break;
		case AACS_ERROR_NO_CERT:
			return "no valid AACS certificate found";
			break;
		case AACS_ERROR_CERT_REVOKED:
			return "AACS certificate revoked";
			break;
		case AACS_ERROR_MMC_OPEN:
			return "MMC open failed";
			break;
		case AACS_ERROR_MMC_FAILURE:
			return "MMC failed";
			break;
		case AACS_ERROR_NO_DK:
			return "no matching device key";
			break;
		default:
			return "unknown AACS error";
			break;
	}
}

const char* disc_path;
const char* output_path;

BLURAY* bd;
AACS* aacs;

int fn(const char* fpath, const struct stat* sb, int typeflag)
{
	const char* path = fpath + strlen(disc_path);
	if(!strncmp(path, "AACS", 4))
	{
		return 0;
	}

	char output_entry[255];
	sprintf(output_entry, "%s/%s", output_path, path);

	if(typeflag == FTW_D)
	{
		mkdir(output_entry, S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
	}
	else
	{
		const char* extension = ".m2ts";
		if(strlen(path) >= strlen(extension) && !strncmp(path + strlen(path) - strlen(extension), extension, strlen(extension)))
		{
			void* data = NULL;
			int64_t size;
			if(!bd_read_file(bd, path, &data, &size))
			{
				fprintf(stderr, "Error reading file from Blu-Ray!]\n");
				return EXIT_FAILURE;
			}

			FILE* file = fopen(output_entry, "wb");
			uint8_t unit[6144];
			size_t unit_count = (size + sizeof(unit) - 1) / sizeof(unit);
			for(int i = 0; i < unit_count; i++)
			{
				size_t unit_size = ((i + 1) * sizeof(unit) <= size) ? sizeof(unit) : size - i * sizeof(unit);
				memcpy(unit, data + i * sizeof(unit), unit_size);
				aacs_decrypt_unit(aacs, data + i * sizeof(unit));
				fwrite(data + i * sizeof(unit), unit_size, 1, file);
			}
			fclose(file);
			free(data);
		}
		else
		{
			//TODO: Error checking
			uint8_t buffer[BUFFER_SIZE];
			FILE* from = fopen(fpath, "rb");
			FILE* to = fopen(output_entry, "wb");
			size_t bytes_read = 0;
			while((bytes_read = fread(buffer, 1, BUFFER_SIZE, from)) > 0)
			{
				fwrite(buffer, 1, bytes_read, to);
			}
			fclose(to);
			fclose(from);
		}
	}

	return 0;
}

int main(int argc, char* argv[])
{
	if(argc < 3)
	{
		fprintf(stderr, "usage: %s <BD root> <output path> [key file]\n", argv[0]);
		return EXIT_FAILURE;
	}

	disc_path = argv[1];
	output_path = argv[2];
	const char* keyfile_path = argc >= 4 ? argv[3] : NULL;

	bd = bd_open(disc_path, keyfile_path);

	if(!bd)
	{
		fprintf(stderr, "Error opening %s!\n", disc_path);
		return EXIT_FAILURE;
	}

	const BLURAY_DISC_INFO* disc_info = bd_get_disc_info(bd);

	if(!disc_info)
	{
		fprintf(stderr, "Error getting disc info!\n");
		return EXIT_FAILURE;
	}

	if(!disc_info->aacs_detected)
	{
		fprintf(stderr, "Disc is not encrypted with AACS.\n");
		return EXIT_SUCCESS;
	}

	if(!disc_info->libaacs_detected)
	{
		fprintf(stderr, "libaacs is not available!\n");
		return EXIT_FAILURE;
	}

	int error_code;
	aacs = aacs_open2(disc_path, keyfile_path, &error_code);
	if(error_code != AACS_SUCCESS)
	{
		fprintf(stderr, "AACS handling failed: %s.\n", aacs_error_string(error_code));
		return EXIT_FAILURE;
	}

	aacs_select_title(aacs, 0xffff);

	ftw(disc_path, &fn, 16);

	aacs_close(aacs);
	bd_close(bd);

	return EXIT_SUCCESS;
}