#include <pslang/parser/parser.hpp>
#include <pslang/parser/error.hpp>
#include <pslang/interpreter/exec.hpp>
#include <pslang/interpreter/error.hpp>
#include <pslang/ast/statement.hpp>
#include <pslang/ast/preprocess.hpp>
#include <pslang/ast/print.hpp>
#include <pslang/jit/jit.hpp>
#include <pslang/jit/executable.hpp>

#include <filesystem>
#include <iostream>
#include <fstream>
#include <cstring>
#include <cmath>

std::string extract_nth_line(std::filesystem::path const & path, std::size_t n)
{
	std::ifstream file(path);
	std::string line;

	for (size_t i = 1; i <= n && std::getline(file, line); ++i) {
		if (i == n) {
			return line;
		}
	}

	return {};
}

std::size_t replace_tabs_with_spaces(std::string & str, std::size_t count)
{
	std::string replaced;
	for (char c : str)
	{
		if (c == '\t')
			replaced.append(count, ' ');
		else
			replaced.append(1, c);
	}
	std::swap(str, replaced);
	return str.size() - replaced.size();
}

void print_error_context(std::filesystem::path const & file, pslang::ast::location const & location)
{
	if (location.begin.line == location.end.line)
	{
		std::size_t max_line_number_digits = std::floor(std::log10(location.begin.line)) + 1;

		auto line = extract_nth_line(file, location.begin.line);
		auto extra = replace_tabs_with_spaces(line, 2);

		std::cerr << std::endl;
		std::cerr << "line " << std::setw(max_line_number_digits) << std::right << location.begin.line << ": ";
		std::cerr << line << std::endl;
		std::cerr << std::string(std::max<std::size_t>(location.begin.column, 1) - 1 + max_line_number_digits + 7 + extra, ' ') << std::string(location.end.column - location.begin.column, '^') << std::endl;
	}
}

int main(int argc, char ** argv)
{
	if (argc == 1)
	{
		std::cout << "Usage: psli [ options ] <file1> [ <file2> ... ]\n";
		std::cout << "Available options:\n";
		std::cout << "    -t, --trace        Trace each line of execution\n";
		std::cout << "    -d, --dump         Dump all variables after processing all files\n";
		std::cout << "    -p, --print        Print the AST after parsing each file\n";
		std::cout << "    -j, --jit          Just-in-time compile & execute the code instead of interpreting\n";
		return 0;
	}

	using namespace pslang;

	auto context = interpreter::empty_context();

	bool dump = false;
	bool dump_ast = false;
	bool jit = false;

	std::vector<std::string> filenames;
	std::vector<ast::statement_list_ptr> parsed;

	for (int arg = 1; arg < argc; ++arg)
	{
		if (std::strcmp(argv[arg], "-d") == 0 || std::strcmp(argv[arg], "--dump") == 0)
		{
			dump = true;
			continue;
		}

		if (std::strcmp(argv[arg], "-t") == 0 || std::strcmp(argv[arg], "--trace") == 0)
		{
			context.trace = true;
			continue;
		}

		if (std::strcmp(argv[arg], "-p") == 0 || std::strcmp(argv[arg], "--print") == 0)
		{
			dump_ast = true;
			continue;
		}

		if (std::strcmp(argv[arg], "-j") == 0 || std::strcmp(argv[arg], "--jit") == 0)
		{
			jit = true;
			continue;
		}

		try
		{
			filenames.push_back(argv[arg]);
			auto ast = parser::parse(filenames.back());
			ast::resolve_identifiers(ast);
			ast::check_and_infer_types(ast);
			parsed.push_back(std::move(ast));
		}
		catch (ast::parse_error const & error)
		{
			std::cerr << "Parse error at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
		catch (ast::type_error const & error)
		{
			std::cerr << "Type error at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
		catch (ast::invalid_ast_error const & error)
		{
			std::cerr << "Invalid AST at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
		catch (interpreter::internal_error const & error)
		{
			std::cerr << "Internal error at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
		catch (interpreter::runtime_error const & error)
		{
			std::cerr << "Runtime error at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
	}

	if (dump_ast)
	{
		for (std::size_t i = 0; i < filenames.size(); ++i)
		{
			std::cout << "Input file " << filenames[i] << " AST dump:\n\n";
			ast::print(std::cout, *parsed[i]);
			std::cout << "\n";
		}
		std::cout << std::flush;
	}

	if (jit)
	{
		auto abi = jit::host_abi();
		std::vector<jit::compiled_module> modules;
		for (auto const & ast : parsed)
			modules.push_back(jit::make_host_executable(jit::compile(ast, abi)));

		for (auto const & module : modules)
		{
			// TODO: remove, testing-only code; should execute entry point instead
			auto offset = module.code.symbol_table.at("test");
			auto fptr = (bool(*)())(module.code.memory.data.get() + offset);
			auto x = fptr();
			std::cout << "Result: " << std::boolalpha << x << std::endl;
		}
	}
	else
	{
		for (auto const & ast : parsed)
			interpreter::exec(context, ast);

		if (dump)
			interpreter::dump(std::cout, context);
	}
}