/*
 * 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 <libavformat/avio.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavcodec/codec_id.h>
#include <libavutil/channel_layout.h>
#include <libavutil/mathematics.h>
#include <libavutil/error.h>
#include <libavutil/opt.h>
#include <libavutil/log.h>
#include <libavutil/rational.h>
#include <libavutil/imgutils.h>
#include <libavcodec/codec.h>
#include <libavcodec/packet.h>
#include <libavutil/frame.h>
#include <libavutil/avutil.h>
#include <libavutil/samplefmt.h>
#include <libavcodec/codec_par.h>
#include <libavutil/log.h>
#include <libswscale/swscale.h> 
#include <libswresample/swresample.h>

extern struct {
	int width;
	int height;
	const char *game;
	int fps;
	int quality;
	int bitrate;
	const char *out;
	const char **demo;
	int demo_count;
	const char *encoder;
	const char *preset;
} args;

struct renderctx {
	AVFormatContext *format;
	AVStream *vstream;
	const AVCodec *vcodec;
	AVCodecContext *vcodec_ctx;
	AVFrame *yuv_frame;
	struct SwsContext *sws;

	AVStream *astream;
	const AVCodec *acodec;
	AVCodecContext *acodec_ctx;
	AVFrame *aframe;
	AVFrame *aframe_tmp;
	struct SwrContext *swr;
};

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

bool encode(AVCodecContext *codec, AVFrame *frame, AVStream *stream) {
	int ret;
	AVPacket *pkt = av_packet_alloc();
 
    if ((ret = avcodec_send_frame(codec, frame)) < 0)
		bail("error sending a frame for encoding");

    while (ret >= 0) {
        ret = avcodec_receive_packet(codec, pkt);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            break;
        else if (ret < 0)
            bail("encoding error!");

		pkt->stream_index = stream->index;

		av_packet_rescale_ts(pkt, codec->time_base, stream->time_base);

		if (av_interleaved_write_frame(ctx.format, pkt) < 0)
			bail("failed to write frame to file!");
        av_packet_unref(pkt);
    }
	av_packet_free(&pkt);
	return true;
}

bool do_frame(struct videomode *this, struct movieinfo *info) {
	(void)(info); // unused param
	is_rendering = true;

	static int nextpts = 0;
	static struct { u8 bgr[3]; } *pixels = 0;
	if (!pixels)
		pixels = malloc(args.width * args.height * sizeof(*pixels));

	if (steampipe)
		this->vt->steampipe_read_screen_pixels(this, 0, 0, args.width,
			args.height, pixels, IMAGE_FORMAT_BGR888);
	else
		this->vt->read_screen_pixels(this, 0, 0, args.width, args.height, pixels,
			IMAGE_FORMAT_BGR888);

	int stride = args.width * sizeof(*pixels);
	sws_scale(ctx.sws, (const u8 **)&pixels, &stride, 0, args.height, 
			ctx.yuv_frame->data, ctx.yuv_frame->linesize);

	ctx.yuv_frame->pts = nextpts++;
	ctx.yuv_frame->time_base = ctx.vcodec_ctx->time_base;
	ctx.yuv_frame->duration = 1;

	encode(ctx.vcodec_ctx, ctx.yuv_frame, ctx.vstream);

	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)) warn("do_frame failed!");
}

bool do_audio_frame(void) {
	if (!is_rendering) return true;
	static u64 nextpts = 0;
	static int frame_idx = 0;
#define CLIP16(n) (i16)(n < -32768 ? -32768 : (n > 32767 ? 32767 : n))

	for (int i = 0; i < *snd->snd_linear_count; i += 2) {
		for (int c = 0; c < 2; ++c) {
			((i16 *)ctx.aframe_tmp->data[c])[frame_idx] =
				CLIP16(((*snd->snd_p)[i + c] * *snd->snd_vol) >> 8);
		}

		if (++frame_idx == ctx.aframe->nb_samples) {
			if (swr_convert(ctx.swr, ctx.aframe->data, ctx.aframe->nb_samples,
					(const u8 *const *)ctx.aframe_tmp->data,
					ctx.aframe_tmp->nb_samples) < 0)
				bail("failed to resample audio frame");

			ctx.aframe->pts = nextpts;
			nextpts += ctx.aframe->nb_samples;
			encode(ctx.acodec_ctx, ctx.aframe, ctx.astream);
			frame_idx = 0;
		}
	}

#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() {
	if (!encode(ctx.vcodec_ctx, 0, ctx.vstream))
		warn("video flush failed");
	if (!encode(ctx.acodec_ctx, 0, ctx.astream))
		warn("audio flush failed!");

	if (av_write_trailer(ctx.format))
		warn("couldn't write trailer!");

	avio_flush(ctx.format->pb);
	avio_close(ctx.format->pb);

	av_frame_free(&ctx.yuv_frame);
	avcodec_free_context(&ctx.vcodec_ctx);
	av_frame_free(&ctx.aframe);
	avcodec_free_context(&ctx.acodec_ctx);
	avformat_free_context(ctx.format);

	sws_freeContext(ctx.sws);

	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!");
	do_stop();
	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);

	info("started recording to %s (%dx%d @ %d)", args.out, args.width,
			args.height, args.fps);
	
	avformat_alloc_output_context2(&ctx.format, 0, 0, args.out);
	if (!ctx.format)
		die("couldn't create output context (does the filename have an ext?)");

	// video codec
	ctx.vcodec = avcodec_find_encoder_by_name(args.encoder);
	ctx.vcodec_ctx = avcodec_alloc_context3(ctx.vcodec);
	ctx.vcodec_ctx->width = args.width;
	ctx.vcodec_ctx->height = args.height;
	ctx.vcodec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
	ctx.vcodec_ctx->time_base = av_make_q(1, args.fps);
	ctx.vcodec_ctx->framerate = av_make_q(args.fps, 1);
	ctx.vcodec_ctx->gop_size = args.fps * 2;
	ctx.vcodec_ctx->bit_rate = args.bitrate * 1000;

	if (!args.bitrate)
		av_opt_set_int(ctx.vcodec_ctx->priv_data, "crf", args.quality, 0);
	if (args.preset)
		av_opt_set(ctx.vcodec_ctx->priv_data, "preset", args.preset, 0);

	if (avcodec_open2(ctx.vcodec_ctx, ctx.vcodec, NULL) < 0)
		die("failed to open video encoder");

	// video stream
	ctx.vstream = avformat_new_stream(ctx.format, NULL);
	if (!ctx.vstream) die("failed to create vide stream");
	avcodec_parameters_from_context(ctx.vstream->codecpar, ctx.vcodec_ctx);
	ctx.vstream->time_base = ctx.vcodec_ctx->time_base;
	ctx.vstream->avg_frame_rate = ctx.vcodec_ctx->framerate;

	// video frame
	ctx.yuv_frame = av_frame_alloc();
	ctx.yuv_frame->width = args.width;
	ctx.yuv_frame->height = args.height;
	ctx.yuv_frame->color_range = AVCOL_RANGE_JPEG;
	ctx.yuv_frame->time_base = ctx.vcodec_ctx->time_base;
	ctx.yuv_frame->format = ctx.vcodec_ctx->pix_fmt;

	ctx.sws = sws_getContext(args.width, args.height, AV_PIX_FMT_BGR24, 
			args.width, args.height, ctx.yuv_frame->format, 
			SWS_FAST_BILINEAR | SWS_FULL_CHR_H_INT | SWS_ACCURATE_RND, 0, 0, 0);

	if (av_frame_get_buffer(ctx.yuv_frame, 0) < 0)
		bail("failed to alloc frame buffer");

	// audio codec
	ctx.acodec = avcodec_find_encoder(AV_CODEC_ID_AAC);
	if (!ctx.acodec) bail("failed to find aac codec");
	ctx.acodec_ctx = avcodec_alloc_context3(ctx.acodec);
	ctx.acodec_ctx->bit_rate = 128000;
	ctx.acodec_ctx->sample_fmt = AV_SAMPLE_FMT_FLTP;
	ctx.acodec_ctx->sample_rate = 44100;
	ctx.acodec_ctx->profile = FF_PROFILE_AAC_MAIN;
	ctx.acodec_ctx->time_base = av_make_q(1, 44100);

	if (av_channel_layout_copy(&ctx.acodec_ctx->ch_layout, 
				&(AVChannelLayout)AV_CHANNEL_LAYOUT_STEREO))
		bail("failed to copy channel layout");

	if (avcodec_open2(ctx.acodec_ctx, ctx.acodec, NULL) < 0)
		die("failed to open audio encoder");

	// audio resampler
	ctx.swr = swr_alloc();
	av_opt_set_int(ctx.swr, "in_channel_count", 2, 0);
	av_opt_set_int(ctx.swr, "in_sample_rate", 44100, 0);
	av_opt_set_sample_fmt(ctx.swr, "in_sample_fmt", AV_SAMPLE_FMT_S16P, 0);
	av_opt_set_chlayout(ctx.swr, "in_chlayout", &ctx.acodec_ctx->ch_layout, 0);
	av_opt_set_int(ctx.swr, "out_channel_count", 2, 0);
	av_opt_set_int(ctx.swr, "out_sample_rate", ctx.acodec_ctx->sample_rate, 0);
	av_opt_set_sample_fmt(ctx.swr, "out_sample_fmt", ctx.acodec_ctx->sample_fmt, 0);
	av_opt_set_chlayout(ctx.swr, "out_chlayout", &ctx.acodec_ctx->ch_layout, 0);

	if (swr_init(ctx.swr) < 0)
		bail("failed to init resampler");

	// audio frame
	ctx.aframe = av_frame_alloc();
	if (!ctx.aframe) bail("failed to alloc audio frame");

	ctx.aframe->nb_samples = ctx.acodec_ctx->frame_size;
	ctx.aframe->format = ctx.acodec_ctx->sample_fmt;
	if (av_channel_layout_copy(&ctx.aframe->ch_layout, &ctx.acodec_ctx->ch_layout))
		bail("failed to copy channel layout");
	if (av_frame_get_buffer(ctx.aframe, 0))
		bail("couldn't alloc audio frame buffer");

	// audio tmp frame (pre resampling)
	ctx.aframe_tmp = av_frame_alloc();
	if (!ctx.aframe_tmp) bail("failed to alloc tmp audio frame");

	ctx.aframe_tmp->nb_samples = ctx.acodec_ctx->frame_size;
	ctx.aframe_tmp->format = ctx.acodec_ctx->sample_fmt;
	if (av_channel_layout_copy(&ctx.aframe_tmp->ch_layout, &ctx.acodec_ctx->ch_layout))
		bail("failed to copy channel layout");
	if (av_frame_get_buffer(ctx.aframe_tmp, 0))
		bail("couldn't alloc audio frame buffer");

	// audio stream
	ctx.astream = avformat_new_stream(ctx.format, NULL);
	if (!ctx.astream) die("failed to create audio stream");
	avcodec_parameters_from_context(ctx.astream->codecpar, ctx.acodec_ctx);
	ctx.astream->time_base = ctx.acodec_ctx->time_base;

	if (avcodec_open2(ctx.acodec_ctx, ctx.acodec, NULL) < 0)
		bail("failed to open audio encoder");
	
	if (avio_open(&ctx.format->pb, args.out, AVIO_FLAG_WRITE) < 0)
		bail("failed to open file");

	if (avformat_write_header(ctx.format, NULL) < 0)
		bail("error when writing to output file");

	return true;
}

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