package com.ttstd.adbloopback.tls;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.security.SecureRandom;
import java.util.Arrays;

/**
 * 最小化 TLS 1.2 PSK 客户端，仅支持 {@code TLS_PSK_WITH_AES_128_GCM_SHA256}。
 *
 * <p>Android 11+ 的「无线调试 → 使用配对码配对设备」走的就是这条 TLS-PSK 通道：
 * 配对码同时作为 PSK 与 PSK identity。标准 Android (JSSE / Conscrypt) 不暴露 PSK 接口，
 * 故这里用 Android 自带密码学原语（AES-GCM / HMAC / SHA-256）手搓实现。
 *
 * <p>握手完成后，通过 {@link #getInputStream()}/{@link #getOutputStream()} 获得透明加密的流，
 * 上层即可在其上跑 adb 的配对公钥交换协议。
 */
public class TlsPskClient {

    // 仅支持这一个密码套件（与 AOSP adbd pairing 一致）
    private static final int CIPHER_PSK_AES128_GCM_SHA256 = 0x00A8;

    private final Socket socket;
    private InputStream in;
    private OutputStream out;
    private final byte[] psk;
    private final byte[] identity;

    private final byte[] clientRandom = new byte[32];
    private final byte[] serverRandom = new byte[32];
    private byte[] clientWriteKey, serverWriteKey, clientWriteIV, serverWriteIV;
    private byte[] masterSecret;
    private long clientSeq = 0;
    private long serverSeq = 0;

    private final ByteArrayOutputStream handshakeBuf = new ByteArrayOutputStream();

    public TlsPskClient(String host, int port, String identityStr, byte[] psk) throws IOException {
        this.psk = psk;
        this.identity = identityStr.getBytes("UTF-8");
        this.socket = new Socket();
        this.socket.connect(new InetSocketAddress(host, port), 8000);
        this.socket.setSoTimeout(8000);
        this.socket.setTcpNoDelay(true);
        this.in = socket.getInputStream();
        this.out = socket.getOutputStream();
        new SecureRandom().nextBytes(clientRandom);
    }

    public InputStream getInputStream() {
        return new TlsInputStream();
    }

    public OutputStream getOutputStream() {
        return new TlsOutputStream();
    }

    public void handshake() throws Exception {
        sendClientHello();
        readServerHelloAndDone();
        deriveKeys();
        sendClientKeyExchange();
        sendChangeCipherSpec();
        sendFinished();
        readServerFinished();
    }

    // ---------- 握手各阶段 ----------

    private void sendClientHello() throws Exception {
        byte[] sigAlgs = {
                0x00, 0x0c,                     // 签名算法列表长度 = 12
                0x04, 0x03,                     // sha256 + ecdsa
                0x05, 0x03,                     // sha384 + ecdsa
                0x06, 0x03,                     // sha512 + ecdsa
                0x04, 0x01,                     // sha256 + rsa
                0x05, 0x01,                     // sha384 + rsa
                0x06, 0x01,                     // sha512 + rsa
                0x02, 0x01                      // sha1 + rsa
        };
        byte[] ext = concat(shortBytes(0x000d), shortBytes(sigAlgs.length), sigAlgs);
        byte[] extensions = concat(shortBytes(ext.length), ext);

        byte[] cipherSuites = concat(shortBytes(2), new byte[]{(byte) 0x00, (byte) 0xA8});
        byte[] clientHello = concat(
                new byte[]{0x03, 0x03},     // client_version TLS 1.2
                clientRandom,               // 32 字节
                new byte[]{0x00},           // session_id 长度 0
                cipherSuites,               // 4 字节
                new byte[]{0x01, 0x00},     // compression: 1 个方法, null(0x00)
                extensions
        );
        byte[] body = concat(new byte[]{0x01}, int24(clientHello.length), clientHello);
        out.write(buildPlainRecord((byte) 22, body));
        out.flush();
        handshakeBuf.write(body);
    }

    private void readServerHelloAndDone() throws Exception {
        boolean gotServerHello = false;
        boolean gotDone = false;
        while (!gotDone) {
            Record rec = readRecordAny();
            if (rec.type == 21) {
                throw new IOException("TLS 握手收到 alert: "
                        + (rec.payload.length > 1 ? (rec.payload[1] & 0xff) : 0));
            }
            if (rec.type != 22) {
                throw new IOException("期望握手记录，实际类型: " + rec.type);
            }
            int off = 0;
            byte[] p = rec.payload;
            while (off < p.length) {
                int htype = p[off] & 0xff;
                int hlen = ((p[off + 1] & 0xff) << 16) | ((p[off + 2] & 0xff) << 8) | (p[off + 3] & 0xff);
                byte[] hbody = Arrays.copyOfRange(p, off + 4, off + 4 + hlen);
                handshakeBuf.write(p, off, 4 + hlen);
                if (htype == 2) {
                    parseServerHello(hbody);
                    gotServerHello = true;
                } else if (htype == 14) {
                    gotDone = true;
                }
                off += 4 + hlen;
            }
        }
        if (!gotServerHello) {
            throw new IOException("未收到 ServerHello");
        }
    }

    private void parseServerHello(byte[] b) throws IOException {
        System.arraycopy(b, 2, serverRandom, 0, 32); // version(2) + random(32)
        int sidLen = b[34] & 0xff;
        int pos = 35 + sidLen;
        int cs = ((b[pos] & 0xff) << 8) | (b[pos + 1] & 0xff);
        if (cs != CIPHER_PSK_AES128_GCM_SHA256) {
            throw new IOException("服务端未选择 PSK_AES_128_GCM 密码套件，实际: 0x"
                    + Integer.toHexString(cs));
        }
    }

    private void deriveKeys() throws Exception {
        // PSK 预主密钥：other_secret(空) || psk
        byte[] premaster = concat(new byte[]{0x00, 0x00}, shortBytes(psk.length), psk);
        masterSecret = TlsCrypto.prf(premaster, "master secret",
                concat(clientRandom, serverRandom), 48);
        byte[] keyBlock = TlsCrypto.prf(masterSecret, "key expansion",
                concat(serverRandom, clientRandom), 40);
        clientWriteKey = Arrays.copyOfRange(keyBlock, 0, 16);
        serverWriteKey = Arrays.copyOfRange(keyBlock, 16, 32);
        clientWriteIV = Arrays.copyOfRange(keyBlock, 32, 36);
        serverWriteIV = Arrays.copyOfRange(keyBlock, 36, 40);
    }

    private void sendClientKeyExchange() throws Exception {
        byte[] body = concat(new byte[]{0x10}, int24(2 + identity.length),
                shortBytes(identity.length), identity);
        out.write(buildPlainRecord((byte) 22, body));
        out.flush();
        handshakeBuf.write(body);
    }

    private void sendChangeCipherSpec() throws IOException {
        out.write(new byte[]{20, 3, 3, 0, 1, 1});
        out.flush();
    }

    private void sendFinished() throws Exception {
        byte[] hash = TlsCrypto.sha256(handshakeBuf.toByteArray());
        byte[] verify = TlsCrypto.prf(masterSecret, "client finished", hash, 12);
        byte[] body = concat(new byte[]{0x14}, int24(12), verify);
        handshakeBuf.write(body); // 计入，供服务端 Finished 校验
        byte[] record = buildEncryptedRecord((byte) 23, body);
        out.write(record);
        out.flush();
    }

    private void readServerFinished() throws Exception {
        Record r = readRecordAny();
        if (r.type == 20) { // 服务端 ChangeCipherSpec（明文）
            r = readRecordAny();
        }
        if (r.type != 23) {
            throw new IOException("期望加密 Finished，实际类型: " + r.type);
        }
        byte[] plaintext = decryptRecord(r);
        byte[] expected = TlsCrypto.prf(masterSecret, "server finished",
                TlsCrypto.sha256(handshakeBuf.toByteArray()), 12);
        byte[] actual = Arrays.copyOfRange(plaintext, 4, 16);
        if (!Arrays.equals(expected, actual)) {
            throw new IOException("TLS Finished 校验失败");
        }
    }

    // ---------- 记录层 ----------

    private static class Record {
        int type;
        byte[] payload;
        Record(int type, byte[] payload) {
            this.type = type;
            this.payload = payload;
        }
    }

    private Record readRecordAny() throws IOException {
        byte[] header = readExactly(5);
        int type = header[0] & 0xff;
        int len = ((header[3] & 0xff) << 8) | (header[4] & 0xff);
        byte[] payload = readExactly(len);
        return new Record(type, payload);
    }

    private byte[] buildPlainRecord(byte contentType, byte[] body) {
        byte[] header = {contentType, 3, 3, (byte) (body.length >> 8), (byte) body.length};
        return concat(header, body);
    }

    private byte[] buildEncryptedRecord(byte contentType, byte[] plaintext) throws Exception {
        long seq = clientSeq++;
        byte[] explicit = longTo8(seq);
        byte[] nonce = concat(clientWriteIV, explicit);
        byte[] aad = {contentType, 3, 3, (byte) (plaintext.length >> 8), (byte) plaintext.length};
        byte[] ct = TlsCrypto.aesGcmEncrypt(clientWriteKey, nonce, plaintext, aad);
        byte[] payload = concat(explicit, ct);
        byte[] header = {contentType, 3, 3, (byte) (payload.length >> 8), (byte) payload.length};
        return concat(header, payload);
    }

    private byte[] decryptRecord(Record r) throws Exception {
        byte[] payload = r.payload;
        byte[] explicit = Arrays.copyOfRange(payload, 0, 8);
        byte[] ct = Arrays.copyOfRange(payload, 8, payload.length);
        byte[] nonce = concat(serverWriteIV, explicit);
        int plainLen = payload.length - 8 - 16;
        byte[] aad = {23, 3, 3, (byte) (plainLen >> 8), (byte) plainLen};
        return TlsCrypto.aesGcmDecrypt(serverWriteKey, nonce, ct, aad);
    }

    private byte[] readExactly(int n) throws IOException {
        byte[] b = new byte[n];
        int off = 0;
        while (off < n) {
            int r = in.read(b, off, n - off);
            if (r < 0) throw new IOException("TLS 连接断开 (EOF)");
            off += r;
        }
        return b;
    }

    // ---------- 透明加密流 ----------

    private class TlsInputStream extends InputStream {
        private byte[] buf = new byte[0];
        private int bufPos = 0;

        @Override
        public int read() throws IOException {
            byte[] b = new byte[1];
            int n = read(b, 0, 1);
            return n <= 0 ? -1 : (b[0] & 0xff);
        }

        @Override
        public int read(byte[] b, int off, int len) throws IOException {
            if (bufPos >= buf.length) {
                Record r = readRecordAny();
                if (r.type != 23) {
                    throw new IOException("TLS 流期望应用数据记录，实际: " + r.type);
                }
                try {
                    buf = decryptRecord(r);
                } catch (Exception e) {
                    throw new IOException("TLS 解密失败", e);
                }
                bufPos = 0;
            }
            int avail = buf.length - bufPos;
            int n = Math.min(len, avail);
            System.arraycopy(buf, bufPos, b, off, n);
            bufPos += n;
            return n;
        }
    }

    private class TlsOutputStream extends OutputStream {
        @Override
        public void write(int b) throws IOException {
            write(new byte[]{(byte) b}, 0, 1);
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            byte[] slice = Arrays.copyOfRange(b, off, off + len);
            try {
                byte[] record = buildEncryptedRecord((byte) 23, slice);
                out.write(record);
                out.flush();
            } catch (Exception e) {
                throw new IOException("TLS 加密失败", e);
            }
        }
    }

    // ---------- 字节工具 ----------

    private static byte[] shortBytes(int v) {
        return new byte[]{(byte) (v >> 8), (byte) v};
    }

    private static byte[] int24(int v) {
        return new byte[]{(byte) (v >> 16), (byte) (v >> 8), (byte) v};
    }

    private static byte[] longTo8(long v) {
        byte[] b = new byte[8];
        for (int i = 7; i >= 0; i--) {
            b[i] = (byte) (v & 0xff);
            v >>= 8;
        }
        return b;
    }

    private static byte[] concat(byte[]... arrays) {
        return TlsCrypto.concat(arrays);
    }
}
