/*
This Source Code Form is subject to the
terms of the Mozilla Public License, v.
2.0. If a copy of the MPL was not
distributed with this file, You can
obtain one at
http://mozilla.org/MPL/2.0/.
*/

#include "scripting/PyInt.h"
#include "stdafx.h"

#include "application/application.h"
#include "utilities/Globals.h"
#include "utilities/Logs.h"
#include "utilities/dictionary.h"

#ifdef __GNUC__
#pragma GCC diagnostic ignored "-Wwrite-strings"
#endif
#include <simulation/simulation.h>

void render_task::run()
{

	// convert provided input to a python dictionary
	auto *input = PyDict_New();
	if (input == nullptr)
	{
		cancel();
		return;
	}
	for (auto const &datapair : m_input->floats)
	{
		auto *value{PyGetFloat(datapair.second)};
		if (value == nullptr)
		{
			PyErr_Clear();
			continue;
		}
		PyDict_SetItemString(input, datapair.first.c_str(), value);
		Py_DECREF(value);
	}
	for (auto const &datapair : m_input->integers)
	{
		auto *value{PyLong_FromLong(datapair.second)};
		if (value == nullptr)
		{
			PyErr_Clear();
			continue;
		}
		PyDict_SetItemString(input, datapair.first.c_str(), value);
		Py_DECREF(value);
	}
	for (auto const &datapair : m_input->bools)
	{
		// Py_True / Py_False sa niesmiertelne, ale PyDict_SetItemString i tak
		// pobiera wlasna referencje - nie zwalniamy
		auto *value{PyGetBool(datapair.second)};
		PyDict_SetItemString(input, datapair.first.c_str(), value);
	}
	for (auto const &datapair : m_input->strings)
	{
		// Nazwy scenerii/wagonow (asName, SceneryFile, asCarName, cCode...) moga
		// byc albo w UTF-8, albo w starym kodowaniu Windows-1250. PyUnicode_FromString
		// wymaga POPRAWNEGO UTF-8 - dla nie-UTF-8 zwracalo NULL, a potem
		// PyDict_SetItemString(dict, key, NULL) robilo Py_INCREF na NULL -> crash.
		//
		// Strategia: najpierw probujemy UTF-8 (strict). Jesli sie nie uda - probujemy
		// cp1250. Jesli oba zawioda - pomijamy klucz, ale nie wywracamy symulatora.
		char const *const str{datapair.second.c_str()};
		Py_ssize_t const len{static_cast<Py_ssize_t>(datapair.second.size())};

		auto *value{PyUnicode_DecodeUTF8(str, len, "strict")};
		if (value == nullptr)
		{
			PyErr_Clear();
			value = PyUnicode_Decode(str, len, "cp1250", "strict");
		}
		if (value == nullptr)
		{
			PyErr_Clear();
			continue;
		}
		PyDict_SetItemString(input, datapair.first.c_str(), value);
		Py_DECREF(value);
	}
	for (auto const &datapair : m_input->vec2_lists)
	{
		PyObject *list = PyList_New(datapair.second.size());

		for (size_t i = 0; i < datapair.second.size(); i++)
		{
			auto const &vec = datapair.second[i];
			WriteLog("passing " + glm::to_string(vec));

			PyObject *tuple = PyTuple_New(2);
			PyTuple_SetItem(tuple, 0, PyGetFloat(vec.x)); // steals ref
			PyTuple_SetItem(tuple, 1, PyGetFloat(vec.y)); // steals ref

			PyList_SetItem(list, i, tuple); // steals ref
		}

		PyDict_SetItemString(input, datapair.first.c_str(), list);
		Py_DECREF(list);
	}
	m_input = nullptr;

	// call the renderer
	auto *output{PyObject_CallMethod(m_renderer, const_cast<char *>("render"), const_cast<char *>("O"), input)};
	Py_DECREF(input);

	if (output != nullptr)
	{
		auto *outputWidth = PyObject_CallMethod(m_renderer, const_cast<char *>("get_width"), nullptr);
		auto *outputHeight = PyObject_CallMethod(m_renderer, const_cast<char *>("get_height"), nullptr);

		if (outputWidth != nullptr && outputHeight != nullptr && m_target != nullptr)
		{
			const int screenWidth = static_cast<int>(PyLong_AsLong(outputWidth));
			const int screenHeight = static_cast<int>(PyLong_AsLong(outputHeight));

			const bool useRgb = false && !Global.gfx_usegles;

			const int glFormat = useRgb ? GL_SRGB8 : GL_SRGB8_ALPHA8;
			const int glComponents = useRgb ? GL_RGB : GL_RGBA;
			const size_t bytesPerPixel = useRgb ? 3u : 4u;
			const size_t expectedBytes = static_cast<size_t>(screenWidth) * static_cast<size_t>(screenHeight) * bytesPerPixel;

			Py_ssize_t pythonBufferBytes = 0;
			char *pythonBufferPtr = nullptr;

			const bool bufferExtracted = PyBytes_AsStringAndSize(output, &pythonBufferPtr, &pythonBufferBytes) == 0 && pythonBufferPtr != nullptr;

			if (!bufferExtracted)
			{
				ErrorLog("Python screen renderer: output is not a valid byte buffer");
			}
			else if (pythonBufferBytes < static_cast<Py_ssize_t>(expectedBytes))
			{
				ErrorLog(std::format("Python screen renderer: output buffer too small ({} bytes, expected {})", pythonBufferBytes, expectedBytes));
			}
			else
			{
				std::lock_guard guard(m_target->mutex);

				if (m_target->image.size() != expectedBytes)
					m_target->image.resize(expectedBytes);

				std::memcpy(m_target->image.data(), pythonBufferPtr, expectedBytes);

				m_target->width = screenWidth;
				m_target->height = screenHeight;
				m_target->components = glComponents;
				m_target->format = glFormat;
				m_target->timestamp = std::chrono::high_resolution_clock::now();
			}
		}

		if (outputHeight != nullptr)
			Py_DECREF(outputHeight);
		if (outputWidth != nullptr)
			Py_DECREF(outputWidth);

		Py_DECREF(output);
	}

	// get commands from renderer
	auto *commandsPO = PyObject_CallMethod(m_renderer, const_cast<char *>("getCommands"), nullptr);
	if (commandsPO != nullptr)
	{
		std::vector<std::string> commands = python_external_utils::PyObjectToStringArray(commandsPO);

		Py_DECREF(commandsPO);
		// we perform any actions ONLY when there are any commands in buffer
		if (!commands.empty())
		{
			for (const auto &cmd : commands)
			{
				std::string baseCmd;
				int p1 = 0, p2 = 0;

				size_t pos1 = cmd.find(';');
				if (pos1 == std::string::npos)
				{
					baseCmd = cmd;
				}
				else
				{
					baseCmd = cmd.substr(0, pos1);

					size_t pos2 = cmd.find(';', pos1 + 1);
					if (pos2 == std::string::npos)
					{
						p1 = std::stoi(cmd.substr(pos1 + 1));
					}
					else
					{
						p1 = std::stoi(cmd.substr(pos1 + 1, pos2 - pos1 - 1));
						p2 = std::stoi(cmd.substr(pos2 + 1));
					}
				}

				auto it = simulation::commandMap.find(baseCmd);
				if (it != simulation::commandMap.end())
				{
					command_data cd;
					cd.command = it->second;
					cd.action = GLFW_PRESS;
					cd.param1 = p1;
					cd.param2 = p2;

					WriteLog("Python: Executing command [" + baseCmd + "] with params: P1=" + std::to_string(p1) + " P2=" + std::to_string(p2) +
					         " Target ID=" + std::to_string(simulation::Train->id()));

					simulation::Commands.push(cd, static_cast<size_t>(command_target::vehicle) | simulation::Train->id());
				}
				else
				{
					ErrorLog("Python: Command [" + baseCmd + "] not found!");
				}
			}
		}
	}
}

void render_task::upload()
{
	if (Global.python_uploadmain && m_target && m_target->shared_tex)
	{
		m_target->shared_tex->update_from_memory(m_target->width, m_target->height, reinterpret_cast<const uint8_t *>(m_target->image.data()));
	}
}

void render_task::cancel() {}

// initializes the module. returns true on success
auto python_taskqueue::init() -> bool
{

	crashreport_add_info("python.threadedupload", Global.python_threadedupload ? "yes" : "no");
	crashreport_add_info("python.uploadmain", Global.python_uploadmain ? "yes" : "no");
#ifdef _WIN32
	const wchar_t *pythonhome = sizeof(void *) == 8 ? L"python64" : L"python";
#elif __linux__
	const wchar_t *pythonhome = (sizeof(void *) == 8) ? L"linuxpython64" : L"linuxpython";
#elif __APPLE__
	const wchar_t *pythonhome = (sizeof(void *) == 8) ? L"macpython64" : L"macpython";
#endif

	{
		// Py_SetPythonHome / Py_InitializeEx are deprecated since Python 3.11.
		// Use the PyConfig init API (PEP 587) instead.
		PyConfig config;
		PyConfig_InitPythonConfig(&config);
		config.install_signal_handlers = 0; // matches former Py_InitializeEx(0)

		PyStatus status = PyConfig_SetString(&config, &config.home, pythonhome);
		if (PyStatus_Exception(status))
		{
			PyConfig_Clear(&config);
			ErrorLog("Python Interpreter: failed to set PYTHONHOME");
			return false;
		}

		status = Py_InitializeFromConfig(&config);
		if (PyStatus_Exception(status))
		{
			std::string msg = "Python Interpreter: Py_InitializeFromConfig failed";
			if (status.err_msg != nullptr)
				msg += std::string(": ") + status.err_msg;
			if (status.func != nullptr)
				msg += std::string(" (in ") + status.func + ")";
			PyConfig_Clear(&config);
			ErrorLog(msg);
			return false;
		}
		PyConfig_Clear(&config);
	}

	// make CWD-local packages (e.g. the asset-side "scripts" package used by
	// texture/screen renderers) importable. Embedded CPython builds sys.path
	// from PYTHONHOME only; unlike the python CLI it does not prepend the
	// working directory, and unlike Windows' getpathp.c it does not add the
	// executable's directory either, so on POSIX "from scripts import ..."
	// would otherwise fail. "." matches how the engine resolves all other
	// assets (relative to CWD). Idempotent/harmless on Windows.
	if (PyObject *syspath = PySys_GetObject("path")) // borrowed ref
	{
		PyObject *cwd = PyUnicode_FromString(".");
		if (cwd != nullptr)
		{
			PyList_Insert(syspath, 0, cwd);
			Py_DECREF(cwd);
		}
	}

	PyObject *stringiomodule{nullptr};
	PyObject *stringioclassname{nullptr};
	PyObject *stringioobject{nullptr};

	// do the setup work while we hold the lock
	m_main = PyImport_ImportModule("__main__");
	if (m_main == nullptr)
	{
		ErrorLog("Python Interpreter: __main__ module is missing");
		goto release_and_exit;
	}

	stringiomodule = PyImport_ImportModule("io");
	stringioclassname = stringiomodule != nullptr ? PyObject_GetAttrString(stringiomodule, "StringIO") : nullptr;
	stringioobject = stringioclassname != nullptr ? PyObject_CallObject(stringioclassname, nullptr) : nullptr;
	m_stderr = {(stringioobject == nullptr ? nullptr : PySys_SetObject(const_cast<char *>("stderr"), stringioobject) != 0 ? nullptr : stringioobject)};

	if (false == run_file("abstractscreenrenderer"))
	{
		goto release_and_exit;
	}

	// release the lock, save the state for future use
	m_mainthread = PyEval_SaveThread();

	WriteLog("Python Interpreter: setup complete");

	// init workers
	for (auto &worker : m_workers)
	{

		GLFWwindow *openglcontextwindow = nullptr;
		if (Global.python_threadedupload)
			openglcontextwindow = Application.window(-1);
		worker = std::jthread(&python_taskqueue::run, this, openglcontextwindow, std::ref(m_tasks), std::ref(m_uploadtasks), std::ref(m_condition), std::ref(m_exit));

		if (false == worker.joinable())
		{
			return false;
		}
	}

	m_initialized = true;

	return true;

release_and_exit:
	PyEval_SaveThread();
	return false;
}

// shuts down the module
void python_taskqueue::exit()
{
	if (!m_initialized)
		return;

	// let the workers know we're done with them
	m_exit = true;
	m_condition.notify_all();
	// let them free up their shit before we proceed
	m_workers = {};
	// drop any pending render/upload tasks. the workers are joined at this
	// point so the locks are strictly defensive, but they cost nothing and
	// document intent. clearing the deque also actually releases the tasks,
	// which the previous code did not do (cancel() is a no-op stub).
	{
		std::lock_guard<std::mutex> lock(m_tasks.mutex);
		for (auto &task : m_tasks.data)
		{
			task->cancel();
		}
		m_tasks.data.clear();
	}
	{
		std::lock_guard<std::mutex> lock(m_uploadtasks.mutex);
		m_uploadtasks.data.clear();
	}
	// reclaim cached python objects while the interpreter is still alive,
	// so no Py_DECREF lands on a finalized interpreter during later teardown
	acquire_lock();
	for (auto &entry : m_renderers)
	{
		Py_XDECREF(entry.second);
	}
	m_renderers.clear();
	Py_XDECREF(m_stderr);
	m_stderr = nullptr;
	Py_XDECREF(m_main);
	m_main = nullptr;
	// take a bow
	Py_Finalize();
	m_initialized = false;
}

// adds specified task along with provided collection of data to the work queue. returns true on success
auto python_taskqueue::insert(task_request const &Task) -> bool
{

	if (!m_initialized || false == Global.python_enabled || Task.renderer.empty() || Task.input == nullptr || Task.target == 0)
	{
		return false;
	}

	auto *renderer{fetch_renderer(Task.renderer)};
	if (renderer == nullptr)
	{
		return false;
	}

	auto newtask = std::make_shared<render_task>(renderer, Task.input, Task.target);
	bool newtaskinserted{false};
	// acquire a lock on the task queue and add the new task
	{
		std::lock_guard<std::mutex> lock(m_tasks.mutex);
		// check the task list for a pending request with the same target
		for (auto &task : m_tasks.data)
		{
			if (task->target() == Task.target)
			{
				// replace pending task in the slot with the more recent one
				task->cancel();
				task = newtask;
				newtaskinserted = true;
				break;
			}
		}
		if (false == newtaskinserted)
		{
			m_tasks.data.emplace_back(newtask);
		}
	}
	// potentially wake a worker to handle the new task
	m_condition.notify_one();
	// all done
	return true;
}

// executes python script stored in specified file. returns true on success
auto python_taskqueue::run_file(std::string const &File, std::string const &Path) -> bool
{

	auto const lookup{FileExists({Path + File, "python/local/" + File}, {".py"})};
	if (lookup.first.empty())
	{
		return false;
	}

	std::ifstream inputfile{lookup.first + lookup.second};
	std::string input;
	input.assign(std::istreambuf_iterator<char>(inputfile), std::istreambuf_iterator<char>());

	if (PyRun_SimpleString(input.c_str()) != 0)
	{
		error();
		return false;
	}

	return true;
}

// acquires the python gil and sets the main thread as current
void python_taskqueue::acquire_lock()
{

	PyEval_RestoreThread(m_mainthread);
}

// releases the python gil and swaps the main thread out
void python_taskqueue::release_lock()
{

	PyEval_SaveThread();
}

auto python_taskqueue::fetch_renderer(std::string const Renderer) -> PyObject *
{

	auto const lookup{m_renderers.find(Renderer)};
	if (lookup != std::end(m_renderers))
	{
		return lookup->second;
	}
	// try to load specified renderer class
	std::string const path{substr_path(Renderer)};
	auto const file{Renderer.substr(path.size())};
	PyObject *renderer{nullptr};
	PyObject *rendererarguments{nullptr};
	PyObject *renderername{nullptr};
	acquire_lock();
	{
		if (m_main == nullptr)
		{
			ErrorLog("Python Renderer: __main__ module is missing");
			goto cache_and_return;
		}

		if (false == run_file(file, path))
		{
			goto cache_and_return;
		}
		renderername = PyObject_GetAttrString(m_main, file.c_str());
		if (renderername == nullptr)
		{
			ErrorLog("Python Renderer: class \"" + file + "\" not defined");
			goto cache_and_return;
		}
		rendererarguments = Py_BuildValue("(s)", path.c_str());
		if (rendererarguments == nullptr)
		{
			ErrorLog("Python Renderer: failed to create initialization arguments");
			goto cache_and_return;
		}
		renderer = PyObject_CallObject(renderername, rendererarguments);

		PyObject_CallMethod(renderer, const_cast<char *>("manul_set_format"), const_cast<char *>("(s)"), "RGBA");

		if (PyErr_Occurred() != nullptr)
		{
			error();
			renderer = nullptr;
		}

	cache_and_return:
		// clean up after yourself
		if (rendererarguments != nullptr)
		{
			Py_DECREF(rendererarguments);
		}
	}
	release_lock();
	// cache the failures as well so we don't try again on subsequent requests
	m_renderers.emplace(Renderer, renderer);
	return renderer;
}

void python_taskqueue::run(GLFWwindow *Context, rendertask_sequence &Tasks, uploadtask_sequence &Upload_Tasks, threading::condition_variable &Condition, std::atomic<bool> &Exit)
{

	if (Context)
		glfwMakeContextCurrent(Context);

	// create a state object for this thread
	PyEval_AcquireThread(m_mainthread);
	auto *threadstate{PyThreadState_New(m_mainthread->interp)};
	PyEval_ReleaseThread(m_mainthread);

	std::shared_ptr<render_task> task{nullptr};

	while (false == Exit.load())
	{
		// regardless of the reason we woke up prime the spurious wakeup flag for the next time
		Condition.spurious(true);
		// keep working as long as there's any scheduled tasks
		do
		{
			task = nullptr;
			// acquire a lock on the task queue and potentially grab a task from it
			{
				std::lock_guard<std::mutex> lock(Tasks.mutex);
				if (false == Tasks.data.empty())
				{
					// fifo
					task = Tasks.data.front();
					Tasks.data.pop_front();
				}
			}
			if (task != nullptr)
			{
				// swap in my thread state
				PyEval_RestoreThread(threadstate);
				{
					// execute python code
					task->run();
					if (Context)
						task->upload();
					else
					{
						std::lock_guard<std::mutex> lock(Upload_Tasks.mutex);
						Upload_Tasks.data.push_back(task);
					}
					if (PyErr_Occurred() != nullptr)
						error();
				}
				// clear the thread state
				PyEval_SaveThread();
			}
			// TBD, TODO: add some idle time between tasks in case we're on a single thread cpu?
		} while (task != nullptr);
		// if there's nothing left to do wait until there is
		// short timeout: notify_all() in exit() can race with the worker's
		// transition into the wait, and the run loop also drives prompt
		// shutdown checks
		Condition.wait_for(std::chrono::milliseconds(250));
	}
	// clean up thread state data
	// Re-acquire the GIL with this worker's thread state, then clear and delete it.
	// PyThreadState_DeleteCurrent() frees the current thread state AND releases the GIL,
	// so we must NOT call PyEval_SaveThread() afterwards: at that point the current
	// thread state is already NULL and PyEval_SaveThread() fatals on a NULL tstate.
	PyEval_RestoreThread(threadstate);
	PyThreadState_Clear(threadstate);
	PyThreadState_DeleteCurrent();

	// detach the GL context before the worker terminates; some drivers
	// (NVIDIA on X11, certain Mesa/Wayland configs) hang in process teardown
	// if a context is still current on a dying thread
	if (Context)
		glfwMakeContextCurrent(nullptr);
}

void python_taskqueue::update()
{
	std::lock_guard<std::mutex> lock(m_uploadtasks.mutex);

	for (auto &task : m_uploadtasks.data)
		task->upload();

	m_uploadtasks.data.clear();
}

void python_taskqueue::error()
{
	if (m_stderr != nullptr)
	{
		// std err pythona jest buforowane
		PyErr_Print();
		auto *errortext{PyObject_CallMethod(m_stderr, const_cast<char *>("getvalue"), nullptr)};
		if (errortext != nullptr)
		{
			const char *errstr = PyUnicode_AsUTF8(errortext);
			if (errstr != nullptr)
				ErrorLog(errstr);
			Py_DECREF(errortext);
		}
		PyObject_CallMethod(m_stderr, const_cast<char *>("truncate"), const_cast<char *>("L"), (long long)0);
		PyObject_CallMethod(m_stderr, const_cast<char *>("seek"), const_cast<char *>("L"), (long long)0);
	}
	else
	{
		// nie dziala buffor pythona
		PyObject *type, *value, *traceback;
		PyErr_Fetch(&type, &value, &traceback);
		if (type == nullptr)
		{
			ErrorLog("Python Interpreter: don't know how to handle null exception");
		}
		PyErr_NormalizeException(&type, &value, &traceback);
		if (type == nullptr)
		{
			ErrorLog("Python Interpreter: don't know how to handle null exception");
		}
		auto *typetext{PyObject_Str(type)};
		if (typetext != nullptr)
		{
			const char *s = PyUnicode_AsUTF8(typetext);
			if (s)
				ErrorLog(s);
			Py_DECREF(typetext);
		}
		if (value != nullptr)
		{
			auto *valuetext{PyObject_Str(value)};
			if (valuetext != nullptr)
			{
				const char *s = PyUnicode_AsUTF8(valuetext);
				if (s)
					ErrorLog(s);
				Py_DECREF(valuetext);
			}
		}
		auto *tracebacktext{PyObject_Str(traceback)};
		if (tracebacktext != nullptr)
		{
			const char *s = PyUnicode_AsUTF8(tracebacktext);
			if (s)
				ErrorLog(s);
			Py_DECREF(tracebacktext);
		}
		else
		{
			WriteLog("Python Interpreter: failed to retrieve the stack traceback");
		}
	}
}

std::vector<std::string> python_external_utils::PyObjectToStringArray(PyObject *pyList)
{
	std::vector<std::string> result;
	std::vector<std::string> emptyIfError = {};
	if (!PySequence_Check(pyList))
	{
		ErrorLog("Python: Failed to convert PyObject -> vector<string>");
		return emptyIfError;
	}

	Py_ssize_t size = PySequence_Size(pyList);
	for (Py_ssize_t i = 0; i < size; ++i)
	{
		PyObject *item = PySequence_GetItem(pyList, i); // Increments reference count
		if (item == nullptr)
		{
			ErrorLog("Python: Failed to get item from sequence.");
			return emptyIfError;
		}

		const char *str = PyUnicode_AsUTF8(item);
		if (str == nullptr)
		{
			Py_DECREF(item);
			ErrorLog("Python: Failed to convert item to string.");
			return emptyIfError;
		}

		result.push_back(std::string(str));
		Py_DECREF(item); // Decrease reference count for the item
	}

	return result;
}

#ifdef __GNUC__
#pragma GCC diagnostic pop
#endif