TCPClientInterface.h

#pragma once

#if HAS_WIFI && defined(TCP_CLIENT_TRANSPORT) && defined(HAS_RNS)

#include <WiFi.h>
#include <microReticulum/Bytes.h>
#include <microReticulum/Interface.h>
#include <microReticulum/Log.h>
#include <microReticulum/Type.h>

#if defined(ESP32)
  #include <esp_heap_caps.h>
  #include <esp_wifi.h>
  #include <esp_task_wdt.h>
#endif

#define URTN_TCP_RECONNECT_INTERVAL_MS 30000UL
#define URTN_TCP_CONNECT_TIMEOUT_MS 5000
#define URTN_TCP_FRAME_MTU 508
#define URTN_TCP_READ_BUDGET_BYTES 512
#define URTN_TCP_READ_BUDGET_FRAMES 1
#define URTN_TCP_READ_BUDGET_MS 2
#define URTN_TCP_WRITE_CHUNK_SIZE 128
#define URTN_TCP_MIN_PROCESS_HEAP 24576
#define URTN_TCP_CRITICAL_HEAP 16384
#define URTN_TCP_LOW_HEAP_BACKOFF_MS 15000UL
#define URTN_TCP_KEEPALIVE_INTERVAL_MS 60000UL
#define URTN_TCP_WIFI_RECOVERY_FAILURES 12
#define URTN_TCP_WIFI_RECOVERY_INTERVAL_MS 300000UL

#define URTN_TCP_STATE_IDLE 0
#define URTN_TCP_STATE_CONNECTED 1
#define URTN_TCP_STATE_CONNECT_FAILED 2
#define URTN_TCP_STATE_DISCONNECTED 3
#define URTN_TCP_STATE_LOW_HEAP 4
#define URTN_TCP_STATE_WRITE_FAILED 5
#define URTN_TCP_STATE_BAD_ALLOC 6
#define URTN_TCP_STATE_LOW_HEAP_DEFER 7

class TCPClientInterface : public RNS::InterfaceImpl {
public:
	TCPClientInterface(const char *host, uint16_t port) : RNS::InterfaceImpl("TCPClientInterface") {
		memset(_host, 0, sizeof(_host));
		if (host) {
			strncpy(_host, host, sizeof(_host) - 1);
		}
		_port = port;
		_IN = true;
		_OUT = true;
		_HW_MTU = URTN_TCP_FRAME_MTU;
		_bitrate = 10000000;
		snprintf(_display_name, sizeof(_display_name), "TCPClientInterface[%s:%u]", _host, _port);
	}

	virtual ~TCPClientInterface() {
		_client.stop();
		_name = "deleted";
	}

	uint8_t status_code() const {
		if (_online && _last_status == URTN_TCP_STATE_LOW_HEAP_DEFER) {
			return URTN_TCP_STATE_LOW_HEAP_DEFER;
		}
		if (_online) return URTN_TCP_STATE_CONNECTED;
		return _last_status;
	}

	uint32_t last_event_age_s() const {
		if (_last_event_ms == 0) return 0xFFFFFFFFUL;
		return (uint32_t)(millis() - _last_event_ms) / 1000UL;
	}

	uint16_t connect_failures() const { return _connect_failures; }
	uint16_t disconnects() const { return _disconnects; }
	uint16_t low_heap_pauses() const { return _low_heap_pauses; }
	uint16_t wifi_recoveries() const { return _wifi_recoveries; }

protected:
	virtual void loop() {
		if (!heap_available_for_tcp()) {
			if (heap_critical_for_tcp()) {
				pause_low_heap("loop");
			} else {
				defer_low_heap("loop");
			}
			return;
		}

		if (_client.connected()) {
			if (!_online || _last_status == URTN_TCP_STATE_LOW_HEAP_DEFER) {
				_online = true;
				record_status(URTN_TCP_STATE_CONNECTED);
			}
			ensure_rx_frame_capacity();
			read_available();
			send_keepalive_if_due();
			return;
		}

		if (_online) {
			_online = false;
			reset_client();
			release_rx_frame();
			_in_frame = false;
			_escape = false;
			_overflow = false;
			record_status(URTN_TCP_STATE_DISCONNECTED);
			increment_u16(_disconnects);
			WARNINGF("TCP uplink to %s:%u disconnected", _host, _port);
		}

		if (WiFi.status() != WL_CONNECTED) {
			return;
		}

		uint32_t now = millis();
		if ((uint32_t)(now - _last_connect_try) < URTN_TCP_RECONNECT_INTERVAL_MS) {
			return;
		}
		_last_connect_try = now;

		TRACEF("Connecting TCP uplink to %s:%u", _host, _port);
		reset_client();
		if (connect_uplink()) {
			_client.setNoDelay(true);
			_online = true;
			_in_frame = false;
			_escape = false;
			_overflow = false;
			reset_rx_frame();
			_last_tx_ms = millis();
			record_status(URTN_TCP_STATE_CONNECTED);
			TRACEF("TCP uplink to %s:%u connected", _host, _port);
		} else {
			reset_client();
			record_status(URTN_TCP_STATE_CONNECT_FAILED);
			increment_u16(_connect_failures);
			if ((_connect_failures % URTN_TCP_WIFI_RECOVERY_FAILURES) == 0) {
				_remote_ip_resolved = false;
				recover_wifi_after_failures(now);
			}
			WARNINGF("TCP uplink to %s:%u failed", _host, _port);
		}
	}

	virtual bool send_outgoing(const RNS::Bytes& data) {
		if (!_client.connected()) {
			_online = false;
			return false;
		}
		if (!heap_available_for_tcp()) {
			if (heap_critical_for_tcp()) {
				pause_low_heap("send");
			} else {
				defer_low_heap("send");
			}
			return false;
		}

		uint8_t chunk[URTN_TCP_WRITE_CHUNK_SIZE];
		size_t chunk_len = 0;
		auto flush_chunk = [&]() -> bool {
			if (chunk_len == 0) return true;
			size_t requested = chunk_len;
			size_t written = _client.write(chunk, requested);
			chunk_len = 0;
			service_tcp_runtime();
			if (written != requested) {
				reset_client();
				_online = false;
				record_status(URTN_TCP_STATE_WRITE_FAILED);
				return false;
			}
			return true;
		};
		auto write_frame_byte = [&](uint8_t b) -> bool {
			chunk[chunk_len++] = b;
			if (chunk_len >= sizeof(chunk)) {
				return flush_chunk();
			}
			return true;
		};

		if (!write_frame_byte(HDLC_FLAG)) return false;
		for (size_t i = 0; i < data.size(); i++) {
			uint8_t b = data.data()[i];
			if (b == HDLC_FLAG || b == HDLC_ESC) {
				if (!write_frame_byte(HDLC_ESC)) return false;
				if (!write_frame_byte(b ^ HDLC_ESC_MASK)) return false;
			} else {
				if (!write_frame_byte(b)) return false;
			}
		}
		if (!write_frame_byte(HDLC_FLAG)) return false;
		if (!flush_chunk()) return false;
		_last_tx_ms = millis();
		InterfaceImpl::handle_outgoing(data);
		return true;
	}

	virtual std::string toString() const {
		return _display_name;
	}

private:
	static const uint8_t HDLC_FLAG = 0x7E;
	static const uint8_t HDLC_ESC = 0x7D;
	static const uint8_t HDLC_ESC_MASK = 0x20;

	void service_tcp_runtime() {
		yield();
#if defined(ESP32)
		esp_task_wdt_reset();
#endif
	}

	static void increment_u16(uint16_t &value) {
		if (value < 0xFFFF) value++;
	}

	void record_status(uint8_t status) {
		if (status != _last_status || _last_event_ms == 0) {
#if defined(ARDUINO)
#if defined(ESP32)
			Serial.printf("URTNDBG tcp_status=%s code=%u fail=%u disc=%u rec=%u defer=%u short=%u heap=%u min_heap=%u largest=%u rx_len=%u rx_cap=%u avail=%d ms=%lu\r\n",
				tcp_status_name(status),
				status,
				_connect_failures,
				_disconnects,
				_wifi_recoveries,
				_low_heap_defers,
				_short_frames,
				ESP.getFreeHeap(),
				ESP.getMinFreeHeap(),
				heap_caps_get_largest_free_block(MALLOC_CAP_8BIT),
				(unsigned int)_rx_frame.size(),
				(unsigned int)_rx_frame.capacity(),
				_client.connected() ? _client.available() : 0,
				(unsigned long)millis());
#else
			Serial.printf("URTNDBG tcp_status=%s code=%u fail=%u disc=%u rec=%u defer=%u short=%u rx_len=%u rx_cap=%u ms=%lu\r\n",
				tcp_status_name(status),
				status,
				_connect_failures,
				_disconnects,
				_wifi_recoveries,
				_low_heap_defers,
				_short_frames,
				(unsigned int)_rx_frame.size(),
				(unsigned int)_rx_frame.capacity(),
				(unsigned long)millis());
#endif
#endif
		}
		_last_status = status;
		_last_event_ms = millis();
	}

	const char *tcp_status_name(uint8_t status) const {
		switch (status) {
			case URTN_TCP_STATE_IDLE: return "idle";
			case URTN_TCP_STATE_CONNECTED: return "connected";
			case URTN_TCP_STATE_CONNECT_FAILED: return "connect-failed";
			case URTN_TCP_STATE_DISCONNECTED: return "disconnected";
			case URTN_TCP_STATE_LOW_HEAP: return "low-heap";
			case URTN_TCP_STATE_WRITE_FAILED: return "write-failed";
			case URTN_TCP_STATE_BAD_ALLOC: return "bad-alloc";
			case URTN_TCP_STATE_LOW_HEAP_DEFER: return "low-heap-defer";
			default: return "unknown";
		}
	}

	void ensure_rx_frame_capacity() {
		if (_rx_frame.capacity() < _HW_MTU) {
			_rx_frame = RNS::Bytes(_HW_MTU);
			_rx_frame.resize(0);
		}
	}

	void reset_rx_frame() {
		ensure_rx_frame_capacity();
		_rx_frame.resize(0);
	}

	void release_rx_frame() {
		_rx_frame.clear();
	}

	void reset_client() {
		_client.stop();
		_client = WiFiClient();
		_client.setTimeout((URTN_TCP_CONNECT_TIMEOUT_MS + 999) / 1000);
	}

	bool connect_uplink() {
		service_tcp_runtime();
		if (!_remote_ip_resolved) {
			if (WiFi.hostByName(_host, _remote_ip) != 1) {
				service_tcp_runtime();
				return false;
			}
			_remote_ip_resolved = true;
		}
		bool connected = _client.connect(_remote_ip, _port, URTN_TCP_CONNECT_TIMEOUT_MS);
		if (!connected) {
			reset_client();
			service_tcp_runtime();
			connected = _client.connect(_host, _port, URTN_TCP_CONNECT_TIMEOUT_MS);
		}
		service_tcp_runtime();
		return connected;
	}

	void recover_wifi_after_failures(uint32_t now) {
		if (WiFi.status() != WL_CONNECTED) {
			return;
		}
		if (_last_wifi_recovery_ms != 0 && (uint32_t)(now - _last_wifi_recovery_ms) < URTN_TCP_WIFI_RECOVERY_INTERVAL_MS) {
			return;
		}
		_last_wifi_recovery_ms = now;
		increment_u16(_wifi_recoveries);
		_remote_ip_resolved = false;
		reset_client();
		WARNING("TCP uplink forcing WiFi STA reconnect after repeated failures");
#if defined(ESP32)
		esp_err_t err = esp_wifi_disconnect();
		service_tcp_runtime();
		delay(100);
		if (err == ESP_OK || err == ESP_ERR_WIFI_NOT_CONNECT) {
			err = esp_wifi_connect();
		}
		if (err != ESP_OK && err != ESP_ERR_WIFI_CONN) {
			WARNINGF("TCP WiFi recovery failed: 0x%x", err);
		}
#else
		WiFi.disconnect(false);
		service_tcp_runtime();
		delay(100);
		WiFi.reconnect();
#endif
		_last_connect_try = now;
	}

	void send_keepalive_if_due() {
		uint32_t now = millis();
		if (_last_tx_ms != 0 && (uint32_t)(now - _last_tx_ms) < URTN_TCP_KEEPALIVE_INTERVAL_MS) {
			return;
		}
		const uint8_t keepalive[2] = { HDLC_FLAG, HDLC_FLAG };
		size_t written = _client.write(keepalive, sizeof(keepalive));
		service_tcp_runtime();
		if (written == sizeof(keepalive)) {
			_last_tx_ms = now;
			return;
		}
		reset_client();
		_online = false;
		record_status(URTN_TCP_STATE_WRITE_FAILED);
		increment_u16(_disconnects);
	}

	bool heap_available_for_tcp() const {
#if defined(ESP32)
		return ESP.getFreeHeap() >= URTN_TCP_MIN_PROCESS_HEAP;
#else
		return true;
#endif
	}

	bool heap_critical_for_tcp() const {
#if defined(ESP32)
		return ESP.getFreeHeap() < URTN_TCP_CRITICAL_HEAP;
#else
		return false;
#endif
	}

	void log_tcp_heap_event(const char *event, const char *reason, uint32_t now, uint32_t threshold) {
#if defined(ARDUINO)
#if defined(ESP32)
		Serial.printf("URTNDBG tcp_heap=%s reason=%s threshold=%lu heap=%u min_heap=%u largest=%u connected=%u online=%u rx_len=%u rx_cap=%u avail=%d low=%u defer=%u short=%u ms=%lu\r\n",
			event ? event : "-",
			reason ? reason : "-",
			(unsigned long)threshold,
			ESP.getFreeHeap(),
			ESP.getMinFreeHeap(),
			heap_caps_get_largest_free_block(MALLOC_CAP_8BIT),
			_client.connected() ? 1 : 0,
			_online ? 1 : 0,
			(unsigned int)_rx_frame.size(),
			(unsigned int)_rx_frame.capacity(),
			_client.connected() ? _client.available() : 0,
			_low_heap_pauses,
			_low_heap_defers,
			_short_frames,
			(unsigned long)now);
#else
		Serial.printf("URTNDBG tcp_heap=%s reason=%s threshold=%lu connected=%u online=%u rx_len=%u rx_cap=%u low=%u defer=%u short=%u ms=%lu\r\n",
			event ? event : "-",
			reason ? reason : "-",
			(unsigned long)threshold,
			_client.connected() ? 1 : 0,
			_online ? 1 : 0,
			(unsigned int)_rx_frame.size(),
			(unsigned int)_rx_frame.capacity(),
			_low_heap_pauses,
			_low_heap_defers,
			_short_frames,
			(unsigned long)now);
#endif
#else
		(void)event;
		(void)reason;
		(void)now;
		(void)threshold;
#endif
	}

	void defer_low_heap(const char *reason) {
		uint32_t now = millis();
		if ((uint32_t)(now - _last_low_heap_notice) >= URTN_TCP_LOW_HEAP_BACKOFF_MS) {
			log_tcp_heap_event("defer", reason, now, URTN_TCP_MIN_PROCESS_HEAP);
			_last_low_heap_notice = now;
		}
		increment_u16(_low_heap_defers);
		record_status(URTN_TCP_STATE_LOW_HEAP_DEFER);
		service_tcp_runtime();
	}

	void pause_low_heap(const char *reason) {
		uint32_t now = millis();
		log_tcp_heap_event("hard_pause", reason, now, URTN_TCP_CRITICAL_HEAP);
		if ((uint32_t)(now - _last_low_heap_notice) >= URTN_TCP_LOW_HEAP_BACKOFF_MS) {
#if defined(ESP32)
			WARNINGF("TCP uplink paused, low heap: %u bytes", ESP.getFreeHeap());
#else
			WARNING("TCP uplink paused, low heap");
#endif
			_last_low_heap_notice = now;
		}
		if (_client.connected()) {
			reset_client();
			_online = false;
			release_rx_frame();
			_in_frame = false;
			_escape = false;
			_overflow = false;
		}
		record_status(URTN_TCP_STATE_LOW_HEAP);
		increment_u16(_low_heap_pauses);
		_last_connect_try = now;
		service_tcp_runtime();
	}

	void read_available() {
		uint32_t start = millis();
		uint16_t bytes_read = 0;
		uint8_t frames_read = 0;

		while (_client.connected() && _client.available() > 0) {
			if (bytes_read >= URTN_TCP_READ_BUDGET_BYTES || frames_read >= URTN_TCP_READ_BUDGET_FRAMES || (uint32_t)(millis() - start) >= URTN_TCP_READ_BUDGET_MS) {
				return;
			}

			int value = _client.read();
			if (value < 0) {
				return;
			}
			bytes_read++;
			uint8_t b = (uint8_t)value;
			if (b == HDLC_FLAG) {
				if (_in_frame && !_overflow && _rx_frame.size() > RNS::Type::Reticulum::HEADER_MINSIZE) {
					if (!heap_available_for_tcp()) {
						if (heap_critical_for_tcp()) {
							pause_low_heap("rx-frame");
						} else {
							defer_low_heap("rx-frame");
						}
						return;
					}
					try {
						InterfaceImpl::handle_incoming(_rx_frame);
						frames_read++;
						service_tcp_runtime();
					}
					catch (const std::bad_alloc&) {
						ERROR("TCPClientInterface::handle_incoming: bad_alloc - out of memory");
						record_status(URTN_TCP_STATE_BAD_ALLOC);
						pause_low_heap("rx-bad-alloc");
						return;
					}
					catch (std::exception& e) {
						ERRORF("TCPClientInterface::handle_incoming: %s", e.what());
						service_tcp_runtime();
					}
				} else if (_in_frame && !_overflow && _rx_frame.size() > 0) {
					increment_u16(_short_frames);
				}
				reset_rx_frame();
				_in_frame = true;
				_escape = false;
				_overflow = false;
				continue;
			}

			if (!_in_frame || _overflow) {
				continue;
			}

			if (b == HDLC_ESC) {
				_escape = true;
				continue;
			}

			if (_escape) {
				b ^= HDLC_ESC_MASK;
				_escape = false;
			}

			if (_rx_frame.size() < _HW_MTU) {
				ensure_rx_frame_capacity();
				_rx_frame.append(b);
			} else {
				_overflow = true;
				reset_rx_frame();
				WARNING("Dropping oversized TCP frame");
			}
		}
	}

	WiFiClient _client;
	IPAddress _remote_ip;
	char _host[RNS_CFG_HOST_SIZE];
	char _display_name[RNS_CFG_HOST_SIZE + 32];
	uint16_t _port = 0;
	uint32_t _last_connect_try = 0;
	uint32_t _last_low_heap_notice = 0;
	uint32_t _last_tx_ms = 0;
	uint32_t _last_event_ms = 0;
	uint32_t _last_wifi_recovery_ms = 0;
	uint16_t _connect_failures = 0;
	uint16_t _disconnects = 0;
	uint16_t _low_heap_pauses = 0;
	uint16_t _low_heap_defers = 0;
	uint16_t _short_frames = 0;
	uint16_t _wifi_recoveries = 0;
	uint8_t _last_status = URTN_TCP_STATE_IDLE;
	bool _remote_ip_resolved = false;
	bool _in_frame = false;
	bool _escape = false;
	bool _overflow = false;
	RNS::Bytes _rx_frame;
};

#endif