/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License, Version 1.0 only
 * (the "License").  You may not use this file except in compliance
 * with the License.
 *
 * You can obtain a copy of the license at
 * trunk/opends/resource/legal-notices/OpenDS.LICENSE
 * or https://OpenDS.dev.java.net/OpenDS.LICENSE.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at
 * trunk/opends/resource/legal-notices/OpenDS.LICENSE.  If applicable,
 * add the following below this CDDL HEADER, with the fields enclosed
 * by brackets "[]" replaced with your own identifying information:
 *      Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 *
 *
 *      Copyright 2010 Sun Microsystems, Inc.
 */

package org.opends.sdk;



import static com.sun.opends.sdk.ldap.LDAPConstants.TYPE_AUTHENTICATION_SASL;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.net.ssl.SSLContext;
import javax.security.auth.callback.*;
import javax.security.sasl.*;

import org.opends.sdk.asn1.ASN1;
import org.opends.sdk.asn1.ASN1Reader;
import org.opends.sdk.controls.Control;
import org.opends.sdk.controls.ControlDecoder;
import org.opends.sdk.requests.*;
import org.opends.sdk.responses.*;

import org.glassfish.grizzly.TransportFactory;
import org.glassfish.grizzly.nio.transport.TCPNIOTransport;
import com.sun.opends.sdk.controls.AccountUsabilityRequestControl;
import com.sun.opends.sdk.controls.AccountUsabilityResponseControl;
import com.sun.opends.sdk.ldap.GrizzlyLDAPListenerOptions;



/**
 * A simple ldap server that manages 1000 entries and used for running
 * testcases. //FIXME: make it MT-safe.
 */
public class LDAPServer implements
    ServerConnectionFactory<LDAPClientContext, Integer>
{
  // Creates an abandonable request from the ordinary requests.
  private static class AbandonableRequest implements Request
  {
    // the request.
    private final Request request;

    // whether is has been cancelled.
    private final AtomicBoolean isCanceled;



    // Ctor.
    AbandonableRequest(final Request request)
    {
      this.request = request;
      this.isCanceled = new AtomicBoolean(false);
    }



    public Request addControl(final Control cntrl)
        throws UnsupportedOperationException, NullPointerException
    {
      return request.addControl(cntrl);
    }



    public <C extends Control> C getControl(final ControlDecoder<C> decoder,
        final DecodeOptions options) throws DecodeException,
        NullPointerException
    {
      return request.getControl(decoder, options);
    }



    public List<Control> getControls()
    {
      return request.getControls();
    }



    void cancel()
    {
      isCanceled.set(true);
    }



    boolean isCanceled()
    {
      return isCanceled.get();
    }
  }



  // The singleton instance.
  private static final LDAPServer instance = new LDAPServer();



  /**
   * Returns the singleton instance.
   *
   * @return Singleton instance.
   */
  public static LDAPServer getInstance()
  {
    return instance;
  }



  private class LDAPServerConnection implements
      ServerConnection<Integer>
  {

    private final LDAPClientContext clientContext;
    private SaslServer saslServer;



    private LDAPServerConnection(LDAPClientContext clientContext)
    {
      this.clientContext = clientContext;
    }



    /**
     * Abandons the request sent by the client.
     *
     * @param context
     * @param request
     * @throws UnsupportedOperationException
     */
    public void handleAbandon(final Integer context,
        final AbandonRequest request) throws UnsupportedOperationException
    {
      // Check if we have any concurrent operation with this message id.
      final AbandonableRequest req = requestsInProgress.get(request
          .getRequestID());
      if (req == null)
      {
        // Nothing to do here.
        return;
      }
      // Cancel the request
      req.cancel();
      // No response is needed.
    }



    /**
     * Adds the request sent by the client.
     *
     * @param context
     * @param request
     * @param handler
     * @param intermediateResponseHandler
     * @throws UnsupportedOperationException
     */
    public void handleAdd(final Integer context,
        final AddRequest request, final ResultHandler<? super Result> handler,
        final IntermediateResponseHandler intermediateResponseHandler)
        throws UnsupportedOperationException
    {
      Result result = null;
      final AbandonableRequest abReq = new AbandonableRequest(request);
      requestsInProgress.put(context, abReq);
      // Get the DN.
      final DN dn = request.getName();
      if (entryMap.containsKey(dn))
      {
        // duplicate entry.
        result = Responses.newResult(ResultCode.ENTRY_ALREADY_EXISTS);
        final ErrorResultException ere = ErrorResultException.wrap(result);
        handler.handleErrorResult(ere);
        // doesn't matter if it was canceled.
        requestsInProgress.remove(context);
        return;
      }

      // Create an entry out of this request.
      final SearchResultEntry entry = Responses.newSearchResultEntry(dn);
      for (final Control control : request.getControls())
      {
        entry.addControl(control);
      }

      for (final Attribute attr : request.getAllAttributes())
      {
        entry.addAttribute(attr);
      }

      if (abReq.isCanceled())
      {
        result = Responses.newResult(ResultCode.CANCELLED);
        final ErrorResultException ere = ErrorResultException.wrap(result);
        handler.handleErrorResult(ere);
        requestsInProgress.remove(context);
        return;
      }
      // Add this to the map.
      entryMap.put(dn, entry);
      requestsInProgress.remove(context);
      result = Responses.newResult(ResultCode.SUCCESS);
      handler.handleResult(result);
    }



    /**
     * @param context
     * @param version
     * @param request
     * @param resultHandler
     * @param intermediateResponseHandler
     * @throws UnsupportedOperationException
     */
    public void handleBind(final Integer context, final int version,
        final BindRequest request,
        final ResultHandler<? super BindResult> resultHandler,
        final IntermediateResponseHandler intermediateResponseHandler)
        throws UnsupportedOperationException
    {
      // TODO: all bind types.
      final AbandonableRequest abReq = new AbandonableRequest(request);
      requestsInProgress.put(context, abReq);
      if (request.getAuthenticationType() == TYPE_AUTHENTICATION_SASL
          && request instanceof GenericBindRequest)
      {
        ASN1Reader reader = ASN1.getReader(((GenericBindRequest) request)
            .getAuthenticationValue());
        try
        {
          String saslMech = reader.readOctetStringAsString();
          ByteString saslCred;
          if (reader.hasNextElement())
          {
            saslCred = reader.readOctetString();
          }
          else
          {
            saslCred = ByteString.empty();
          }

          if (saslServer == null
              || !saslServer.getMechanismName().equalsIgnoreCase(saslMech))
          {
            final Map<String, String> props = new HashMap<String, String>();
            props.put(Sasl.QOP, "auth-conf,auth-int,auth");
            saslServer = Sasl.createSaslServer(saslMech, "ldap", clientContext
                .getLocalAddress().getHostName(), props, new CallbackHandler()
            {
              public void handle(Callback[] callbacks) throws IOException,
                  UnsupportedCallbackException
              {
                for (final Callback callback : callbacks)
                {
                  if (callback instanceof NameCallback)
                  {
                    // Do nothing
                  }
                  else if (callback instanceof PasswordCallback)
                  {
                    ((PasswordCallback) callback).setPassword("password"
                        .toCharArray());
                  }
                  else if (callback instanceof AuthorizeCallback)
                  {
                    ((AuthorizeCallback) callback).setAuthorized(true);
                  }
                  else if (callback instanceof RealmCallback)
                  {
                    // Do nothing
                  }
                  else
                  {
                    throw new UnsupportedCallbackException(callback);

                  }
                }
              }
            });
          }

          byte[] challenge = saslServer
              .evaluateResponse(saslCred.toByteArray());
          if (saslServer.isComplete())
          {
            resultHandler.handleResult(Responses.newBindResult(
                ResultCode.SUCCESS).setServerSASLCredentials(
                ByteString.wrap(challenge)));

            String qop = (String) saslServer.getNegotiatedProperty(Sasl.QOP);
            if (qop != null
                && (qop.equalsIgnoreCase("auth-int") || qop
                    .equalsIgnoreCase("auth-conf")))
            {
              ConnectionSecurityLayer csl = new ConnectionSecurityLayer()
              {
                public void dispose()
                {
                  try
                  {
                    saslServer.dispose();
                  }
                  catch (SaslException e)
                  {
                    e.printStackTrace();
                  }
                }



                public byte[] unwrap(byte[] incoming, int offset, int len)
                    throws ErrorResultException
                {
                  try
                  {
                    return saslServer.unwrap(incoming, offset, len);
                  }
                  catch (SaslException e)
                  {
                    throw ErrorResultException.wrap(Responses.newResult(
                        ResultCode.OPERATIONS_ERROR).setCause(e));
                  }
                }



                public byte[] wrap(byte[] outgoing, int offset, int len)
                    throws ErrorResultException
                {
                  try
                  {
                    return saslServer.wrap(outgoing, offset, len);
                  }
                  catch (SaslException e)
                  {
                    throw ErrorResultException.wrap(Responses.newResult(
                        ResultCode.OPERATIONS_ERROR).setCause(e));
                  }
                }
              };

              clientContext.startSASL(csl);
            }

          }
          else
          {
            resultHandler.handleResult(Responses.newBindResult(
                ResultCode.SASL_BIND_IN_PROGRESS).setServerSASLCredentials(
                ByteString.wrap(challenge)));
          }
        }
        catch (Exception e)
        {
          resultHandler.handleErrorResult(ErrorResultException.wrap(Responses
              .newBindResult(ResultCode.OPERATIONS_ERROR).setCause(e)
              .setDiagnosticMessage(e.toString())));
        }
      }
      else
      {
        resultHandler.handleResult(Responses.newBindResult(ResultCode.SUCCESS));
      }
      requestsInProgress.remove(context);
    }



    /**
     * {@inheritDoc}
     */
    public void handleConnectionClosed(final Integer context,
        final UnbindRequest request)
    {
      close();
    }



    private void close()
    {
      if (saslServer != null)
      {
        try
        {
          saslServer.dispose();
        }
        catch (SaslException e)
        {
          e.printStackTrace();
        }
      }
    }



    /**
     * {@inheritDoc}
     */
    public void handleConnectionDisconnected(ResultCode resultCode,
        String message)
    {
      close();
    }



    /**
     * {@inheritDoc}
     */
    public void handleConnectionError(final Throwable error)
    {
      close();
    }



    /**
     * @param context
     * @param request
     * @param resultHandler
     * @param intermediateResponseHandler
     * @throws UnsupportedOperationException
     */
    public void handleCompare(final Integer context,
        final CompareRequest request,
        final ResultHandler<? super CompareResult> resultHandler,
        final IntermediateResponseHandler intermediateResponseHandler)
        throws UnsupportedOperationException
    {
      CompareResult result = null;
      final AbandonableRequest abReq = new AbandonableRequest(request);
      requestsInProgress.put(context, abReq);
      // Get the DN.
      final DN dn = request.getName();
      if (!entryMap.containsKey(dn))
      {
        // entry not found.
        result = Responses.newCompareResult(ResultCode.NO_SUCH_ATTRIBUTE);
        final ErrorResultException ere = ErrorResultException.wrap(result);
        resultHandler.handleErrorResult(ere);
        // doesn't matter if it was canceled.
        requestsInProgress.remove(context);
        return;
      }

      // Get the entry.
      final Entry entry = entryMap.get(dn);
      final AttributeDescription attrDesc = request.getAttributeDescription();
      for (final Attribute attr : entry.getAllAttributes(attrDesc))
      {
        final Iterator<ByteString> it = attr.iterator();
        while (it.hasNext())
        {
          final ByteString s = it.next();
          if (abReq.isCanceled())
          {
            final Result r = Responses.newResult(ResultCode.CANCELLED);
            final ErrorResultException ere = ErrorResultException.wrap(r);
            resultHandler.handleErrorResult(ere);
            requestsInProgress.remove(context);
            return;
          }
          if (s.equals(request.getAssertionValue()))
          {
            result = Responses.newCompareResult(ResultCode.COMPARE_TRUE);
            resultHandler.handleResult(result);
          }
        }
      }
      result = Responses.newCompareResult(ResultCode.COMPARE_FALSE);
      resultHandler.handleResult(result);
      requestsInProgress.remove(context);
    }



    /**
     * @param context
     * @param request
     * @param handler
     * @param intermediateResponseHandler
     * @throws UnsupportedOperationException
     */
    public void handleDelete(final Integer context,
        final DeleteRequest request,
        final ResultHandler<? super Result> handler,
        final IntermediateResponseHandler intermediateResponseHandler)
        throws UnsupportedOperationException
    {
      Result result = null;
      final AbandonableRequest abReq = new AbandonableRequest(request);
      requestsInProgress.put(context, abReq);
      // Get the DN.
      final DN dn = request.getName();
      if (!entryMap.containsKey(dn))
      {
        // entry is not found.
        result = Responses.newResult(ResultCode.NO_SUCH_OBJECT);
        final ErrorResultException ere = ErrorResultException.wrap(result);
        handler.handleErrorResult(ere);
        // doesn't matter if it was canceled.
        requestsInProgress.remove(context);
        return;
      }

      if (abReq.isCanceled())
      {
        result = Responses.newResult(ResultCode.CANCELLED);
        final ErrorResultException ere = ErrorResultException.wrap(result);
        handler.handleErrorResult(ere);
        requestsInProgress.remove(context);
        return;
      }
      // Remove this from the map.
      entryMap.remove(dn);
      requestsInProgress.remove(context);
    }



    /**
     * @param context
     * @param request
     * @param resultHandler
     * @param intermediateResponseHandler
     * @throws UnsupportedOperationException
     */
    public <R extends ExtendedResult> void handleExtendedRequest(
        final Integer context, final ExtendedRequest<R> request,
        final ResultHandler<? super R> resultHandler,
        final IntermediateResponseHandler intermediateResponseHandler)
        throws UnsupportedOperationException
    {
      if (request.getOID().equals(StartTLSExtendedRequest.OID))
      {
        final R result = request.getResultDecoder().newExtendedErrorResult(
            ResultCode.SUCCESS, "", "");
        resultHandler.handleResult(result);
        clientContext.startTLS(sslContext, null, sslContext.getSocketFactory()
            .getSupportedCipherSuites(), false, false);
      }
    }



    /**
     * @param context
     * @param request
     * @param resultHandler
     * @param intermediateResponseHandler
     * @throws UnsupportedOperationException
     */
    public void handleModify(final Integer context,
        final ModifyRequest request,
        final ResultHandler<? super Result> resultHandler,
        final IntermediateResponseHandler intermediateResponseHandler)
        throws UnsupportedOperationException
    {
      // TODO:
    }



    /**
     * @param context
     * @param request
     * @param resultHandler
     * @param intermediateResponseHandler
     * @throws UnsupportedOperationException
     */
    public void handleModifyDN(final Integer context,
        final ModifyDNRequest request,
        final ResultHandler<? super Result> resultHandler,
        final IntermediateResponseHandler intermediateResponseHandler)
        throws UnsupportedOperationException
    {
      // TODO
    }



    /**
     * @param context
     * @param request
     * @param resultHandler
     * @param intermediateResponseHandler
     * @throws UnsupportedOperationException
     */
    public void handleSearch(final Integer context,
        final SearchRequest request, final SearchResultHandler resultHandler,
        final IntermediateResponseHandler intermediateResponseHandler)
        throws UnsupportedOperationException
    {
      Result result = null;
      final AbandonableRequest abReq = new AbandonableRequest(request);
      requestsInProgress.put(context, abReq);
      // Get the DN.
      final DN dn = request.getName();
      if (!entryMap.containsKey(dn))
      {
        // Entry not found.
        result = Responses.newResult(ResultCode.NO_SUCH_OBJECT);
        final ErrorResultException ere = ErrorResultException.wrap(result);
        resultHandler.handleErrorResult(ere);
        // Should searchResultHandler handle anything?

        // doesn't matter if it was canceled.
        requestsInProgress.remove(context);
        return;
      }

      if (abReq.isCanceled())
      {
        result = Responses.newResult(ResultCode.CANCELLED);
        final ErrorResultException ere = ErrorResultException.wrap(result);
        resultHandler.handleErrorResult(ere);
        requestsInProgress.remove(context);
        return;
      }

      final SearchResultEntry e = Responses
          .newSearchResultEntry(new LinkedHashMapEntry(entryMap.get(dn)));
      // Check we have had any controls in the request.
      for (final Control control : request.getControls())
      {
        if (control.getOID().equals(AccountUsabilityRequestControl.OID))
        {
          e.addControl(AccountUsabilityResponseControl.newControl(false, false,
              false, 10, false, 0));
        }
      }
      resultHandler.handleEntry(e);
      result = Responses.newResult(ResultCode.SUCCESS);
      resultHandler.handleResult(result);
      requestsInProgress.remove(context);
    }
  }



  // The mapping between entry DNs and the corresponding entries.
  private final ConcurrentHashMap<DN, Entry> entryMap = new ConcurrentHashMap<DN, Entry>();

  // The grizzly transport.
  private final TCPNIOTransport transport = TransportFactory.getInstance()
      .createTCPTransport();

  // The LDAP listener.
  private LDAPListener listener = null;

  // whether the server is running.
  private volatile boolean isRunning;

  // The mapping between the message id and the requests the server is currently
  // handling.
  private final ConcurrentHashMap<Integer, AbandonableRequest> requestsInProgress = new ConcurrentHashMap<Integer, AbandonableRequest>();

  // The Set used for locking dns.
  private final HashSet<DN> lockedDNs = new HashSet<DN>();

  private SSLContext sslContext;



  private LDAPServer()
  {
    // Add the root dse first.
    entryMap.put(DN.rootDN(),
        Entries.unmodifiableEntry(new LinkedHashMapEntry()));
    for (int i = 0; i < 1000; i++)
    {
      final String dn = String.format("uid=user.%d,ou=people,o=test", i);
      final String cn = String.format("cn: user.%d", i);
      final String sn = String.format("sn: %d", i);
      final String uid = String.format("uid: user.%d", i);

      final DN d = DN.valueOf(dn);
      final Entry e = new LinkedHashMapEntry("dn: " + dn,
          "objectclass: person", "objectclass: inetorgperson",
          "objectclass: top", cn, sn, uid);
      entryMap.put(d, Entries.unmodifiableEntry(e));
    }
  }



  /**
   * @param context
   * @return
   */
  public ServerConnection<Integer> handleAccept(final LDAPClientContext context)
  {
    return new LDAPServerConnection(context);
  }



  /**
   * Returns whether the server is running or not.
   *
   * @return Whether the server is running.
   */
  public boolean isRunning()
  {
    return isRunning;
  }



  /**
   * Starts the server.
   *
   * @param port
   * @exception IOException
   */
  public synchronized void start(final int port) throws Exception
  {
    if (isRunning)
    {
      return;
    }
    sslContext = new SSLContextBuilder().getSSLContext();

    transport.setSelectorRunnersCount(2);
    listener = new LDAPListener(port, getInstance(),
        new GrizzlyLDAPListenerOptions().setTCPNIOTransport(transport)
            .setBacklog(4096));
    transport.start();
    isRunning = true;
  }



  /**
   * Stops the server.
   */
  public synchronized void stop()
  {
    if (!isRunning)
    {
      return;
    }
    listener.close();
    try
    {
      transport.stop();
    }
    catch (final IOException e)
    {
      e.printStackTrace();
    }
    TransportFactory.getInstance().close();
    isRunning = false;
  }
}