package com.zt.plat.module.qms.iot.tcpserver.proxy;

import com.zt.plat.module.qms.iot.tcpserver.IgnorePacket;
import com.zt.plat.module.qms.iot.tcpserver.inf.DecoderFunction;
import com.zt.plat.module.qms.iot.tcpserver.utils.ByteBufferUtil;
import tech.zzjc.tio.core.ChannelContext;
import tech.zzjc.tio.core.Node;
import tech.zzjc.tio.core.exception.TioDecodeException;
import tech.zzjc.tio.core.intf.Packet;

import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;

/**
 * <b>ProxyProtocolDecoder</b>
 * 开启 nginx 代理协议时需要开启，转发代理 ip 信息
 *
 * <p>
 * proxy_protocol on;
 * </p>
 * <p>
 * 更新历史：
 * <pre> 版本         更新时间            更新者        更新内容<hr/>
 * V1.0      2024年9月10日          wxr         Add</pre>
 * <b>Copyright (C)  云南志者竟成科技有限公司</b> 
 * </p>
 * @author 王兴荣<wxr@wangxingrong.com>
 * @version V1.0 
 * @since 2024年9月10日
 */
public class ProxyProtocolDecoder {
	/**
	 * 最小头 “PROXY ” 用来判定是否 v1 的协议
	 */
	private static final int V1_MIN_HEAD_LENGTH = 6;
	/**
	 * Maximum possible length of a v1 proxy message header per spec
	 */
	private static final int V1_MAX_LENGTH = 108;
	/**
	 * 开启 proxy_protocol 的 key
	 */
	private static final String PROXY_PROTOCOL_KEY = "proxy_protocol_key";
	/**
	 * PROXY UNKNOWN\r\n
	 */
	private static final String UNKNOWN = "UNKNOWN";

	private ProxyProtocolDecoder() {
	}

	/**
	 * 开始 proxy message
	 *
	 * @param context ChannelContext
	 */
	public static void enableProxyProtocol(ChannelContext context) {
		context.set(PROXY_PROTOCOL_KEY, (byte) 1);
	}

	/**
	 * 去掉 proxy message
	 *
	 * @param context ChannelContext
	 */
	public static void removeProxyProtocol(ChannelContext context) {
		context.remove(PROXY_PROTOCOL_KEY);
	}

	/**
	 * 判断是否开启 proxy message
	 *
	 * @param context ChannelContext
	 * @return 是否开启
	 */
	public static boolean isProxyProtocolEnabled(ChannelContext context) {
		Object obj = context.get(PROXY_PROTOCOL_KEY);
		if (obj == null) {
			return false;
		}
		return true;
	}

	/**
	 * 解码，如果开启了 proxy message
	 *
	 * @param context        ChannelContext
	 * @param buffer         ByteBuffer
	 * @param readableLength readableLength
	 * @param next           下一个解码器
	 * @return ProxyProtocolMessage
	 * @throws TioDecodeException TioDecodeException
	 */
	public static Packet decodeIfEnable(ChannelContext context, ByteBuffer buffer, int readableLength,
										DecoderFunction next) throws TioDecodeException {
		if (isProxyProtocolEnabled(context)) {
			return decode(context, buffer, readableLength, next);
		} else {
			return next.apply(context, buffer, readableLength);
		}
	}

	/**
	 * 解码 proxy message
	 *
	 * @param context        ChannelContext
	 * @param buffer         ByteBuffer
	 * @param readableLength readableLength
	 * @param next           下一个解码器
	 * @return ProxyProtocolMessage
	 * @throws TioDecodeException TioDecodeException
	 */
	public static Packet decode(ChannelContext context, ByteBuffer buffer, int readableLength, DecoderFunction next) throws TioDecodeException {
		// 如果小于最小长度，尝试解析下一个
		if (readableLength < V1_MIN_HEAD_LENGTH) {
			return next.apply(context, buffer, readableLength);
		}
		// 标记
		buffer.mark();
		// PROXY TCP4 192.168.0.1 192.168.0.11 56324 443\r\n
		String proxyPrefix = ByteBufferUtil.readString(buffer, V1_MIN_HEAD_LENGTH, StandardCharsets.US_ASCII);
		// 非 PROXY 协议，直接返回
		if (!"PROXY ".equals(proxyPrefix)) {
			// 清除协议 key，重置 buffer
			context.remove(PROXY_PROTOCOL_KEY);
			buffer.reset();
			return next.apply(context, buffer, readableLength);
		}
		// 解析 proxy 协议
		ProxyProtocolMessage message = decodeMessage(buffer, readableLength);
		// 半包的情况
		if (message == null) {
			return null;
		}
		// 清除协议 key
		context.remove(PROXY_PROTOCOL_KEY);
		// 设置客户端代理节点
		String protocol = message.getProtocol();
		if (UNKNOWN.equals(protocol)) {
			context.setProxyClientNode(new Node(UNKNOWN, message.getDestinationPort()));
		} else {
			context.setClientNode(new Node(message.getSourceAddress(), message.getSourcePort()));
			context.setProxyClientNode(new Node(message.getDestinationAddress(), message.getDestinationPort()));
		}
		if (buffer.hasRemaining()) {
			return next.apply(context, buffer, readableLength);
		} else {
			return IgnorePacket.INSTANCE;
		}
	}

	/**
	 * 解码 proxy message
	 *
	 * @param buffer         ByteBuffer
	 * @param readableLength readableLength
	 * @return ProxyProtocolMessage
	 * @throws TioDecodeException TioDecodeException
	 */
	public static ProxyProtocolMessage decodeForTest(ByteBuffer buffer, int readableLength) throws TioDecodeException {
		// PROXY TCP4 192.168.0.1 192.168.0.11 56324 443\r\n
		String proxyPrefix = ByteBufferUtil.readString(buffer, V1_MIN_HEAD_LENGTH, StandardCharsets.US_ASCII);
		// 非 PROXY 协议，直接返回
		if (!"PROXY ".equals(proxyPrefix)) {
			throw new TioDecodeException("unknown identifier: " + proxyPrefix);
		}
		return decodeMessage(buffer, readableLength);
	}

	/**
	 * 解码 proxy message
	 *
	 * @param buffer         ByteBuffer
	 * @param readableLength readableLength
	 * @return ProxyProtocolMessage
	 * @throws TioDecodeException TioDecodeException
	 */
	public static ProxyProtocolMessage decodeMessage(ByteBuffer buffer, int readableLength) throws TioDecodeException {
		int endOfLine = findEndOfLine(buffer);
		// 判断超长的情况，有可能是半包，多次进入
		if (endOfLine > V1_MAX_LENGTH || (readableLength > V1_MAX_LENGTH && endOfLine == -1)) {
			throw new TioDecodeException("Error v1 proxy message, readableLength: " + readableLength);
		}
		// 有可能半包，所以返回 null
		if (endOfLine == -1) {
			return null;
		}
		// PROXY TCP4 192.168.0.1 192.168.0.11 56324 443\r\n 去除前缀 PROXY
		// TCP4 192.168.0.1 192.168.0.11 56324 443\r\n
		String header = ByteBufferUtil.readString(buffer, endOfLine - V1_MIN_HEAD_LENGTH, StandardCharsets.US_ASCII);
		// 跳过 \r\n
		ByteBufferUtil.skipBytes(buffer, 2);
		String[] parts = header.split(" ");
		int numParts = parts.length;
		if (numParts < 1) {
			throw new TioDecodeException("invalid header: PROXY " + header + " (expected: 'PROXY' and proxied message values)");
		}
		String proxyProtocol = parts[0];
		if (!"TCP4".equals(proxyProtocol) && !"TCP6".equals(proxyProtocol) && !UNKNOWN.equals(proxyProtocol)) {
			throw new TioDecodeException("unsupported v1 proxy message: " + proxyProtocol);
		}
		if (UNKNOWN.equals(proxyProtocol)) {
			return unknownMsg();
		}
		if (numParts != 5) {
			throw new TioDecodeException("invalid TCP4/6 header: PROXY " + header + " (expected: 6 parts)");
		}
		return new ProxyProtocolMessage(proxyProtocol, parts[1], parts[2], parts[3], parts[4]);
	}

	/**
	 * Proxy message message for 'UNKNOWN' proxied protocols. Per spec, when the proxied message is
	 * 'UNKNOWN' we must discard all other header values.
	 */
	private static ProxyProtocolMessage unknownMsg() {
		return new ProxyProtocolMessage(UNKNOWN, null, null, 0, 0);
	}

	/**
	 * Returns the index in the buffer of the end of line found.
	 * Returns -1 if no end of line was found in the buffer.
	 */
	private static int findEndOfLine(final ByteBuffer buffer) {
		final int n = buffer.limit();
		for (int i = buffer.position(); i < n; i++) {
			final byte b = buffer.get(i);
			if (b == '\r' && i < n - 1 && buffer.get(i + 1) == '\n') {
				return i;  // \r\n
			}
		}
		return -1;  // Not found.
	}
}
