#include "classic_stream.h"

#include <arpa/inet.h>
#include <cerrno>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <fcntl.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <sys/socket.h>
#include <unistd.h>

namespace mtnet::classic {

namespace {
constexpr size_t RECV_CHUNK = 32 * 1024;

uint32_t now_ms() {
	using namespace std::chrono;
	return static_cast<uint32_t>(
			duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count());
}

void set_nonblocking(int fd) {
	int fl = fcntl(fd, F_GETFL, 0);
	fcntl(fd, F_SETFL, fl | O_NONBLOCK);
}

bool plausible_character_add(const uint8_t *p, size_t available) {
	if (available < sizeof(GCCharacterAdd) || p[0] != HDR_GC_CHARACTER_ADD) {
		return false;
	}
	GCCharacterAdd packet{};
	std::memcpy(&packet, p, sizeof(packet));
	return packet.vid != 0 && packet.race != 0 && packet.type <= 8 &&
			std::isfinite(packet.angle) && std::fabs(packet.angle) <= 360.0f &&
			packet.moving_speed <= 255 && packet.attack_speed <= 255;
}
} // namespace

ClassicStream::~ClassicStream() {
	disconnect();
}

void ClassicStream::set_state(State s) {
	if (m_state == s) {
		return;
	}
	m_state = s;
	if (on_state_change) {
		on_state_change(s);
	}
}

// --------------------------------------------------------------------- socket
bool ClassicStream::connect(const std::string &host, uint16_t port) {
	disconnect();

	addrinfo hints{};
	hints.ai_family = AF_INET;
	hints.ai_socktype = SOCK_STREAM;
	addrinfo *res = nullptr;
	char portbuf[8];
	std::snprintf(portbuf, sizeof(portbuf), "%u", port);
	if (getaddrinfo(host.c_str(), portbuf, &hints, &res) != 0 || !res) {
		m_last_error = "getaddrinfo(" + host + ") failed";
		return false;
	}

	m_sock = ::socket(res->ai_family, res->ai_socktype, res->ai_protocol);
	if (m_sock < 0) {
		m_last_error = std::string("socket: ") + std::strerror(errno);
		freeaddrinfo(res);
		return false;
	}
	set_nonblocking(m_sock);
	int one = 1;
	setsockopt(m_sock, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));

	int rc = ::connect(m_sock, res->ai_addr, res->ai_addrlen);
	freeaddrinfo(res);
	if (rc == 0) {
		set_state(State::Online);
	} else if (errno == EINPROGRESS || errno == EWOULDBLOCK) {
		set_state(State::Connecting);
	} else {
		m_last_error = std::string("connect: ") + std::strerror(errno);
		::close(m_sock);
		m_sock = -1;
		return false;
	}
	return true;
}

void ClassicStream::disconnect() {
	if (m_sock >= 0) {
		::close(m_sock);
		m_sock = -1;
	}
	m_recv.clear();
	m_send.clear();
	m_seq_idx = 0;
	m_seq_on = false;
	m_handshake_seen = false;
	m_phase = PHASE_HANDSHAKE;
	m_cipher.clean_up();
	set_state(State::Offline);
}

void ClassicStream::discard_recv(size_t n) {
	m_recv.discard(n);
}

void ClassicStream::decrypt_appended(size_t n) {
	if (n == 0 || !m_cipher.activated()) {
		return;
	}
	// the n bytes just written sit at the tail of the unread region
	m_cipher.decrypt(m_recv.mutable_unread() + (m_recv.readable() - n), n);
}

bool ClassicStream::recv_into_buffer() {
	uint8_t *dst = m_recv.reserve_write(RECV_CHUNK);
	ssize_t n = ::recv(m_sock, dst, RECV_CHUNK, 0);
	if (n > 0) {
		m_recv.commit_write(static_cast<size_t>(n), RECV_CHUNK);
		decrypt_appended(static_cast<size_t>(n));
		return true;
	}
	m_recv.commit_write(0, RECV_CHUNK);
	if (n == 0) {
		m_last_error = "peer closed";
		return false;
	}
	if (errno == EAGAIN || errno == EWOULDBLOCK) {
		return true;
	}
	m_last_error = std::string("recv: ") + std::strerror(errno);
	return false;
}

bool ClassicStream::flush_send() {
	while (m_send.readable() > 0) {
		ssize_t n = ::send(m_sock, m_send.read_ptr(), m_send.readable(), 0);
		if (n > 0) {
			m_send.discard(static_cast<size_t>(n));
			continue;
		}
		if (n < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
			break;
		}
		m_last_error = std::string("send: ") + std::strerror(errno);
		return false;
	}
	return true;
}

void ClassicStream::process() {
	if (m_sock < 0) {
		return;
	}
	if (m_state == State::Connecting) {
		int err = 0;
		socklen_t len = sizeof(err);
		if (getsockopt(m_sock, SOL_SOCKET, SO_ERROR, &err, &len) < 0 || err != 0) {
			m_last_error = "connect failed";
			disconnect();
			return;
		}
		set_state(State::Online);
	}
	if (!recv_into_buffer()) {
		disconnect();
		return;
	}
	dispatch();
	if (m_sock >= 0 && !flush_send()) {
		disconnect();
	}
}

// ----------------------------------------------------------------- framing core
void ClassicStream::feed(const void *data, size_t n) {
	m_recv.write(data, n);
	decrypt_appended(n);
	dispatch();
}

void ClassicStream::emit_bytes(const void *p, size_t n) {
	if (!m_cipher.activated()) {
		m_send.write(p, n);
		return;
	}
	// CTR stream: encrypt in call order; the encoder keeps its counter across calls.
	std::string tmp(static_cast<const char *>(p), n);
	m_cipher.encrypt(tmp.data(), n);
	m_send.write(tmp.data(), n);
}

size_t ClassicStream::take_outgoing(void *dst, size_t cap) {
	size_t n = m_send.readable() < cap ? m_send.readable() : cap;
	if (n) {
		std::memcpy(dst, m_send.read_ptr(), n);
		m_send.discard(n);
	}
	return n;
}

void ClassicStream::append_sequence_if_needed(uint8_t header) {
	if (!m_seq_on || !is_sequence_cg(header)) {
		return;
	}
	uint8_t seq = SEQUENCE_TABLE[m_seq_idx];
	m_seq_idx = (m_seq_idx + 1) % SEQUENCE_TABLE_SIZE;
	emit_bytes(&seq, 1);
}

bool ClassicStream::send_fixed(const void *struct_bytes, size_t n) {
	if (n == 0) {
		return false;
	}
	uint8_t header = *static_cast<const uint8_t *>(struct_bytes);
	emit_bytes(struct_bytes, n);
	append_sequence_if_needed(header);
	if (m_trace) {
		std::fprintf(stderr, "[classic] send hdr=%u n=%zu seq=%d\n", header, n,
				(m_seq_on && is_sequence_cg(header)));
	}
	return true;
}

bool ClassicStream::send_dynamic(const void *bytes, size_t n) {
	if (n < sizeof(DynHead)) {
		return false;
	}
	uint8_t header = *static_cast<const uint8_t *>(bytes);
	emit_bytes(bytes, n);
	append_sequence_if_needed(header);
	return true;
}

// ----------------------------------------------------------------- dispatch
bool ClassicStream::handle_control(uint8_t header, bool &consumed) {
	consumed = false;
	switch (header) {
		case HDR_HANDSHAKE: {
			Handshake hs{};
			if (!m_recv.peek(&hs, sizeof(hs))) {
				return true; // need more
			}
			discard_recv(sizeof(hs));
			consumed = true;
			// mirror EterLib/PythonNetworkStreamPhaseHandshake.cpp + desc.cpp:
			m_server_time_base = hs.time + static_cast<uint32_t>(hs.delta);
			m_client_time_base = now_ms();
			hs.time = hs.time + 2u * static_cast<uint32_t>(hs.delta);
			hs.delta = 0;
			const bool initial = !m_handshake_seen;
			m_handshake_seen = true;
			if (!initial && m_time_sync_mode) {
				// PythonNetworkStreamPhaseHandShake.cpp uses 0xFF only for the
				// first handshake. Select/Game phase resyncs are sent as the
				// same struct with header 0xFC and the next sequence byte.
				hs.header = HDR_CG_TIME_SYNC;
				send_fixed(&hs, sizeof(hs));
			} else {
				// Initial handshake (and the mark side's repeated login
				// handshake) stays 0xFF and has no sequence byte.
				emit_bytes(&hs, sizeof(hs));
			}
			if (on_server_handshake) {
				on_server_handshake();
			}
			if (m_trace) {
				std::fprintf(stderr, "[classic] handshake echo time=%u\n", hs.time);
			}
			return true;
		}
		case HDR_GC_TIME_SYNC: { // 0xFC — 1-byte "handshake ok" blank (re-sync path)
			Blank b{};
			if (!m_recv.peek(&b, sizeof(b))) {
				return true;
			}
			discard_recv(sizeof(b));
			consumed = true;
			if (on_handshake_ok) {
				on_handshake_ok();
			}
			return true;
		}
		case HDR_GC_PHASE: {
			Phase_ p{};
			if (!m_recv.peek(&p, sizeof(p))) {
				return true;
			}
			discard_recv(sizeof(p));
			consumed = true;
			m_phase = p.phase;
			if (m_trace) {
				std::fprintf(stderr, "[classic] phase -> %u\n", p.phase);
			}
			if (on_phase) {
				on_phase(p.phase);
			}
			return true;
		}
		case HDR_GC_PING: {
			Blank b{};
			if (!m_recv.peek(&b, sizeof(b))) {
				return true;
			}
			discard_recv(sizeof(b));
			consumed = true;
			uint8_t pong = HDR_CG_PONG; // 0xFE
			send_fixed(&pong, 1);        // CG_PONG is bSeq=true (sizeof(BYTE)+seq)
			return true;
		}
		case HDR_GC_BINDUDP: {
			BindUDP u{};
			if (!m_recv.peek(&u, sizeof(u))) {
				return true;
			}
			discard_recv(sizeof(u));
			consumed = true;
			return true; // UDP not used
		}
		case HDR_KEY_AGREEMENT: { // 0xFB — GC_KEY_AGREEMENT (server's DH2 blob)
			if (m_recv.readable() < sizeof(KeyAgreement)) {
				return true; // need the whole 261-byte packet (consumed stays false)
			}
			consumed = true;
			return handle_key_agreement();
		}
		case HDR_GC_KEY_AGREEMENT_COMPLETED: { // 0xFA
			KeyAgreementCompleted kac{};
			if (!m_recv.peek(&kac, sizeof(kac))) {
				return true;
			}
			discard_recv(sizeof(kac));
			consumed = true;
			m_cipher.set_activated(true);
			// the server may have pipelined encrypted bytes right behind this;
			// decrypt whatever's already buffered before we dispatch it.
			if (m_recv.readable() > 0) {
				m_cipher.decrypt(m_recv.mutable_unread(), m_recv.readable());
			}
			if (m_trace) {
				std::fprintf(stderr, "[classic] cipher activated\n");
			}
			if (on_cipher_active) {
				on_cipher_active();
			}
			return true;
		}
		default:
			return true; // not a control header
	}
}

bool ClassicStream::handle_key_agreement() {
	KeyAgreement pkt{};
	if (!m_recv.peek(&pkt, sizeof(pkt))) {
		return true; // need the whole 261-byte packet
	}
	discard_recv(sizeof(pkt));

	// generate our own DH2 public blob
	KeyAgreement out{};
	out.header = HDR_KEY_AGREEMENT; // CG_KEY_AGREEMENT (same 0xFB)
	size_t data_len = sizeof(out.data);
	size_t agreed = m_cipher.prepare(out.data, &data_len);
	if (agreed == 0) {
		m_last_error = "classic cipher: DH2 Prepare failed";
		if (on_error) {
			on_error(m_last_error);
		}
		return false;
	}
	out.agreed_length = static_cast<uint16_t>(agreed);
	out.data_length = static_cast<uint16_t>(data_len);

	// derive keys from the peer's blob (does NOT flip activated_ yet)
	if (!m_cipher.activate(m_polarity, pkt.agreed_length, pkt.data, pkt.data_length)) {
		m_last_error = "classic cipher: DH2 Agree/SetUp failed";
		if (on_error) {
			on_error(m_last_error);
		}
		return false;
	}

	// reply with our blob — still plaintext (peer activates on its own COMPLETED)
	emit_bytes(&out, sizeof(out));
	if (m_trace) {
		std::fprintf(stderr, "[classic] key agreement: replied, keys ready\n");
	}
	return true;
}

void ClassicStream::dispatch() {
	for (;;) {
		if (m_recv.readable() == 0) {
			return;
		}
		uint8_t header = *m_recv.read_ptr();

		// zero padding between packets (cipher block alignment on the m2dev side;
		// harmless to skip here too).
		if (header == 0) {
			discard_recv(1);
			continue;
		}

		bool consumed = false;
		if (!handle_control(header, consumed)) {
			if (on_error && m_last_error.empty()) {
				on_error("dispatch aborted");
			}
			disconnect();
			return;
		}
		if (consumed) {
			continue;
		}
		// control handler said "need more bytes" for a control header it owns?
		if (header == HDR_HANDSHAKE || header == HDR_GC_TIME_SYNC || header == HDR_GC_PHASE ||
				header == HDR_GC_PING || header == HDR_GC_BINDUDP || header == HDR_KEY_AGREEMENT ||
				header == HDR_GC_KEY_AGREEMENT_COMPLETED) {
			return; // wait for the rest of this control packet
		}

		// Some side connections share the classic cipher/control handshake but
		// have their own packet framing.  Give that connection a chance before
		// the normal GC static/dynamic table is consulted.
		if (on_raw_packet) {
			RawPacketResult raw = on_raw_packet(header, m_recv.read_ptr(), m_recv.readable());
			switch (raw.status) {
				case RawPacketStatus::NeedMore:
					return;
				case RawPacketStatus::Consumed:
					if (raw.bytes == 0 || raw.bytes > m_recv.readable()) {
						m_last_error = "raw packet consumed invalid byte count";
						if (on_error) {
							on_error(m_last_error);
						}
						disconnect();
						return;
					}
					discard_recv(raw.bytes);
					continue;
				case RawPacketStatus::Error:
					m_last_error = "raw packet framing error";
					if (on_error) {
						on_error(m_last_error);
					}
					disconnect();
					return;
				case RawPacketStatus::NotHandled:
					break;
			}
		}

		// --- dynamic-size GC packet ---
		if (is_dynamic_gc(header)) {
			DynHead dh{};
			if (!m_recv.peek(&dh, sizeof(dh))) {
				return;
			}
			if (dh.size < sizeof(DynHead)) {
				m_last_error = "dynamic packet size underflow";
				if (on_error) {
					on_error(m_last_error);
				}
				disconnect();
				return;
			}
			if (m_recv.readable() < dh.size) {
				return; // whole packet not here yet
			}
			m_recv.discard(sizeof(DynHead));
			uint32_t body_len = dh.size - sizeof(DynHead);
			const uint8_t *body = m_recv.read_ptr();
			if (m_trace) {
				std::fprintf(stderr, "[classic] recv hdr=%u n=%u dynamic=1\n", header, dh.size);
			}
			bool ok = !on_packet || on_packet(header, body, body_len);
			discard_recv(body_len);
			if (!ok) {
				disconnect();
				return;
			}
			continue;
		}

		// --- static-size GC packet --- (packet_size_gc = FULL size incl. header)
		int total = packet_size_gc(header);
		if (total <= 0) {
			m_last_error = "unknown/unsupported GC header " + std::to_string(header);
			// Some 40250 deployments append one legacy entity packet variant that
			// is not present in the checked-in header map.  It has no length field,
			// but the following normal GC_CHARACTER_ADD stream is self-describing.
			// During PHASE_LOADING, recover only at a strongly validated character
			// boundary; this keeps a deployment-specific spawn packet from aborting
			// the whole login while leaving malformed encrypted data fatal elsewhere.
			if (m_phase == PHASE_LOADING) {
				const uint8_t *bytes = m_recv.read_ptr();
				bool recovered = false;
				for (size_t i = 1; i + sizeof(GCCharacterAdd) <= m_recv.readable(); ++i) {
					if (!plausible_character_add(bytes + i, m_recv.readable() - i)) {
						continue;
					}
					if (m_trace) {
						std::fprintf(stderr, "[classic] recovering loading stream at +%zu\n", i);
					}
					discard_recv(i);
					m_last_error.clear();
					recovered = true;
					break;
				}
				if (recovered) {
					continue;
				}
			}
			if (on_error) {
				on_error(m_last_error);
			}
			disconnect();
			return;
		}
		if (m_recv.readable() < static_cast<size_t>(total)) {
			return;
		}
		const uint8_t *p = m_recv.read_ptr();
		if (m_trace) {
			std::fprintf(stderr, "[classic] recv hdr=%u n=%d dynamic=0\n", header, total);
		}
		bool ok = !on_packet || on_packet(header, p + 1, static_cast<uint32_t>(total - 1));
		discard_recv(static_cast<size_t>(total));
		if (!ok) {
			disconnect();
			return;
		}
	}
}

} // namespace mtnet::classic