#include "wifi_module.hpp"

#include "util.hpp"

#include <Arduino.h>
#include <string.h>

WiFiModule::WiFiModule(uint8_t receive_pin, uint8_t transmit_pin) noexcept
	: _serial(receive_pin, transmit_pin)
{
}

void WiFiModule::begin(size_t baudrate) noexcept
{
	_serial.begin(baudrate);
}

int WiFiModule::get_available() noexcept
{
	return _serial.available();
}

void WiFiModule::reset() noexcept
{
	_send_serial("AT+RST");

	char response[MAX_NETWORK_MODULE_RESPONSE_LENGTH] = "";

	const auto response_status = _read(2000U, response);

	Serial.println(
		response_status == ResponseStatus::OK ? "Reset wifi module"
											  : "Failed to reset wifi module"
	);
}

bool WiFiModule::connect(const char *ssid, const char *password) noexcept
{
	const char cmd[] = "AT+CWJAP";

	auto command_length = strlen(cmd) + strlen(ssid) + strlen(password) + 6U + 1U;

	auto *command = util::malloc<char>(command_length);

	if (command == nullptr)
	{
		Serial.println("Memory allocation failed for creating full connection command");
		return false;
	}

	snprintf(command, command_length, "%s=\"%s\",\"%s\"", cmd, ssid, password);

	const auto send_success = _send_serial(command);

	if (!send_success)
	{
		Serial.println("Failed to send connection command");
		return false;
	}

	free(command);

	char response[MAX_NETWORK_MODULE_RESPONSE_LENGTH] = "";

	const auto response_status = _read(20000U, response);

	if (response_status != ResponseStatus::OK)
	{
		Serial.print("Connection command responded with status ");
		Serial.println(response_status);
		return false;
	}

	return true;
}

void WiFiModule::set_wifi_mode(WifiMode wifi_mode) noexcept
{
	const char cmd[] = "AT+CWMODE_CUR";

	auto command_length = strlen(cmd) + 2U + 1U;

	auto *command = util::malloc<char>(command_length);

	snprintf(command, command_length, "%s=%u", cmd, static_cast<int>(wifi_mode));

	_send_serial(command);

	free(command);

	char response[MAX_NETWORK_MODULE_RESPONSE_LENGTH] = "";

	_read(1500U, response);
}

void WiFiModule::set_multiple_connections_enabled(bool is_enabled) noexcept
{
	const char cmd[] = "AT+CIPMUX";

	auto command_length = strlen(cmd) + 2U + 1U;

	auto *command = util::malloc<char>(command_length);

	snprintf(command, command_length, "%s=%u", cmd, is_enabled ? 1U : 0U);

	_send_serial(command);

	free(command);

	char response[MAX_NETWORK_MODULE_RESPONSE_LENGTH] = "";

	_read(1500U, response);

	// Serial.println(response);
}

void WiFiModule::set_echo_enabled(bool is_enabled) noexcept
{
	const char cmd[] = "ATE";

	auto command_length = strlen(cmd) + 1U + 1U;

	auto *command = util::malloc<char>(command_length);

	snprintf(command, command_length, "%s%u", cmd, is_enabled ? 1U : 0U);

	_send_serial(command);

	free(command);

	char response[MAX_NETWORK_MODULE_RESPONSE_LENGTH] = "";

	const auto response_status = _read(1500U, response);

	if (response_status == ResponseStatus::OK)
	{
		Serial.print("Turned ");
		Serial.print(is_enabled ? "on" : "off");
		Serial.println(" AT commands echo");
	}
	else
	{
		Serial.print("Failed to turn ");
		Serial.print(is_enabled ? "on" : "off");
		Serial.println(" AT commands echo");
	}
}

void WiFiModule::create_tcp_server(size_t port) noexcept
{
	const auto cmd = "AT+CIPSERVER";

	auto command_length = strlen(cmd) + 7U + 1U;

	auto *command = util::malloc<char>(command_length);

	snprintf(command, command_length, "%s=1,%u", cmd, port);

	_send_serial(command);

	free(command);

	char response[MAX_NETWORK_MODULE_RESPONSE_LENGTH] = "";

	_read(4000U, response);

	// Serial.println(response);
}

bool WiFiModule::test() noexcept
{
	const auto send_success = _send_serial("AT");

	if (!send_success)
	{
		return false;
	}

	char response[MAX_NETWORK_MODULE_RESPONSE_LENGTH] = "";

	const auto response_status = _read(5000U, response);

	if (response_status == ResponseStatus::TIMEOUT)
	{
		return false;
	}

	return strcmp(response, "") != 0;
}

const char *WiFiModule::get_local_ip(char *local_ip_out) noexcept
{
	const auto command = "AT+CIFSR";

	const auto send_success = _send_serial(command);

	if (!send_success)
	{
		return local_ip_out;
	}

	auto *buf = util::malloc<char>(strlen(local_ip_out));

	if (buf == nullptr)
	{
		strcpy(local_ip_out, "Memory allocation failure");
		return local_ip_out;
	}

	const auto response_status = _read(10000U, buf);

	if (response_status != ResponseStatus::OK)
	{
		free(buf);

		sprintf(
			local_ip_out,
			"Response status was not OK. Was %d",
			static_cast<int>(response_status)
		);

		return local_ip_out;
	}

	const auto *local_ip_end = strstr(buf, "\r\n+CIFSR:STAMAC");

	if (local_ip_end == nullptr)
	{
		free(buf);
		strcpy(local_ip_out, "Response parsing error");
		return local_ip_out;
	}

	const auto staip_title_length = 16U; // The length of '+CIFSR:STAIP,"'

	util::substr(buf + staip_title_length, local_ip_end, local_ip_out);

	free(buf);

	return local_ip_out;
}

bool WiFiModule::has_incoming_request() noexcept
{
	return get_available() != 0 && _serial.find("+IPD,");
}

size_t WiFiModule::read_incoming_request(char *raw_request_out) noexcept
{
	// Wait for the data buffer a bit
	while (get_available() < 5)
	{
	}

	const auto connection_id = _serial.read() - ASCII_TO_CHAR;

	const auto data_length = _serial.read();

	Serial.print("Data length: ");
	Serial.print(data_length);

	auto read_bytes = 0U;
	const auto start_time = millis();

	while (read_bytes != data_length && (start_time + 10000) > millis())
	{
		if (get_available() == 0)
		{
			continue;
		}

		char character = _read_byte();

		strncat(raw_request_out, &character, 1U);

		read_bytes++;
	}

	// Skip the comma at the beginning
	raw_request_out++;

	return connection_id;
}

bool WiFiModule::close_connection(size_t connection_id) noexcept
{
	const auto cmd = "AT+CIPCLOSE";

	auto command_length = strlen(cmd) + 50U + 1U;

	auto *command = util::malloc<char>(command_length);

	if (command == nullptr)
	{
		Serial.println("Memory allocation failed for creating close connection command");
		return false;
	}

	snprintf(command, command_length, "%s=%u", cmd, connection_id);

	_send_serial(command);

	free(command);

	char response[MAX_NETWORK_MODULE_RESPONSE_LENGTH] = "";

	const auto response_status = _read(4000U, response);

	if (response_status == ResponseStatus::OK)
	{
		Serial.print("Closed connection to ");
		Serial.println(connection_id);
	}
	else
	{
		Serial.print("Failed to close connection to ");
		Serial.println(connection_id);
		return false;
	}

	return true;
}

bool WiFiModule::send(size_t connection_id, const char *data)
{
	const auto cmd = "AT+CIPSEND";

	const auto data_length = strlen(data);

	auto command_length = strlen(cmd) + 50U + data_length + 1U;

	auto *command = util::malloc<char>(command_length);

	if (command == nullptr)
	{
		Serial.println("Memory allocation failed for creating close connection command");
		return false;
	}

	snprintf(command, command_length, "%s=%u,%u", cmd, connection_id, data_length);

	_send_serial(command);

	free(command);

	char response[MAX_NETWORK_MODULE_RESPONSE_LENGTH] = "";

	_read(4000U, response);

	_serial.print(data);

	strcpy(response, "");

	_read(4000U, response);

	return true;
}

bool WiFiModule::_send_serial(const char *command) noexcept
{
	auto full_command_length = strlen(command) + 2U + 1U;

	auto *full_command = util::malloc<char>(full_command_length);

	if (full_command == nullptr)
	{
		return false;
	}

	snprintf(full_command, full_command_length, "%s\r\n", command);

	_serial.print(full_command);

	// Serial.print("Sent command: ");
	// Serial.println(full_command);

	free(full_command);

	return true;
}

ResponseStatus WiFiModule::_read(uint64_t timeout, char *response_out) noexcept
{
	const auto start_time = millis();

	bool has_end = false;
	auto status = ResponseStatus::TIMEOUT;

	// auto *response_out_head = response_out;

	while (!has_end && (start_time + timeout) > millis())
	{
		while (_serial.available() != 0)
		{
			char character = _read_byte();

			strncat(response_out, &character, 1U);

			// *response_out_head = character;
			// response_out_head++;

			if (util::str_ends_with(response_out, "OK"))
			{
				status = ResponseStatus::OK;
				has_end = true;
			}

			if (util::str_ends_with(response_out, "ERROR"))
			{
				status = ResponseStatus::ERROR;
				has_end = true;
			}

			if (util::str_ends_with(response_out, "FAIL"))
			{
				status = ResponseStatus::FAIL;
				has_end = true;
			}
		}
	}

	return status;
}

char WiFiModule::_read_byte() noexcept
{
	return static_cast<char>(_serial.read());
}