package study.iot.support;

import java.io.ByteArrayOutputStream;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

/**
 * 原始 MQTT 5 报文探针：直接读写 TCP socket 上的 MQTT 定长头/剩余长度/报文，
 * 用来观察客户端库不总是暴露的协议细节（CONNACK reason code、SUBACK、DUP 位、
 * QoS 2 的 PUBREC/PUBREL/PUBCOMP 握手）。
 *
 * <p>这是实验工具，不是生产客户端；它只支持本实验用到的报文子集。</p>
 */
public final class RawMqttProbe implements AutoCloseable {

    public static final int CONNACK = 2;
    public static final int PUBLISH = 3;
    public static final int PUBACK = 4;
    public static final int PUBREC = 5;
    public static final int PUBREL = 6;
    public static final int PUBCOMP = 7;
    public static final int SUBACK = 9;

    public record Packet(int type, int flags, byte[] body) {
        public int packetId() {
            if (body.length < 2) {
                return -1;
            }
            return ((body[0] & 0xff) << 8) | (body[1] & 0xff);
        }

        public boolean dup() {
            return (flags & 0x08) != 0;
        }

        public int qos() {
            return (flags >> 1) & 0x03;
        }
    }

    private final Socket socket;
    private final InputStream in;
    private final OutputStream out;

    public RawMqttProbe(String host, int port) throws IOException {
        this.socket = new Socket(host, port);
        this.socket.setSoTimeout(5000);
        this.in = socket.getInputStream();
        this.out = socket.getOutputStream();
    }

    public int connect(String clientId, String username, String password, boolean cleanStart, long sessionExpiry)
            throws IOException {
        return connect(clientId, username, password, cleanStart, sessionExpiry, null, null, 1, false);
    }

    /** 建立连接；可选地在 CONNECT 中携带遗嘱（topic/payload/qos/retain）。 */
    public int connect(String clientId, String username, String password, boolean cleanStart, long sessionExpiry,
                       String willTopic, String willPayload, int willQos, boolean willRetain) throws IOException {
        ByteArrayOutputStream body = new ByteArrayOutputStream();
        body.writeBytes(utf8("MQTT"));
        body.write(5);
        int connectFlags = cleanStart ? 0x02 : 0x00;
        if (willTopic != null) {
            connectFlags |= 0x04 | (willQos << 3) | (willRetain ? 0x20 : 0x00);
        }
        if (username != null) {
            connectFlags |= 0xC0; // username(0x80) + password(0x40)
        }
        body.write(connectFlags);
        body.write(0);
        body.write(60); // keep-alive 60 秒
        ByteArrayOutputStream props = new ByteArrayOutputStream();
        props.write(0x11); // session expiry interval
        props.write((int) (sessionExpiry >> 24) & 0xff);
        props.write((int) (sessionExpiry >> 16) & 0xff);
        props.write((int) (sessionExpiry >> 8) & 0xff);
        props.write((int) sessionExpiry & 0xff);
        writeProperties(body, props.toByteArray());
        body.writeBytes(utf8(clientId));
        if (willTopic != null) {
            ByteArrayOutputStream willProps = new ByteArrayOutputStream();
            writeProperties(body, willProps.toByteArray());
            body.writeBytes(utf8(willTopic));
            body.writeBytes(binary(willPayload.getBytes(StandardCharsets.UTF_8)));
        }
        if (username != null) {
            body.writeBytes(utf8(username));
            body.writeBytes(utf8(password));
        }
        writePacket(0x10, body.toByteArray());
        Packet connack = read();
        require(connack.type() == CONNACK, "expected CONNACK, got type " + connack.type());
        return connack.body().length >= 2 ? connack.body()[1] & 0xff : -1;
    }

    public List<Integer> subscribe(String topicFilter, int qos, int packetId) throws IOException {
        ByteArrayOutputStream body = new ByteArrayOutputStream();
        body.write(packetId >> 8);
        body.write(packetId & 0xff);
        body.write(0);
        body.writeBytes(utf8(topicFilter));
        body.write(qos);
        writePacket(0x82, body.toByteArray());
        Packet suback = read();
        require(suback.type() == SUBACK, "expected SUBACK, got type " + suback.type());
        List<Integer> codes = new ArrayList<>();
        for (byte code : suback.body()) {
            codes.add(code & 0xff);
        }
        return codes;
    }

    public void publish(String topic, byte[] payload, int qos, boolean retain, int packetId, boolean dup)
            throws IOException {
        ByteArrayOutputStream body = new ByteArrayOutputStream();
        body.writeBytes(utf8(topic));
        if (qos > 0) {
            body.write(packetId >> 8);
            body.write(packetId & 0xff);
        }
        body.write(0);
        body.writeBytes(payload);
        int header = 0x30 | (qos << 1) | (dup ? 0x08 : 0) | (retain ? 0x01 : 0);
        writePacket(header, body.toByteArray());
    }

    public void ack(int type, int packetId) throws IOException {
        writePacket(type, new byte[]{(byte) (packetId >> 8), (byte) (packetId & 0xff), 0});
    }

    public void disconnect() throws IOException {
        writePacket(0xE0, new byte[]{0});
    }

    public Packet read() throws IOException {
        int header = in.read();
        if (header < 0) {
            throw new EOFException("connection closed");
        }
        int remaining = 0;
        int multiplier = 1;
        int value;
        do {
            value = in.read();
            if (value < 0) {
                throw new EOFException("connection closed in remaining length");
            }
            remaining += (value & 0x7f) * multiplier;
            multiplier *= 128;
        } while ((value & 0x80) != 0);
        byte[] body = in.readNBytes(remaining);
        if (body.length != remaining) {
            throw new EOFException("truncated packet");
        }
        return new Packet(header >> 4, header & 0x0f, body);
    }

    /** 读取下一条 PUBLISH；跳过无关的 ACK 类报文不会发生，调用方应按协议顺序读取。 */
    public Packet readPublish() throws IOException {
        Packet packet = read();
        require(packet.type() == PUBLISH, "expected PUBLISH, got type " + packet.type());
        return packet;
    }

    public static String publishPayload(Packet publish) {
        int index = 2 + ((publish.body()[0] & 0xff) << 8 | (publish.body()[1] & 0xff));
        if (publish.qos() > 0) {
            index += 2;
        }
        int propsLength = publish.body()[index] & 0xff;
        index += 1 + propsLength;
        return new String(publish.body(), index, publish.body().length - index, StandardCharsets.UTF_8);
    }

    private void writePacket(int header, byte[] body) throws IOException {
        out.write(header);
        int remaining = body.length;
        do {
            int digit = remaining % 128;
            remaining /= 128;
            if (remaining > 0) {
                digit |= 0x80;
            }
            out.write(digit);
        } while (remaining > 0);
        out.write(body);
        out.flush();
    }

    private static void writeProperties(ByteArrayOutputStream body, byte[] properties) {
        int length = properties.length;
        do {
            int digit = length % 128;
            length /= 128;
            if (length > 0) {
                digit |= 0x80;
            }
            body.write(digit);
        } while (length > 0);
        body.writeBytes(properties);
    }

    /** MQTT Binary Data：两字节长度前缀 + 原始字节。 */
    private static byte[] binary(byte[] payload) {
        byte[] result = new byte[payload.length + 2];
        result[0] = (byte) (payload.length >> 8);
        result[1] = (byte) (payload.length & 0xff);
        System.arraycopy(payload, 0, result, 2, payload.length);
        return result;
    }

    private static byte[] utf8(String value) {
        byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
        byte[] result = new byte[bytes.length + 2];
        result[0] = (byte) (bytes.length >> 8);
        result[1] = (byte) (bytes.length & 0xff);
        System.arraycopy(bytes, 0, result, 2, bytes.length);
        return result;
    }

    private static void require(boolean condition, String message) {
        if (!condition) {
            throw new IllegalStateException(message);
        }
    }

    @Override
    public void close() throws IOException {
        socket.close();
    }
}