package com.ttstd.adbloopback.adb;

import com.ttstd.adbloopback.tls.TlsPskClient;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

/**
 * 无线调试配对：在 TLS-PSK 通道之上，与 adbd pairing 守护进程交换 RSA 公钥。
 *
 * <p>配对成功后，本 App 的公钥被加入 adbd 的授权列表，之后即可用同一密钥对
 * 无线调试主端口发起 ADB 连接（无需电脑 / Shizuku）。
 *
 * <p>线规（AOSP pairing 协议）：每条消息 = 4 字节头 [type(1)][reserved(1)][len(2, 大端)] + payload。
 */
public class AdbPairing {

    // AOSP pairing.h MsgType
    private static final byte MSG_PING = 0;
    private static final byte MSG_PONG = 1;
    private static final byte MSG_PAIRING_REQ = 2;
    private static final byte MSG_PAIRING_RESULT = 3;
    private static final byte MSG_PAIRING_ACK = 4;

    public static void pair(String host, int port, String pairingCode, AdbKeyPair keyPair) throws Exception {
        // 配对码同时作为 TLS-PSK 的 PSK 与 identity（与 AOSP adb pair 一致）
        byte[] psk = pairingCode.getBytes("UTF-8");
        TlsPskClient tls = new TlsPskClient(host, port, pairingCode, psk);
        tls.handshake();

        InputStream in = tls.getInputStream();
        OutputStream out = tls.getOutputStream();

        // 向服务端提交本 App 的 ADB 公钥
        byte[] pubKey = keyPair.getAdbPublicKey().getBytes("UTF-8");
        writePairingMsg(out, MSG_PAIRING_REQ, pubKey);

        // 读取服务端回应（可能先有 Ping，然后 PairingResult）
        boolean paired = false;
        for (int i = 0; i < 4; i++) {
            PairingMsg msg = readPairingMsg(in);
            if (msg.type == MSG_PING) {
                writePairingMsg(out, MSG_PONG, new byte[0]);
            } else if (msg.type == MSG_PAIRING_RESULT) {
                paired = true;
                break;
            } else {
                throw new IOException("无线调试配对收到意外消息类型: " + msg.type);
            }
        }
        if (!paired) {
            throw new IOException("无线调试配对未完成");
        }

        // 回 ACK（部分实现需要）
        try {
            writePairingMsg(out, MSG_PAIRING_ACK, new byte[0]);
        } catch (IOException ignored) {
        }
    }

    private static void writePairingMsg(OutputStream out, byte type, byte[] payload) throws IOException {
        byte[] header = {type, 0x00, (byte) (payload.length >> 8), (byte) payload.length};
        out.write(header);
        out.write(payload);
        out.flush();
    }

    private static PairingMsg readPairingMsg(InputStream in) throws IOException {
        byte[] header = AdbMessage.readExactly(in, 4);
        byte type = header[0];
        int len = ((header[2] & 0xff) << 8) | (header[3] & 0xff);
        byte[] payload = len > 0 ? AdbMessage.readExactly(in, len) : new byte[0];
        return new PairingMsg(type, payload);
    }

    private static class PairingMsg {
        final byte type;
        @SuppressWarnings("unused")
        final byte[] payload;

        PairingMsg(byte type, byte[] payload) {
            this.type = type;
            this.payload = payload;
        }
    }
}
