package com.lutao.cloud.iot.mqtt.core.server.http.websocket;

import com.lutao.cloud.iot.mqtt.codec.ByteBufferUtil;
import com.lutao.cloud.iot.mqtt.codec.MqttMessage;
import com.lutao.cloud.iot.mqtt.codec.WriteBuffer;
import com.lutao.cloud.iot.mqtt.core.server.MqttServerCreator;
import org.tio.core.ChannelContext;
import org.tio.core.Tio;
import org.tio.core.TioConfig;
import org.tio.core.intf.AioHandler;
import org.tio.core.intf.Packet;
import org.tio.http.common.HttpRequest;
import org.tio.http.common.HttpResponse;
import org.tio.websocket.common.WsRequest;
import org.tio.websocket.common.WsResponse;
import org.tio.websocket.server.handler.IWsMsgHandler;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;

/**
 * mqtt websocket 消息处理
 *

 */
public class MqttWsMsgHandler implements IWsMsgHandler {
	/**
	 * mqtt websocket message body key
	 */
	private static final String MQTT_WS_MSG_BODY_KEY = "MQTT_WS_MSG_BODY_KEY";
	/**
	 * MqttServerCreator
	 */
	private final MqttServerCreator serverCreator;
	/**
	 * websocket 握手端点
	 */
	private final String[] supportedSubProtocols;
	private final AioHandler mqttServerAioHandler;

	public MqttWsMsgHandler(MqttServerCreator serverCreator, AioHandler aioHandler) {
		this(serverCreator, new String[]{"mqtt", "mqttv3.1", "mqttv3.1.1"}, aioHandler);
	}

	public MqttWsMsgHandler(MqttServerCreator serverCreator,
							String[] supportedSubProtocols,
							AioHandler aioHandler) {
		this.serverCreator = serverCreator;
		this.supportedSubProtocols = supportedSubProtocols;
		this.mqttServerAioHandler = aioHandler;
	}

	@Override
	public String[] getSupportedSubProtocols() {
		return this.supportedSubProtocols;
	}

	@Override
	public HttpResponse handshake(HttpRequest request, HttpResponse httpResponse, ChannelContext channelContext) {
		if (serverCreator.isWebsocketEnable()) {
			return httpResponse;
		}
		return null;
	}

	/**
	 * 握手后处理
	 *
	 * @param request  HttpRequest
	 * @param response HttpResponse
	 * @param context  ChannelContext
	 */
	@Override
	public void onAfterHandshaked(HttpRequest request, HttpResponse response, ChannelContext context) {
		// 在连接中添加 WriteBuffer 用来处理半包消息
		WriteBuffer wsBody = (WriteBuffer) context.get(MQTT_WS_MSG_BODY_KEY);
		if (wsBody == null) {
			wsBody = new WriteBuffer();
			context.set(MQTT_WS_MSG_BODY_KEY, wsBody);
		}
	}

	/**
	 * 字节消息（binaryType = arraybuffer）过来后会走这个方法
	 */
	@Override
	public Object onBytes(WsRequest wsRequest, byte[] bytes, ChannelContext context) throws Exception {
		WriteBuffer wsBody = (WriteBuffer) context.get(MQTT_WS_MSG_BODY_KEY);
		ByteBuffer buffer = getMqttBody(wsBody, bytes);
		if (buffer == null) {
			return null;
		}
		// 可能会一次有多个包，所以需要进行拆包
		List<Packet> packets = new ArrayList<>(); // 创建一个列表来存储解析出的Packet对象
		while (buffer.hasRemaining()) {
			// 解析 mqtt 消息
			Packet packet = mqttServerAioHandler.decode(buffer, 0, 0, buffer.remaining(), context);
			if (packet == null) {
				// 如果拆包之后还有剩余，写回到 WriteBuffer
				int remaining = buffer.remaining();
				if (remaining > 0) {
					byte[] data = new byte[remaining];
					buffer.get(data);
					wsBody.writeBytes(data);
				}
				// 如果无法解析出Packet，抛出异常或返回一个特定的值
				throw new Exception("Unable to decode MQTT packet");
			} else {
				// 成功解析出Packet，添加到列表中
				packets.add(packet);
			}
		}
		// 如果有解析出的Packet，返回第一个；否则返回null
		return packets.isEmpty() ? null : packets.get(0);
	}

	@Override
	public WsResponse encodeSubProtocol(Packet packet, TioConfig tioConfig, ChannelContext context) {
		if (packet instanceof MqttMessage) {
			ByteBuffer buffer = mqttServerAioHandler.encode(packet, null, context);
			return WsResponse.fromBytes(buffer.array());
		}
		return null;
	}

	/**
	 * 当客户端发 close flag 时，会走这个方法
	 */
	@Override
	public Object onClose(WsRequest wsRequest, byte[] bytes, ChannelContext context) {
		Tio.remove(context, "Mqtt websocket close.");
		return null;
	}

	/**
	 * 字符消息（binaryType = blob）过来后会走这个方法
	 */
	@Override
	public Object onText(WsRequest wsRequest, String text, ChannelContext context) {
		return null;
	}

	/**
	 * 读取 mqtt 消息体处理半包的情况
	 *
	 * @param bytes 消息类容
	 * @return ByteBuffer
	 */
	private static synchronized ByteBuffer getMqttBody(WriteBuffer wsBody, byte[] bytes) {
		wsBody.writeBytes(bytes);
		int length = wsBody.size();
		if (length < 2) {
			return null;
		}
		ByteBuffer buffer = wsBody.toBuffer();
		int mqttLength = getMqttLength(buffer) + 2;
		if (length < mqttLength) {
			return null;
		}
		// 数据已经读取完毕，此处需要重构
		wsBody.reset();
		// 重置 buffer
		buffer.rewind();
		return buffer;
	}

	/**
	 * 解析 mqtt 消息长度
	 *
	 * @param buffer the buffer to decode from
	 * @return mqtt 消息长度
	 */
	private static int getMqttLength(ByteBuffer buffer) {
		ByteBufferUtil.skipBytes(buffer, 1);
		int remainingLength = 0;
		int multiplier = 1;
		short digit;
		int loops = 0;
		do {
			digit = ByteBufferUtil.readUnsignedByte(buffer);
			remainingLength += (digit & 127) * multiplier;
			multiplier *= 128;
			loops++;
		} while ((digit & 128) != 0 && loops < 4);
		return remainingLength;
	}

}
