package thorium;

import java.io.IOException;
import java.net.ConnectException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;

/**
 * The main connection class. This is instanced once thorium is started up.
 * @param <T> the td of this class.
 * @date 080213
 * @author Stacx
 */
public class Connection<T> implements EventSource {

  public interface Callback<T> {
    void messages(Connection<T> connection, Iterator<T> messages);

    void idle(Connection<T> connection);

    void closed(Connection<T> connection);

    void garbledMessage(String message, byte[] data);
  }

  private List<ByteBuffer> txBuffers = Collections.synchronizedList(new LinkedList<ByteBuffer>());
  private ByteBuffer rxBuffer = ByteBuffer.allocate(4096);
  private SelectionKey selectionKey;
  private SocketChannel channel;
  private MessageParser<T> parser;
  private Callback<T> callback;

  public static <T> Connection<T> attemptToConnect(InetSocketAddress address, MessageParser<T> parser, Callback<T> callback)
      throws IOException {
    return attemptToConnect(address, parser, callback, 3, 100);
  }

  public static <T> Connection<T> attemptToConnect(InetSocketAddress address, MessageParser<T> parser, Callback<T> callback,
      int maxNumAttempts, int retryDelay) throws IOException {
    for (int numAttempts = 0; numAttempts < maxNumAttempts; numAttempts++) {
      try {
        return Connection.connect(address, parser, callback);
      } catch (ConnectException e1) {
        try { Thread.sleep(retryDelay); } catch (InterruptedException e2) { }
      }
    }
    throw new ConnectException("Could not be connected");
  }

  public static <T> Connection<T> connect(InetSocketAddress address, MessageParser<T> parser, Callback<T> callback)
      throws IOException {
    SocketChannel channel = SocketChannel.open();
    channel.connect(address);
    channel.configureBlocking(false);
    return new Connection<T>(channel, parser, callback);
  }

  public static <T> Connection<T> accept(InetSocketAddress address, MessageParser<T> parser, Callback<T> callback) throws IOException {
    ServerSocketChannel serverChannel = ServerSocketChannel.open();
    ServerSocket socket = serverChannel.socket();
    socket.bind(address);
    SocketChannel channel = serverChannel.accept();
    serverChannel.close();
    channel.configureBlocking(false);
    return new Connection<T>(channel, parser, callback);
  }

  public Connection(SocketChannel channel, MessageParser<T> parser, Callback<T> callback) {
    this.channel = channel;
    this.callback = callback;
    this.parser = parser;
  }

  @Override 
  public SelectionKey register(Selector selector, int ops) throws IOException {
    return selectionKey = channel.register(selector, ops);
  }

  @Override public void unregister() {
    callback.closed(this);
  }

  @Override 
  public void read(SelectionKey key) throws IOException {
    SocketChannel sc = (SocketChannel) key.channel();
    if (sc.isOpen()) {
      int len;
      try {
        len = sc.read(rxBuffer);
      } catch (IOException e) {
        len = -1;
      }
      if (len > 0) {
        Iterator<T> messages = parse();
        if (messages.hasNext()) {
          callback.messages(this, messages);
        }
      } else if (len < 0) {
        close();
      }
    }
  }

  private Iterator<T> parse() throws IOException {
    rxBuffer.flip();
    List<T> result = new ArrayList<T>();
    while (rxBuffer.hasRemaining()) {
      rxBuffer.mark();
      try {
        result.add(parser.parse(rxBuffer));
      } catch (PartialMessageException e) {
        rxBuffer.reset();
        break;
      } catch (GarbledMessageException e) {
        callback.garbledMessage(e.getMessage(), e.getMessageData());
      }
    }
    rxBuffer.compact();
    return result.iterator();
  }

  public void send(byte[] byteArray) {
    send(ByteBuffer.wrap(byteArray));
  }

  public void send(Message message) {
    send(message.toByteBuffer());
  }

  public void send(ByteBuffer buffer) {
    txBuffers.add(buffer);
    if (selectionKey == null)
      throw new IllegalStateException("Connection is not registered");
    selectionKey.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
    selectionKey.selector().wakeup();
  }

  @Override public void write(SelectionKey key) throws IOException {
    try {
      flush();
    } catch (IOException e) {
      close();
    }
    if (txBuffers.isEmpty())
      key.interestOps(SelectionKey.OP_READ);
  }

  public void close() {
    try {
      while (!txBuffers.isEmpty())
        flush();
    } catch (IOException e) {
    }
    SocketChannel sc = (SocketChannel) selectionKey.channel();
    Socket socket = sc.socket();
    if (socket != null) {
      try {
        socket.shutdownInput();
      } catch (IOException e) {
      }
      try {
        socket.shutdownOutput();
      } catch (IOException e) {
      }
      try {
        socket.close();
      } catch (IOException e) {
      }
    }
    try {
      sc.close();
    } catch (IOException e) {
    }
    selectionKey.attach(null);
    selectionKey.cancel();
    selectionKey.selector().wakeup();
  }

  private void flush() throws IOException {
    synchronized(txBuffers) {
      while (!txBuffers.isEmpty()) {
        ByteBuffer txBuffer = txBuffers.get(0);
        if (!write(txBuffer)) {
          break;
        }
        txBuffers.remove(0);
      }
    }
  }

  private boolean write(ByteBuffer txBuffer) throws IOException {
    while (txBuffer.hasRemaining()) {
      if (channel.write(txBuffer) == 0) {
        return false;
      }
    }
    return true;
  }

  @Override public void timeout() {
    callback.idle(this);
  }

  @Override public EventSource accept(SelectionKey key) throws IOException {
    throw new UnsupportedOperationException();
  }

  @Override public boolean isClosed() {
    SocketChannel sc = (SocketChannel) selectionKey.channel();
    return !sc.isOpen();
  }
}