package com.ttstd.adbloopback.adb;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.util.HashMap;
import java.util.Map;

/**
 * 自建 ADB 客户端连接：与系统 adbd（无线调试守护进程）建立明文 TCP 连接，
 * 完成 CNXN / AUTH 握手，然后可开启 shell 流执行指令。
 */
public class AdbConnection {
    private final Socket socket;
    private final AdbKeyPair keyPair;
    private InputStream in;
    private OutputStream out;
    private boolean connected = false;
    private int lastLocalId = 1;

    public AdbConnection(Socket socket, AdbKeyPair keyPair) {
        this.socket = socket;
        this.keyPair = keyPair;
    }

    public static AdbConnection connect(String host, int port, AdbKeyPair keyPair) throws Exception {
        Socket s = new Socket();
        s.connect(new InetSocketAddress(host, port), 5000);
        s.setSoTimeout(2000); // 读超时，供 shell 读取在无输出时及时返回
        s.setTcpNoDelay(true);
        AdbConnection c = new AdbConnection(s, keyPair);
        c.connect();
        return c;
    }

    private void connect() throws Exception {
        in = socket.getInputStream();
        out = socket.getOutputStream();

        // 1) 发送 CNXN
        AdbMessage.write(out, new AdbMessage(AdbConstants.CNXN, AdbConstants.VERSION,
                AdbConstants.MAX_PAYLOAD, "device::adb-loopback-demo".getBytes("UTF-8")));

        // 2) 处理 AUTH 挑战，直到收到 CNXN 确认
        while (true) {
            AdbMessage resp = readMessage();
            if (resp.command == AdbConstants.AUTH) {
                int sub = resp.arg0;
                if (sub == AdbConstants.AUTH_TOKEN) {
                    byte[] signature = keyPair.sign(resp.data);
                    AdbMessage.write(out, new AdbMessage(AdbConstants.AUTH,
                            AdbConstants.AUTH_SIGNATURE, 0, signature));
                } else if (sub == AdbConstants.AUTH_RSAPUBLICKEY) {
                    // 某些场景下服务端主动索要公钥，一并提交
                    AdbMessage.write(out, new AdbMessage(AdbConstants.AUTH,
                            AdbConstants.AUTH_RSAPUBLICKEY, 0,
                            keyPair.getAdbPublicKey().getBytes("UTF-8")));
                } else {
                    throw new IOException("ADB AUTH 失败，sub=" + sub);
                }
            } else if (resp.command == AdbConstants.CNXN) {
                connected = true;
                break;
            } else {
                throw new IOException("ADB 握手收到未知消息: 0x"
                        + Integer.toHexString(resp.command));
            }
        }
        if (!connected) {
            throw new IOException("ADB 连接建立失败");
        }
    }

    public AdbStream open(String destination) throws Exception {
        int localId = lastLocalId++;
        AdbMessage.write(out, new AdbMessage(AdbConstants.OPEN, localId, 0,
                (destination + "\0").getBytes("UTF-8")));

        AdbMessage resp = readMessage();
        if (resp.command == AdbConstants.OKAY) {
            AdbStream stream = new AdbStream(this, localId, destination);
            stream.setRemoteId(resp.arg0);
            return stream;
        } else if (resp.command == AdbConstants.CLSE) {
            throw new IOException("ADB 流打开被拒绝");
        } else {
            throw new IOException("ADB 打开流收到未知响应: 0x"
                    + Integer.toHexString(resp.command));
        }
    }

    synchronized void sendInternal(int command, int arg0, int arg1, byte[] data) throws IOException {
        AdbMessage.write(out, new AdbMessage(command, arg0, arg1, data));
    }

    synchronized AdbMessage readMessage() throws IOException {
        return AdbMessage.read(in);
    }

    public void close() throws IOException {
        socket.close();
        connected = false;
    }

    public boolean isConnected() {
        return connected;
    }
}
