// Copyright (c) 2005 Brian Wellington (bwelling@xbill.org)

package org.xbill.DNS;

import java.io.*;
import java.net.*;
import java.nio.*;
import java.nio.channels.*;

final class TCPClient extends Client {

    public TCPClient(long endTime) throws IOException {
        super(SocketChannel.open(), endTime);
    }

    void
    bind(SocketAddress addr) throws IOException {
        SocketChannel channel = (SocketChannel) key.channel();
        channel.socket().bind(addr);
    }

    void
    connect(SocketAddress addr) throws IOException {
        SocketChannel channel = (SocketChannel) key.channel();
        if (channel.connect(addr))
            return;
        key.interestOps(SelectionKey.OP_CONNECT);
        try {
            while (!channel.finishConnect()) {
                if (!key.isConnectable())
                    blockUntil(key, endTime);
            }
        } finally {
            if (key.isValid())
                key.interestOps(0);
        }
    }

    void
    send(byte[] data) throws IOException {
        SocketChannel channel = (SocketChannel) key.channel();
        verboseLog("TCP write", data);
        byte[] lengthArray = new byte[2];
        lengthArray[0] = (byte) (data.length >>> 8);
        lengthArray[1] = (byte) (data.length & 0xFF);
        ByteBuffer[] buffers = new ByteBuffer[2];
        buffers[0] = ByteBuffer.wrap(lengthArray);
        buffers[1] = ByteBuffer.wrap(data);
        int nsent = 0;
        key.interestOps(SelectionKey.OP_WRITE);
        try {
            while (nsent < data.length + 2) {
                if (key.isWritable()) {
                    long n = channel.write(buffers);
                    if (n < 0)
                        throw new EOFException();
                    nsent += (int) n;
                    if (nsent < data.length + 2 &&
                            System.currentTimeMillis() > endTime)
                        throw new SocketTimeoutException();
                } else
                    blockUntil(key, endTime);
            }
        } finally {
            if (key.isValid())
                key.interestOps(0);
        }
    }

    private byte[]
    _recv(int length) throws IOException {
        SocketChannel channel = (SocketChannel) key.channel();
        int nrecvd = 0;
        byte[] data = new byte[length];
        ByteBuffer buffer = ByteBuffer.wrap(data);
        key.interestOps(SelectionKey.OP_READ);
        try {
            while (nrecvd < length) {
                if (key.isReadable()) {
                    long n = channel.read(buffer);
                    if (n < 0)
                        throw new EOFException();
                    nrecvd += (int) n;
                    if (nrecvd < length &&
                            System.currentTimeMillis() > endTime)
                        throw new SocketTimeoutException();
                } else
                    blockUntil(key, endTime);
            }
        } finally {
            if (key.isValid())
                key.interestOps(0);
        }
        return data;
    }

    byte[]
    recv() throws IOException {
        byte[] buf = _recv(2);
        int length = ((buf[0] & 0xFF) << 8) + (buf[1] & 0xFF);
        byte[] data = _recv(length);
        verboseLog("TCP read", data);
        return data;
    }

    static byte[]
    sendrecv(SocketAddress local, SocketAddress remote, byte[] data, long endTime)
            throws IOException {
        TCPClient client = new TCPClient(endTime);
        try {
            if (local != null)
                client.bind(local);
            client.connect(remote);
            client.send(data);
            return client.recv();
        } finally {
            client.cleanup();
        }
    }

    static byte[]
    sendrecv(SocketAddress addr, byte[] data, long endTime) throws IOException {
        return sendrecv(null, addr, data, endTime);
    }

}
