/*
 * Copyright © 2024 Matthew Wozniak <me@woz.blue>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED “AS IS” AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
 * REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
 * AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
 * INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
 * LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
 * OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
 * PERFORMANCE OF THIS SOFTWARE.
 */

#include "api.h"
#include "log.h"
#include "render.h"
#include "os.h"

#include <codecapi.h>
#include <mfapi.h>
#include <mfobjects.h>
#include <mfreadwrite.h>
#include <stdlib.h>
#include <string.h>
#include <strmif.h>
#include <winerror.h>

static struct renderctx ctx = { 0 };
static bool is_rendering = false;

#define HR(expr) { \
	HRESULT hr = expr; \
	if (!SUCCEEDED(hr)) \
		bail(#expr " failed!!! (code %lu)", hr); \
}

bool do_frame(struct videomode *this, struct movieinfo *info) {
	is_rendering = true;
	IMFSample *sample = NULL;
	IMFMediaBuffer *imf_buffer = NULL;
	int szpx = 4;
	u32 *data = NULL;

	HR(MFCreateMemoryBuffer(ctx.width * ctx.height * szpx, &imf_buffer));
	HR(imf_buffer->lpVtbl->Lock(imf_buffer, (u8 **)&data, NULL, NULL));
	// THIS IS SLOW!!
	u32 *raw = malloc(ctx.width * ctx.height * szpx);
	this->vt->read_screen_pixels(this, 0, 0, ctx.width, ctx.height, raw, IMAGE_FORMAT_BGRA8888);
	for (int y = ctx.height - 1; y > 0; y--) {
		memcpy(data + ((ctx.height - y) * ctx.width), raw + y * ctx.width,
				szpx * ctx.width);
		/* for (int x = 0; x < ctx.width; x++) {
			data[(ctx.height - y) * ctx.width + x] = raw[y * ctx.width + x];
		} */
	}
	free(raw);
	HR(imf_buffer->lpVtbl->Unlock(imf_buffer));
	HR(imf_buffer->lpVtbl->SetCurrentLength(imf_buffer, ctx.width * ctx.height * szpx));

	HR(MFCreateSample(&sample));
	HR(sample->lpVtbl->AddBuffer(sample, imf_buffer));
	u64 sample_duration = (u64)((1.0 / (double)ctx.fps) * 10000000.0);
	HR(sample->lpVtbl->SetSampleDuration(sample, sample_duration));
	u64 sample_time = sample_duration * info->curframe;
	HR(sample->lpVtbl->SetSampleTime(sample, sample_time));
	HR(ctx.sink_writer->lpVtbl->WriteSample(ctx.sink_writer, ctx.stream_index, sample));

	sample->lpVtbl->Release(sample);
	imf_buffer->lpVtbl->Release(imf_buffer);
	return true;
}

typeof((*videomode)->vt->write_movie_frame) orig_write_movie_frame;
void VIRTUAL hook_write_movie_frame(struct videomode *this,
		struct movieinfo *info) {
	if (!do_frame(this, info)) 
		die("oopsie!");
}


bool do_stop() {
	HR(ctx.sink_writer->lpVtbl->Flush(ctx.sink_writer, ctx.stream_index));
	HR(ctx.sink_writer->lpVtbl->Finalize(ctx.sink_writer));
	HR(MFShutdown());
	return true;
}

static void (VIRTUAL *orig_stop_playback)(struct demoplayer *);
void VIRTUAL hook_stop_playback(struct demoplayer *this) {
	orig_stop_playback(this);

	if (!is_rendering) return;
	is_rendering = false;
	if (!do_stop())
		die("oopsie!");
	// TODO: IPC mode so we don't have to restart the game for every demo
	exit(0);
}

bool render_init(int width, int height, int framerate, bool use_bitrate,
		int quality, const char *output_file, int qvs) {
	bool r = 
		os_mprot((*videomode)->vt, 28 * sizeof(void *), PAGE_EXECUTE_READWRITE);
	if (!r) bail("couldn't mprotect videomode vtable");
	orig_write_movie_frame = (*videomode)->vt->write_movie_frame;
	(*videomode)->vt->write_movie_frame = hook_write_movie_frame;

	r = os_mprot(demoplayer->vt, 18 * sizeof(void *), PAGE_EXECUTE_READWRITE);
	if (!r) bail("couldn't mprotect demoplayer vtable");
	orig_stop_playback = demoplayer->vt->stop_playback;
	demoplayer->vt->stop_playback = hook_stop_playback;

	// make the game thing we are recording a movie (we are!)
	memcpy(movieinfo->name, "a", 2);
	movieinfo->type = 0;

	char framerate_cmd[32] = { 0 };
	sprintf(framerate_cmd, "host_framerate %d;", framerate);
	cbuf_addtext(framerate_cmd);

	ctx.fps = framerate;
	ctx.width = width;
	ctx.height = height;
	ctx.bit_rate = use_bitrate ? quality : 99999999;

	MFStartup(MF_VERSION, 0);

	// init sinkwriter
	{
		IMFMediaType *mt_out = NULL;
		IMFMediaType *mt_in = NULL;

		u16 path[MAX_PATH] = {0};
		info("strlen = %d", strlen(output_file));
		mbstowcs(path, output_file, strlen(output_file));

		info("rendering to %S", path);

		HR(MFCreateSinkWriterFromURL(path, NULL, NULL, &ctx.sink_writer));
		HR(MFCreateMediaType(&mt_out));
		HR(mt_out->lpVtbl->SetGUID(mt_out, &MF_MT_MAJOR_TYPE, &MFMediaType_Video));
		HR(mt_out->lpVtbl->SetGUID(mt_out, &MF_MT_SUBTYPE, &MFVideoFormat_H264));
		HR(mt_out->lpVtbl->SetUINT32(mt_out, &MF_MT_AVG_BITRATE, ctx.bit_rate));
		HR(mt_out->lpVtbl->SetUINT32(mt_out, &MF_MT_MPEG2_PROFILE , eAVEncH264VProfile_Main));
		HR(mt_out->lpVtbl->SetUINT32(mt_out, &MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive));
		HR(mt_out->lpVtbl->SetUINT64(mt_out, &MF_MT_FRAME_SIZE, ((u64)ctx.width << 32) | ctx.height));
		HR(mt_out->lpVtbl->SetUINT64(mt_out, &MF_MT_FRAME_RATE, ((u64)ctx.fps << 32) | 1));
		HR(mt_out->lpVtbl->SetUINT64(mt_out, &MF_MT_PIXEL_ASPECT_RATIO, (((u64)1 << 32) | 1)));
		HR(ctx.sink_writer->lpVtbl->AddStream(ctx.sink_writer, mt_out, &ctx.stream_index));

		HR(MFCreateMediaType(&mt_in));
		HR(mt_in->lpVtbl->SetGUID(mt_in, &MF_MT_MAJOR_TYPE, &MFMediaType_Video));
		HR(mt_in->lpVtbl->SetGUID(mt_in, &MF_MT_SUBTYPE, &MFVideoFormat_ARGB32));
		HR(mt_in->lpVtbl->SetUINT32(mt_in, &MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive));
		HR(mt_in->lpVtbl->SetUINT64(mt_in, &MF_MT_FRAME_SIZE, ((u64)ctx.width << 32) | ctx.height));
		HR(mt_in->lpVtbl->SetUINT64(mt_in, &MF_MT_FRAME_RATE, ((u64)ctx.fps << 32) | 1));
		HR(mt_in->lpVtbl->SetUINT64(mt_in, &MF_MT_PIXEL_ASPECT_RATIO, (((u64)1 << 32) | 1)));
		HR(ctx.sink_writer->lpVtbl->SetInputMediaType(ctx.sink_writer, ctx.stream_index, mt_in, NULL));

		HR(ctx.sink_writer->lpVtbl->BeginWriting(ctx.sink_writer));
		HR(ctx.sink_writer->lpVtbl->AddRef(ctx.sink_writer));
	}

	// video encoding parameters
	{
		ICodecAPI *codec;
		HR(ctx.sink_writer->lpVtbl->GetServiceForStream(ctx.sink_writer, 0, &GUID_NULL, &IID_ICodecAPI, (void*)&codec));

		// use a quality VBR
		int ratecontrol = use_bitrate ? eAVEncCommonRateControlMode_UnconstrainedVBR : eAVEncCommonRateControlMode_Quality;
		VARIANT _ratecontrol = { .vt = VT_UI4, .ulVal = ratecontrol };
		HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncCommonRateControlMode, &_ratecontrol));

		// set the quality
		VARIANT _quality = { .vt = VT_UI4, .ulVal = quality };
		HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncCommonQuality, &_quality));

		VARIANT _qvs = { .vt = VT_UI4, .ulVal = qvs };
		HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncCommonQualityVsSpeed, &_qvs));

		// set the gop size to 2 seconds
		VARIANT gop = { .vt = VT_UI4, .ulVal = 2 * ctx.fps };
		HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncMPVGOPSize, &gop));

		// enable 2 b-frames
		VARIANT bframes = { .vt = VT_UI4, .ulVal = 2 };
		HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncMPVDefaultBPictureCount, &bframes));

		codec->lpVtbl->Release(codec);
	}

	return true;
}

// vi: sw=4 ts=4 noet tw=80 cc=80