#include <JuceHeader.h>
#include <cstdio>
#include <cstdlib>
#include <cmath>

#include "PluginHost.h"
#include "AudioInputSource.h"
#include "WavRecorder.h"
#include "TestScenario.h"
#include "PerformanceStats.h"

namespace
{
	void printUsage()
	{
		printf (
			"testhost - headless plugin test host\n"
			"\n"
			"Usage:\n"
			"  testhost --plugin <path.so> --scenario <file.json> [--perfCsv <file>]\n"
			"  testhost --plugin <path.so> [ad-hoc options] [--perfCsv <file>]\n"
			"\n"
			"With --scenario, a JSON file drives everything (sample rate, block size,\n"
			"duration, input signal, a timed sequence of MIDI/param/patch/bank events) -\n"
			"see TestScenario.h for the format. All other ad-hoc options below are ignored\n"
			"when --scenario is given.\n"
			"\n"
			"Ad-hoc options:\n"
			"  --patch <file>          load a patch (single-program) state before rendering\n"
			"  --bank <file>           load a bank (full plugin) state before rendering\n"
			"  --savePatch <file>      save the patch state after rendering\n"
			"  --saveBank <file>       save the bank state after rendering\n"
			"  --note <note> <vel>     MIDI note-on (0-127, 0-127) at sample 0, held for the\n"
			"                          whole render, note-off in the final block\n"
			"  --param <index> <value> set parameter <index> to <value> (0.0-1.0) before rendering\n"
			"  --input <mode>          fill the plugin's input channels with: silence (default),\n"
			"                          sine[:freqHz], white, pink, or impulse\n"
			"  --duration <seconds>    how much audio to render (default 2.0)\n"
			"  --sampleRate <hz>       default 44100\n"
			"  --blockSize <n>         default 512\n"
			"  --wav <file>            record the plugin's output to a WAV file\n"
			"  --printParams           print every parameter's index and current value\n"
			"                          (after any --patch/--bank/--param has been applied,\n"
			"                          before rendering) as \"PARAM <index> <value>\" lines\n"
			"\n"
			"Patch/bank state is exactly what JUCE's AudioPluginInstance::get/setCurrentProgram-\n"
			"StateInformation and get/setStateInformation read and write for a hosted VST2\n"
			"plugin - i.e. real .fxp/.fxb file bytes.\n"
		);
	}

	struct Args
	{
		String pluginPath;
		String scenarioFile;

		// Ad-hoc mode only (ignored if scenarioFile is set):
		String patchFile, bankFile, savePatchFile, saveBankFile, wavFile;
		double durationSeconds = 2.0;
		double sampleRate = 44100.0;
		int blockSize = 512;
		bool hasNote = false;
		int note = 60;
		int velocity = 100;
		bool hasParam = false;
		int paramIndex = 0;
		float paramValue = 0.0f;
		AudioInputSource::Type inputType = AudioInputSource::silence;
		double inputFrequencyHz = 440.0;

		String perfCsvFile;
		bool printParams = false;
	};

	bool parseArgs (int argc, char *argv[], Args &args)
	{
		for (int i = 1; i < argc; ++i)
		{
			const String arg (argv[i]);

			if (arg == "--plugin" && i + 1 < argc)
				args.pluginPath = argv[++i];
			else if (arg == "--scenario" && i + 1 < argc)
				args.scenarioFile = argv[++i];
			else if (arg == "--patch" && i + 1 < argc)
				args.patchFile = argv[++i];
			else if (arg == "--bank" && i + 1 < argc)
				args.bankFile = argv[++i];
			else if (arg == "--savePatch" && i + 1 < argc)
				args.savePatchFile = argv[++i];
			else if (arg == "--saveBank" && i + 1 < argc)
				args.saveBankFile = argv[++i];
			else if (arg == "--wav" && i + 1 < argc)
				args.wavFile = argv[++i];
			else if (arg == "--duration" && i + 1 < argc)
				args.durationSeconds = String (argv[++i]).getDoubleValue();
			else if (arg == "--sampleRate" && i + 1 < argc)
				args.sampleRate = String (argv[++i]).getDoubleValue();
			else if (arg == "--blockSize" && i + 1 < argc)
				args.blockSize = String (argv[++i]).getIntValue();
			else if (arg == "--perfCsv" && i + 1 < argc)
				args.perfCsvFile = argv[++i];
			else if (arg == "--printParams")
				args.printParams = true;
			else if (arg == "--note" && i + 2 < argc)
			{
				args.hasNote = true;
				args.note = String (argv[++i]).getIntValue();
				args.velocity = String (argv[++i]).getIntValue();
			}
			else if (arg == "--param" && i + 2 < argc)
			{
				args.hasParam = true;
				args.paramIndex = String (argv[++i]).getIntValue();
				args.paramValue = String (argv[++i]).getFloatValue();
			}
			else if (arg == "--input" && i + 1 < argc)
			{
				const String mode (argv[++i]);
				if (mode.startsWith ("sine"))
				{
					args.inputType = AudioInputSource::sine;
					const int colon = mode.indexOfChar (':');
					if (colon >= 0)
						args.inputFrequencyHz = mode.substring (colon + 1).getDoubleValue();
				}
				else if (mode == "white")
					args.inputType = AudioInputSource::whiteNoise;
				else if (mode == "pink")
					args.inputType = AudioInputSource::pinkNoise;
				else if (mode == "impulse")
					args.inputType = AudioInputSource::impulse;
				else if (mode == "silence")
					args.inputType = AudioInputSource::silence;
				else
				{
					fprintf (stderr, "testhost: unknown --input mode '%s'\n", mode.toRawUTF8());
					return false;
				}
			}
			else if (arg == "--help" || arg == "-h")
				return false;
			else
			{
				fprintf (stderr, "testhost: unrecognised argument '%s'\n", arg.toRawUTF8());
				return false;
			}
		}

		if (args.pluginPath.isEmpty())
		{
			fprintf (stderr, "testhost: --plugin is required\n");
			return false;
		}

		return true;
	}

	/** Builds the internal TestScenario representation from the simple ad-hoc
		CLI flags, so the render loop below only has to be written once,
		whether it's driven by a JSON file or by --note/--param/etc. */
	void buildScenarioFromArgs (const Args &args, TestScenario &scenario)
	{
		scenario.sampleRate = args.sampleRate;
		scenario.blockSize = args.blockSize;
		scenario.durationSamples = (int64) (args.durationSeconds * args.sampleRate);
		scenario.inputType = args.inputType;
		scenario.inputFrequencyHz = args.inputFrequencyHz;
		scenario.inputAmplitude = 0.5f;
		scenario.initialLoadPatchFile = args.patchFile;
		scenario.initialLoadBankFile = args.bankFile;
		scenario.recordWavFile = args.wavFile;

		if (args.hasNote)
		{
			TestScenario::Event on;
			on.sample = 0;
			on.type = TestScenario::Event::noteOn;
			on.note = args.note;
			on.velocity = args.velocity;
			scenario.events.add (on);

			TestScenario::Event off;
			off.sample = jmax ((int64) 0, scenario.durationSamples - 1);
			off.type = TestScenario::Event::noteOff;
			off.note = args.note;
			scenario.events.add (off);
		}

		if (args.hasParam)
		{
			TestScenario::Event p;
			p.sample = 0;
			p.type = TestScenario::Event::param;
			p.paramIndex = args.paramIndex;
			p.paramValue = args.paramValue;
			scenario.events.add (p);
		}
	}

	bool loadStateFile (const String &path, bool isBank, PluginHost &host)
	{
		MemoryBlock data;
		if (! File (path).loadFileAsData (data) || data.getSize() == 0)
		{
			fprintf (stderr, "testhost: could not read '%s'\n", path.toRawUTF8());
			return false;
		}

		if (isBank)
			host.setStateInformation (data.getData(), (int) data.getSize());
		else
			host.setCurrentProgramStateInformation (data.getData(), (int) data.getSize());

		return true;
	}

	void saveStateFile (const String &path, bool isBank, PluginHost &host)
	{
		MemoryBlock data;
		if (isBank)
			host.getStateInformation (data);
		else
			host.getCurrentProgramStateInformation (data);

		if (! File (path).replaceWithData (data.getData(), data.getSize()))
			fprintf (stderr, "testhost: could not write '%s'\n", path.toRawUTF8());
	}

	void applyEvent (PluginHost &host, const TestScenario::Event &event, MidiBuffer &midi, int sampleOffsetInBlock)
	{
		switch (event.type)
		{
			case TestScenario::Event::noteOn:
				midi.addEvent (MidiMessage::noteOn (event.channel, event.note, (uint8) event.velocity), sampleOffsetInBlock);
				break;
			case TestScenario::Event::noteOff:
				midi.addEvent (MidiMessage::noteOff (event.channel, event.note), sampleOffsetInBlock);
				break;
			case TestScenario::Event::controller:
				midi.addEvent (MidiMessage::controllerEvent (event.channel, event.controllerNumber, event.controllerValue), sampleOffsetInBlock);
				break;
			case TestScenario::Event::param:
				host.setParameter (event.paramIndex, event.paramValue);
				break;
			case TestScenario::Event::loadPatch:
				loadStateFile (event.file, false, host);
				break;
			case TestScenario::Event::loadBank:
				loadStateFile (event.file, true, host);
				break;
			case TestScenario::Event::savePatch:
				saveStateFile (event.file, false, host);
				break;
			case TestScenario::Event::saveBank:
				saveStateFile (event.file, true, host);
				break;
		}
	}
}

int main (int argc, char *argv[])
{
	Args args;
	if (! parseArgs (argc, argv, args))
	{
		printUsage();
		return 1;
	}

	// Even a headless console app must bring up JUCE's message/GUI plumbing
	// before touching any AudioProcessor/AudioPluginFormat class - no window
	// is ever created, but juce_audio_processors.h unconditionally depends
	// on juce_gui_basics.
	ScopedJuceInitialiser_GUI juceInit;

	TestScenario scenario;
	if (args.scenarioFile.isNotEmpty())
	{
		if (! scenario.loadFromFile (File (args.scenarioFile)))
			return 1;
	}
	else
	{
		buildScenarioFromArgs (args, scenario);
	}

	PluginHost host;
	if (! host.load (args.pluginPath, scenario.sampleRate, scenario.blockSize))
		return 1;

	const int numInputChannels = host.getNumInputChannels();
	const int numOutputChannels = host.getNumOutputChannels();
	printf ("testhost: loaded '%s' (%d in / %d out, %d parameters)\n",
		args.pluginPath.toRawUTF8(), numInputChannels, numOutputChannels, host.getNumParameters());

	if (scenario.initialLoadBankFile.isNotEmpty() && ! loadStateFile (scenario.initialLoadBankFile, true, host))
		return 1;

	if (scenario.initialLoadPatchFile.isNotEmpty() && ! loadStateFile (scenario.initialLoadPatchFile, false, host))
		return 1;

	if (args.printParams)
	{
		for (int i = 0; i < host.getNumParameters(); ++i)
			printf ("PARAM %d %.6f\n", i, host.getParameter (i));
	}

	AudioInputSource inputSource;
	inputSource.prepare (scenario.sampleRate);
	switch (scenario.inputType)
	{
		case AudioInputSource::sine:       inputSource.setSine (scenario.inputFrequencyHz, scenario.inputAmplitude); break;
		case AudioInputSource::whiteNoise: inputSource.setWhiteNoise (scenario.inputAmplitude); break;
		case AudioInputSource::pinkNoise:  inputSource.setPinkNoise (scenario.inputAmplitude); break;
		case AudioInputSource::impulse:    inputSource.setImpulse (scenario.inputAmplitude); break;
		case AudioInputSource::silence:
		default:                          inputSource.setSilence(); break;
	}

	WavRecorder recorder;
	if (scenario.recordWavFile.isNotEmpty())
	{
		if (! recorder.start (File (scenario.recordWavFile), scenario.sampleRate, numOutputChannels))
			return 1;
	}

	const int numChannels = jmax (numInputChannels, numOutputChannels, 1);
	AudioSampleBuffer buffer (numChannels, scenario.blockSize);

	PerformanceStats perfStats;
	perfStats.reset (scenario.sampleRate);

	int nextEventIndex = 0;
	int64 samplesDone = 0;
	bool sawNaN = false;

	while (samplesDone < scenario.durationSamples)
	{
		const int numThisBlock = (int) jmin ((int64) scenario.blockSize, scenario.durationSamples - samplesDone);

		buffer.clear();
		if (numInputChannels > 0)
			inputSource.fillNextBlock (buffer, numInputChannels, numThisBlock);

		MidiBuffer midi;

		// Non-MIDI actions (param/patch/bank) scheduled inside this block take
		// effect at the start of the block, not at their exact sample offset -
		// good enough for a test tool; MIDI events keep full per-sample timing
		// via MidiBuffer's own sample-offset parameter.
		while (nextEventIndex < scenario.events.size()
				&& scenario.events.getReference (nextEventIndex).sample < samplesDone + numThisBlock)
		{
			const TestScenario::Event &event = scenario.events.getReference (nextEventIndex);
			const int offset = (int) jlimit ((int64) 0, (int64) (numThisBlock - 1), event.sample - samplesDone);
			applyEvent (host, event, midi, offset);
			++nextEventIndex;
		}

		perfStats.timeBlock (numThisBlock, [&] { host.process (buffer, midi); });

		for (int ch = 0; ch < numOutputChannels && ! sawNaN; ++ch)
		{
			const float *data = buffer.getReadPointer (ch);
			for (int i = 0; i < numThisBlock; ++i)
			{
				if (! std::isfinite (data[i]))
				{
					sawNaN = true;
					break;
				}
			}
		}

		if (recorder.isRecording())
			recorder.write (buffer, numThisBlock);

		samplesDone += numThisBlock;
	}

	// Any trailing savePatch/saveBank/etc. events scheduled at or after the
	// final rendered sample still get applied, in order.
	while (nextEventIndex < scenario.events.size())
	{
		MidiBuffer unused;
		applyEvent (host, scenario.events.getReference (nextEventIndex), unused, 0);
		++nextEventIndex;
	}

	recorder.stop();

	if (args.savePatchFile.isNotEmpty())
		saveStateFile (args.savePatchFile, false, host);
	if (args.saveBankFile.isNotEmpty())
		saveStateFile (args.saveBankFile, true, host);

	host.releaseResources();

	if (sawNaN)
		fprintf (stderr, "testhost: WARNING - output contained NaN/Inf samples\n");

	printf ("testhost: rendered %.3fs (%lld samples)%s%s\n",
		(double) samplesDone / scenario.sampleRate, (long long) samplesDone,
		scenario.recordWavFile.isNotEmpty() ? " -> " : "",
		scenario.recordWavFile.isNotEmpty() ? scenario.recordWavFile.toRawUTF8() : "");

	perfStats.printSummary();
	if (args.perfCsvFile.isNotEmpty())
		perfStats.writeCsv (File (args.perfCsvFile));

	return sawNaN ? 2 : 0;
}
