package net.tootallnate.websocket;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.URI;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Iterator;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;

import net.tootallnate.websocket.WebSocketListener.Draft;

/**
 * The <tt>WebSocketClient</tt> is an abstract class that expects a valid
 * "ws://" URI to connect to. When connected, an instance recieves important
 * events related to the life of the connection. A subclass must implement
 * <var>onOpen</var>, <var>onClose</var>, and <var>onMessage</var> to be
 * useful. An instance can send messages to it's connected server via the
 * <var>send</var> method.
 * @author Nathan Rajlich
 */
public abstract class WebSocketClient implements Runnable, WebSocketListener {


  // INSTANCE PROPERTIES /////////////////////////////////////////////////////
  /**
   * The URI this client is supposed to connect to.
   */
  private URI uri;
  /**
   * The WebSocket instance this client object wraps.
   */
  private WebSocket conn;
  /**
   * @author The Websocket mode this client is in.
   */
  private Draft draft;
  /**
   * Number 1 used in handshake 
   */
  private int number1=0;
  /** Number 2 used in handshake
   */
  private int number2=0;
  /** Key3 used in handshake
   */
  private byte[] key3=null;
  public static enum Draft{
    DRAFT75,
    DRAFT76
  }

  // CONSTRUCTOR /////////////////////////////////////////////////////////////
  /**
   * Constructs a WebSocketClient instance and sets it to the connect to the
   * specified URI. The client does not attampt to connect automatically. You
   * must call <var>connect</var> first to initiate the socket connection.
   * @param serverUri The <tt>URI</tt> of the WebSocket server to connect to.
   */
  public WebSocketClient(URI serverUri,Draft draft) {
    this.uri = serverUri;
    this.draft = draft;
  }

  // PUBLIC INSTANCE METHODS /////////////////////////////////////////////////
  /**
   * Gets the URI that this WebSocketClient is connected to.
   * @return The <tt>URI</tt> for this WebSocketClient.
   */
  public URI getURI() {
    return uri;
  }

  /**
   * Starts a background thread that attempts and maintains a WebSocket
   * connection to the URI specified in the constructor or via <var>setURI</var>.
   * <var>setURI</var>.
   */
  public void connect() {
    (new Thread(this)).start();
  }

  /**
   * Calls <var>close</var> on the underlying SocketChannel, which in turn
   * closes the socket connection, and ends the client socket thread.
   * @throws IOException When socket related I/O errors occur.
   */
  public void close() throws IOException {
    conn.close();
  }

  /**
   * Sends <var>text</var> to the connected WebSocket server.
   * @param text The String to send to the WebSocket server.
   * @throws IOException When socket related I/O errors occur.
   */
  public void send(String text) throws IOException {
    conn.send(text);
  }

  // Runnable IMPLEMENTATION /////////////////////////////////////////////////
  public void run() {
    try {
      int port = uri.getPort();
      if (port == -1) {
        port = WebSocket.DEFAULT_PORT;
      }

      // The WebSocket constructor expects a SocketChannel that is
      // non-blocking, and has a Selector attached to it.
      SocketChannel client = SocketChannel.open();
      client.configureBlocking(false);
      client.connect(new InetSocketAddress(uri.getHost(), port));

      Selector selector = Selector.open();

      this.conn = new WebSocket(client, new LinkedBlockingQueue<ByteBuffer>(), this);
      client.register(selector, client.validOps());

      // Continuous loop that is only supposed to end when close is called
      while (selector.select(500) > 0) {

        Set<SelectionKey> keys = selector.selectedKeys();
        Iterator<SelectionKey> i = keys.iterator();

        while (i.hasNext()) {
          SelectionKey key = i.next();
          i.remove();

          // When 'conn' has connected to the host
          if (key.isConnectable()) {
            
            // Ensure connection is finished
            if (client.isConnectionPending()) {
              client.finishConnect();
            }

            // Now send WebSocket client-side handshake
            Random r=new Random();
            String path = "/" + uri.getPath();
            String host = uri.getHost() + (port != WebSocket.DEFAULT_PORT ? ":" + port : "");
            String origin = null; // TODO: Make 'origin' configurable
            String request = "GET " + path + " HTTP/1.1\r\n" +
                              "Upgrade: WebSocket\r\n" +
                              "Connection: Upgrade\r\n" +
                              "Host: " + host + "\r\n" +
                              "Origin: " + origin + "\r\n";
                      if(this.draft==Draft.DRAFT76)
                      {
                        request+="Sec-WebSocket-Key1: " + this.generateKey() + "\r\n";
                        request+="Sec-WebSocket-Key2: " + this.generateKey() + "\r\n";
                        this.key3=new byte[8];
                        r.nextBytes(this.key3);
                      }
                              //extraHeaders.toString() +
                   request+="\r\n";
            conn.socketChannel().write(ByteBuffer.wrap(request.getBytes(WebSocket.UTF8_CHARSET)));
            if(this.key3 !=null){
              conn.socketChannel().write(ByteBuffer.wrap(this.key3));
            }
          }

          // When 'conn' has recieved some data
          if (key.isReadable()) {
            conn.handleRead();
          }
        }
      }
    } catch (IOException ex) {
      ex.printStackTrace();
    } catch (NoSuchAlgorithmException e) {
      // TODO Auto-generated catch block
      e.printStackTrace();
    }
  }
  private String generateKey(){
    Random r=new Random();
    long maxNumber=4294967295L;
    long spaces=r.nextInt(12)+1;
    int max=new Long(maxNumber/spaces).intValue();
    max=Math.abs(max);
    int number=r.nextInt(max)+1;
    if(this.number1==0){
      this.number1=number;
    }
    else{
      this.number2=number;
    }   
    long product=number*spaces;
    String key=Long.toString(product);
    int numChars=r.nextInt(12);
    for (int i=0;i<numChars;i++){
      int position=r.nextInt(key.length());
      position=Math.abs(position);
      char randChar=(char)(r.nextInt(95)+33);
      //exclude numbers here
      if(randChar >= 48 && randChar <=57){
        randChar-=15;
      }
      key=new StringBuilder(key).insert(position, randChar).toString();
    }
    for (int i=0;i<spaces;i++){
      int position=r.nextInt(key.length()-1)+1;
      position=Math.abs(position);
      key=new StringBuilder(key).insert(position,"\u0020").toString();
    }
    return key;
  }

  // WebSocketListener IMPLEMENTATION ////////////////////////////////////////
  /**
   * Parses the server's handshake to verify that it's a valid WebSocket
   * handshake.
   * @param conn The {@link WebSocket} instance who's handshake has been recieved.
   *             In the case of <tt>WebSocketClient</tt>, this.conn == conn.
   * @param handshake The entire UTF-8 decoded handshake from the connection.
   * @return <var>true</var> if <var>handshake</var> is a valid WebSocket server
   *         handshake, <var>false</var> otherwise.
   * @throws IOException When socket related I/O errors occur.
 * @throws NoSuchAlgorithmException 
   */
  public boolean onHandshakeRecieved(WebSocket conn, String handshake,byte[] reply) throws IOException, NoSuchAlgorithmException {
    // TODO: Do some parsing of the returned handshake, and close connection
    // (return false) if we recieved anything unexpected.
  if(this.draft==Draft.DRAFT76){
    if(reply==null){
      return false;
    }
    byte[] challenge=new byte[]{
        (byte)(  this.number1 >> 24 ),
              (byte)( (this.number1 << 8) >> 24 ),
              (byte)( (this.number1 << 16) >> 24 ),
              (byte)( (this.number1 << 24) >> 24 ),
              (byte)(  this.number2 >> 24 ),
              (byte)( (this.number2 << 8) >> 24 ),
              (byte)( (this.number2 << 16) >> 24 ),
              (byte)( (this.number2 << 24) >> 24 ),
              this.key3[0],
              this.key3[1],
              this.key3[2],
              this.key3[3],
              this.key3[4],
              this.key3[5],
              this.key3[6],
              this.key3[7]
              };
    MessageDigest md5=MessageDigest.getInstance("MD5");
    byte[] expected=md5.digest(challenge);
    for(int i=0;i<reply.length;i++){
        if(expected[i]!=reply[i]){
          return false;
        }
      } 
  }
    return true;
  }

  /**
   * Calls subclass' implementation of <var>onMessage</var>.
   * @param conn
   * @param message
   */
  public void onMessage(WebSocket conn, String message) {
    onMessage(message);
  }

  /**
   * Calls subclass' implementation of <var>onOpen</var>.
   * @param conn
   */
  public void onOpen(WebSocket conn) {
    onOpen();
  }

  /**
   * Calls subclass' implementation of <var>onClose</var>.
   * @param conn
   */
  public void onClose(WebSocket conn) {
    onClose();
  }
  
  @Override
  public net.tootallnate.websocket.WebSocketListener.Draft getDraft() {
    return (net.tootallnate.websocket.WebSocketListener.Draft)net.tootallnate.websocket.WebSocketListener.Draft.valueOf(this.draft.name());
  }

  // ABTRACT METHODS /////////////////////////////////////////////////////////
  public abstract void onMessage(String message);
  public abstract void onOpen();
  public abstract void onClose();
}
