/*******************************************************************************
 * Copyright 2018 572682
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License.  You may obtain a copy
 * of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
 * License for the specific language governing permissions and limitations under
 * the License.
 ******************************************************************************/
package us.dot.its.jpo.ode.wrapper;

import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import javax.net.ssl.SSLContext;
import javax.websocket.ClientEndpointConfig;
import javax.websocket.ClientEndpointConfig.Builder;
import javax.websocket.ClientEndpointConfig.Configurator;
import javax.websocket.CloseReason;
import javax.websocket.ContainerProvider;
import javax.websocket.Decoder;
import javax.websocket.Endpoint;
import javax.websocket.EndpointConfig;
import javax.websocket.Session;
import javax.websocket.WebSocketContainer;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import us.dot.its.jpo.ode.util.WebSocketUtils;

public class WebSocketEndpoint<T> extends Endpoint {
   private static final Logger logger = LoggerFactory
         .getLogger(WebSocketEndpoint.class);

   private URI uri;
   private SSLContext sslContext;
   private Map<String, Object> userProperties;
   private Map<String, Map<String, String>> wsHeaders;
   private Session wsSession = null;
   private WebSocketMessageHandler<T> handler;
   private List<Class<? extends Decoder>> decoders = new ArrayList<Class<? extends Decoder>>();;
   private Configurator configurator;
   private ClientEndpointConfig wsConfig;

   /**
    * General Constructor
    * 
    * @param uri
    *           - the URI of the WebSocket server
    * @param sslContext
    *           - SSL context or null if not running on SSL
    * @param userProperties
    *           - map of user properties to be included in the session object
    * @param wsHeaders
    *           - map of headers sent to in the request
    * @param handler
    *           - the incoming message handler
    * @param decoders
    *           - the incoming message decoders (or null if no decoder needed)
    * @throws URISyntaxException 
    *    - If uri string is null 
    *    - If the uri string violates RFC 2396, as augmented by the above deviations 
    */
   public WebSocketEndpoint(String uri, SSLContext sslContext,
         Map<String, Object> userProperties,
         Map<String, Map<String, String>> wsHeaders,
         WebSocketMessageHandler<T> handler,
         List<Class<? extends WebSocketMessageDecoder<?>>> decoders) 
               throws URISyntaxException {
      
      this(new URI(uri), sslContext, userProperties, wsHeaders, handler, decoders);
   }

   /**
    * General Constructor
    * 
    * @param uri
    *           - the URI of the WebSocket server
    * @param sslContext
    *           - SSL context or null if not running on SSL
    * @param userProperties
    *           - map of user properties to be included in the session object
    * @param wsHeaders
    *           - map of headers sent to in the request
    * @param handler
    *           - the incoming message handler
    * @param decoders
    *           - the incoming message decoders (or null if no decoder needed)
    */
   public WebSocketEndpoint(URI uri, SSLContext sslContext,
         Map<String, Object> userProperties,
         Map<String, Map<String, String>> wsHeaders,
         WebSocketMessageHandler<T> handler,
         List<Class<? extends WebSocketMessageDecoder<?>>> decoders) {
      super();
      this.uri = uri;
      this.sslContext = sslContext;
      this.userProperties = userProperties;
      this.wsHeaders = wsHeaders;
      this.handler = handler;
      
      this.configurator = createConfigurator(wsHeaders);

      Builder builder = ClientEndpointConfig.Builder.create().configurator(
            this.configurator);

      if (decoders != null) {
         this.decoders.addAll(decoders);
         builder.decoders(this.decoders);
      }

      this.wsConfig = builder.build();

      if (sslContext != null)
         wsConfig.getUserProperties().put(
               "org.apache.tomcat.websocket.SSL_CONTEXT", sslContext);

      if (this.userProperties != null) {
         for (String key : userProperties.keySet()) {
            Object prop = userProperties.get(key);
            if (prop != null)
               wsConfig.getUserProperties().put(key, prop);
         }
      }
   }

   public URI getUri() {
      return uri;
   }

   /**
    * @param uri - URI of the WebSocket server
    * @return - this client object
    */
   public WebSocketEndpoint<T> setUri(URI uri) {
      this.uri = uri;
      return this;
   }

   /**
    * @param uri - URI of the WebSocket server
    * @return - this client object
    * @throws WebSocketException
    *    - If uri string is null 
    *    - If the uri string violates RFC 2396, as augmented by the above deviations 
    */
   public WebSocketEndpoint<T> setUri(String uri) throws WebSocketException {
      try {
         this.uri = new URI(uri);
      } catch (Exception e) {
         throw new WebSocketException("Error setting URI.", e);
      }
      return this;
   }

   public SSLContext getSslContext() {
      return sslContext;
   }

   public WebSocketEndpoint<T> setSslContext(SSLContext sslContext) {
      this.sslContext = sslContext;
      return this;
   }

   public Map<String, Object> getUserProperties() {
      return userProperties;
   }

   public WebSocketEndpoint<T> setUserProperties(
         Map<String, Object> userProperties) {
      this.userProperties = userProperties;
      return this;
   }

   public Map<String, Map<String, String>> getWsHeaders() {
      return wsHeaders;
   }

   public WebSocketEndpoint<T> setWsHeaders(
         Map<String, Map<String, String>> wsHeaders) {
      this.wsHeaders = wsHeaders;
      return this;
   }

   public WebSocketMessageHandler<T> getHandler() {
      return handler;
   }

   public WebSocketEndpoint<T> setHandler(WebSocketMessageHandler<T> handler) {
      this.handler = handler;
      return this;
   }

   public List<Class<? extends Decoder>> getDecoders() {
      return decoders;
   }

   public WebSocketEndpoint<T> setDecoders(List<Class<? extends Decoder>> decoders) {
      this.decoders = decoders;
      return this;
   }

   /**
    * Connects to the WebSocketServer creating a WebSocket session
    * 
    * @return - the WebSocket session created
    * @throws WebSocketException
    *            - If a message handler has already been registered for the
    *            associated message type
    */
   public Session connect() throws WebSocketException {
      try {
         WebSocketContainer container = ContainerProvider
               .getWebSocketContainer();

         logger.info("Opening connection to {}", uri.toString());
         wsSession = container.connectToServer(this, wsConfig, uri);
         wsSession.addMessageHandler(this.handler);
      } catch (Exception e) {
         throw new WebSocketException(e);
      }
      return wsSession;
   }

   /**
    * Creates a javax.websocket.ClientEndpointConfig.Configurator configured
    * with the given headers.
    * 
    * @param wsHeaders
    *           - map of key value pairs. The key is the header name, the value
    *           is a list of name-value pairs.
    * 
    * @return - a javax.websocket.ClientEndpointConfig.Configurator object
    */
   private Configurator createConfigurator(
         final Map<String, Map<String, String>> wsHeaders) {
      return new Configurator() {
         @Override
         public void beforeRequest(Map<String, List<String>> headers) {
            super.beforeRequest(headers);
            for (String header : wsHeaders.keySet()) {
               Map<String, String> hmap = wsHeaders.get(header);
               ArrayList<String> headerValue = new ArrayList<String>();
               for (Entry<String, String> entry : hmap.entrySet()) {
                  StringBuilder nvp = new StringBuilder(entry.getKey());
                  nvp.append('=').append(entry.getValue());
                  headerValue.add(nvp.toString());
               }
               headers.put(header, headerValue);
            }
         }
      };
   }

   /*
    * (non-Javadoc)
    * 
    * @see javax.websocket.Endpoint#onOpen(javax.websocket.Session,
    * javax.websocket.EndpointConfig)
    */
   @Override
   public void onOpen(Session session, EndpointConfig config) {
      handler.onOpen(session, config);
      logger.info("Connection opened. Session ID {}", session.getId());
   }

   /*
    * (non-Javadoc)
    * 
    * @see javax.websocket.Endpoint#onClose(javax.websocket.Session,
    * javax.websocket.CloseReason)
    */
   @Override
   public void onClose(Session session, CloseReason reason) {
      handler.onClose(session, reason);
      logger.info("Connection closed. Session ID {}. Reason: {}", session.getId(), reason);
      this.wsSession = null;
   }

   /*
    * (non-Javadoc)
    * 
    * @see javax.websocket.Endpoint#onError(javax.websocket.Session,
    * java.lang.Throwable)
    */
   @Override
   public void onError(Session session, Throwable t) {
      handler.onError(session, t);
      logger.error("WebSocket Error. Session ID " + session.getId(), t);
   }

   /**
    * Sends a message to the server end-point.
    *
    * @param message
    *           - the message to be sent
    * @throws WebSocketException - wraps IOException 
    */
   public void send(String message) throws WebSocketException {
      try {
         WebSocketUtils.sendSync(this.wsSession, message);
      } catch (IOException e) {
         throw new WebSocketException("Error sending message", e);
      }
   }

   /**
    * Closes the WebSocket connection. This method should be used to gracefully
    * close a WebSocket connection.
    * 
    * Note that the method closes the session but does not set the wsSession to
    * null. wsSession is set to null in the onClose method.
    * 
    * @throws WebSocketException - if an I/O error occurs while the WebSocket 
    * session is being closed.
    */
   public void close() throws WebSocketException {
      try {
         if (null != wsSession) {
            logger.info("Closing WebSocket session {}", wsSession.getId());
            wsSession.close();
         }
      } catch (IOException e) {
         throw new WebSocketException(e);
      }
      
      
   }

   public Session getWsSession() {
      return wsSession;
   }

   public ClientEndpointConfig getWsConfig() {
      return wsConfig;
   }

   public static class WebSocketException extends Exception {

      private static final long serialVersionUID = 1L;

      public WebSocketException(Exception e) {
         super(e);
      }

      public WebSocketException(String msg) {
         super(msg);
      }

      public WebSocketException(String msg, Exception e) {
         super(msg, e);
      }

   }
}
