/*
 * 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 "hook.h"
#include "render.h"
#include "os.h"

#include <combaseapi.h>
#include <mfapi.h>
#include <mfidl.h>
#include <mfreadwrite.h>
#include <codecapi.h>
#include <mftransform.h>
#include <stdlib.h>
#include <string.h>
#include <strmif.h>
#include <winerror.h>

extern struct {
	int width;
	int height;
	const char *game;
	int fps;
	int quality;
	int bitrate;
	int qvs;
	const char *out;
	const char **demo;
	int demo_count;
	bool combine;
} args;

struct renderctx {
	GUID enc_format;
	GUID in_fmt;
	IMFSinkWriter *sink_writer;
	ulong stream_index;
	ulong audio_index;
	int cur_demo;
};

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;
	const int szpx = 3;

	static struct { u8 bgr[3]; } *raw = 0;
	if (!raw) raw = malloc(args.width * args.height * szpx);

	IMFMediaBuffer *imf_buffer = 0;
	HR(MFCreateMemoryBuffer(args.width * args.height * szpx, &imf_buffer));
	struct { u8 bgr[3]; } *data = 0;
	HR(imf_buffer->lpVtbl->Lock(imf_buffer, (u8 **)&data, NULL, NULL));
	// THIS IS SLOW!!
	if (steampipe)
		this->vt->steampipe_read_screen_pixels(this, 0, 0, args.width,
			args.height, raw, IMAGE_FORMAT_BGR888);
	else
		this->vt->read_screen_pixels(this, 0, 0, args.width, args.height, raw,
			IMAGE_FORMAT_BGR888);

	for (int y = args.height - 1; y > 0; y--) {
		memcpy(data + ((args.height - y) * args.width), raw + y * args.width,
				szpx * args.width);
	}
	HR(imf_buffer->lpVtbl->Unlock(imf_buffer));
	HR(imf_buffer->lpVtbl->SetCurrentLength(imf_buffer, args.width * args.height * szpx));

	IMFSample *sample = 0;
	HR(MFCreateSample(&sample));
	HR(sample->lpVtbl->AddBuffer(sample, imf_buffer));
	u64 sample_duration = (u64)((1.0 / (double)args.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));
	HR(sample->lpVtbl->Release(sample));
	HR(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_audio_frame(void) {
	// TODO: audio
#define CLIP16(n) (i16)(n < -32768 ? -32768 : (n > 32767 ? 32767 : n))
#undef CLIP16
	return true;
}

void (*orig_snd_recordbuffer)(void) = 0;
void hook_snd_recordbuffer(void) {
	if (!is_rendering) {
		orig_snd_recordbuffer();
		return;
	}
	do_audio_frame();
}

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;

	// do we have any more demos to play?
	if (args.demo_count--) {
		demoplayer->vt->start_playback(demoplayer, *++args.demo, false);
		return;
	}

	info("finished!");

	if (!do_stop())
		die("oopsie!");

	exit(0);
}

bool render_init(void) {
	bool r = 
		os_mprot((*videomode)->vt, 28 * sizeof(void *), PAGE_EXECUTE_READWRITE);
	if (!r) bail("couldn't mprotect videomode vtable");

	typeof(orig_write_movie_frame) *wmf = steampipe ? 
		&(*videomode)->vt->steampipe_write_movie_frame :
		&(*videomode)->vt->write_movie_frame;
	orig_write_movie_frame = *wmf;
	*wmf = 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;

	// hook the audio frame
	orig_snd_recordbuffer = (void (*)(void))
		hook_inline((void *)snd_recordbuffer, (void *)hook_snd_recordbuffer);

	// 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;", args.fps);
	cbuf_addtext(framerate_cmd);

	if (CoInitializeEx(NULL, COINIT_MULTITHREADED) == RPC_E_CHANGED_MODE) 
		warn("changed COM concurrency mode!");

	HR(MFStartup(MF_VERSION, 0));

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

		u16 path[MAX_PATH] = {0};
		mbstowcs(path, args.out, strlen(args.out));

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

		IMFAttributes *sinkwriter_attr;
		HR(MFCreateAttributes(&sinkwriter_attr, 1));
		HR(sinkwriter_attr->lpVtbl->SetUINT32(sinkwriter_attr, &MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, TRUE));
		HR(MFCreateSinkWriterFromURL(path, NULL, sinkwriter_attr, &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, args.bitrate ? args.bitrate : 9999999));
		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)args.width << 32) | args.height));
		HR(mt_out->lpVtbl->SetUINT64(mt_out, &MF_MT_FRAME_RATE, ((u64)args.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_RGB24));
		HR(mt_in->lpVtbl->SetUINT32(mt_in, &MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive));
		HR(mt_in->lpVtbl->SetUINT32(mt_in, &MF_MT_ALL_SAMPLES_INDEPENDENT, TRUE));
		HR(mt_in->lpVtbl->SetUINT64(mt_in, &MF_MT_FRAME_SIZE, ((u64)args.width << 32) | args.height));
		HR(mt_in->lpVtbl->SetUINT64(mt_in, &MF_MT_FRAME_RATE, ((u64)args.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));
	}

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

		int ratecontrol = !!args.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 = !!args.bitrate ? args.bitrate : args.quality
		};
		HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncCommonQuality, &_quality));

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

		// set the gop size to 2 seconds
		VARIANT gop = { .vt = VT_UI4, .ulVal = 2 * args.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);
	}

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

	return true;
}

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