/*
 * JBoss, the OpenSource WebOS
 *
 * Distributable under LGPL license.
 * See terms of license at gnu.org.
 */
package org.jboss.web.loadbalancer;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Set;
import javax.management.ObjectName;
import javax.servlet.ServletException;
import javax.servlet.http.Cookie;

import org.apache.commons.httpclient.Header;
import org.apache.commons.httpclient.HttpClient;
import org.apache.commons.httpclient.HttpMethod;
import org.apache.commons.httpclient.HttpState;
import org.apache.commons.httpclient.MultiThreadedHttpConnectionManager;
import org.apache.commons.httpclient.cookie.CookiePolicy;
import org.apache.commons.httpclient.methods.DeleteMethod;
import org.apache.commons.httpclient.methods.GetMethod;
import org.apache.commons.httpclient.methods.HeadMethod;
import org.apache.commons.httpclient.methods.OptionsMethod;
import org.apache.commons.httpclient.methods.PostMethod;
import org.apache.commons.httpclient.methods.PutMethod;
import org.jboss.logging.Logger;
import org.jboss.web.loadbalancer.scheduler.NoHostAvailableException;
import org.jboss.web.loadbalancer.scheduler.SchedulerMBean;
import org.jboss.web.loadbalancer.util.Constants;
import org.jboss.web.loadbalancer.util.Request;

/**
 * The Loadbalancer core class.
 *
 * @jmx:mbean name="jboss.web.loadbalancer: service=Loadbalancer"
 *
 * @author Thomas Peuss <jboss@peuss.de>
 * @version $Revision: 1.9 $
 */
public class Loadbalancer
    implements LoadbalancerMBean
{
  protected static Logger log = Logger.getLogger(Loadbalancer.class);

  // The connection manager
  protected MultiThreadedHttpConnectionManager connectionManager;

  // request header elements that must not be copied to client
  protected static Set ignorableHeader = new HashSet();

  protected int connectionTimeout = 20000;

  protected SchedulerMBean scheduler;
  protected ObjectName schedulerName;

  protected static final int MAX_RETRIES = 5;

  static
  {
    // this header elements are not copied from
    // the HttpClient response to the request client.
    // and vice versa because they are generated by
    // the servlet-engine
    ignorableHeader.add("content-length");
    ignorableHeader.add("server");
    ignorableHeader.add("transfer-encoding");
    ignorableHeader.add("cookie");
    ignorableHeader.add("set-cookie");
    ignorableHeader.add("host");
  }

  protected Loadbalancer(SchedulerMBean scheduler, int timeout) throws ServletException
  {
    this.scheduler=scheduler;
    this.connectionTimeout=timeout;

    connectionManager = new MultiThreadedHttpConnectionManager();

    // We disable this because the web-container limits the maximum connection count anyway
    connectionManager.setMaxConnectionsPerHost(Integer.MAX_VALUE);
    connectionManager.setMaxTotalConnections(Integer.MAX_VALUE);
  }

  /**
   * Create a HttpMethod object for the given request.
   * @param request
   * @param response
   * @param requestMethod
   * @return
   * @throws NoHostAvailableException
   */
  protected void createMethod(Request schedRequest) throws NoHostAvailableException
  {
    String url = null;
    HttpMethod method = null;

    scheduler.getHost(schedRequest);

    // get target host from scheduler
    url = schedRequest.getHost().getUrl().toExternalForm();

    String path = url.substring(0, url.length() - 1) + schedRequest.getRequest().getRequestURI();

    switch (schedRequest.getRequestMethod())
    {
      case Constants.HTTP_METHOD_GET:
        method = new GetMethod(path);
        break;
      case Constants.HTTP_METHOD_POST:
        method = new PostMethod(path);
        break;
      case Constants.HTTP_METHOD_DELETE:
        method = new DeleteMethod(path);
        break;
      case Constants.HTTP_METHOD_HEAD:
        method = new HeadMethod(path);
        break;
      case Constants.HTTP_METHOD_OPTIONS:
        method = new OptionsMethod(path);
        break;
      case Constants.HTTP_METHOD_PUT:
        method = new PutMethod(path);
        break;
      default:
        throw new IllegalStateException("Unknown Request Method " +
                                        schedRequest.getRequest().getMethod());
    }
    schedRequest.setMethod(method);
  }

  /**
   * Add the request information to the HttpMethod
   * @param request
   * @param method
   * @return
   */
  protected void addRequestData(Request schedRequest)
  {
    HttpMethod method=schedRequest.getMethod();

    // add GET-data to query string
    if (schedRequest.getRequest().getQueryString() != null)
    {
      method.setQueryString(schedRequest.getRequest().getQueryString());
    }

    // add POST-data to the request
    if (method instanceof PostMethod)
    {
      PostMethod postMethod = (PostMethod) method;

      Enumeration paramNames = schedRequest.getRequest().getParameterNames();
      while (paramNames.hasMoreElements())
      {
        String paramName = (String) paramNames.nextElement();
        postMethod.addParameter(paramName, schedRequest.getRequest().getParameter(paramName));
      }
    }
  }

  /**
   * Prepare the request to the target node.
   * @param request
   * @param response
   * @param method
   * @return
   */
  protected void prepareServerRequest(Request schedRequest)
  {
    HttpMethod method=schedRequest.getMethod();

    // Initialize client
    HttpClient client = new HttpClient(connectionManager);
    client.setStrictMode(false);
    client.setTimeout(connectionTimeout);
    client.setConnectionTimeout(connectionTimeout);
    client.getState().setCookiePolicy(CookiePolicy.COMPATIBILITY);

    // Initialize Method
    method.setFollowRedirects(false);
    method.setDoAuthentication(false);

    // Add request header to request (minus ignored header values)
    Enumeration reqHeaders = schedRequest.getRequest().getHeaderNames();
    while (reqHeaders.hasMoreElements())
    {
      String headerName = (String) reqHeaders.nextElement();
      String headerValue = schedRequest.getRequest().getHeader(headerName);

      if (!ignorableHeader.contains(headerName.toLowerCase()))
      {
        method.setRequestHeader(headerName, headerValue);
      }
    }

    // Copy cookies into the request
    Cookie[] cookies = schedRequest.getRequest().getCookies();
    HttpState state = client.getState();
    for (int i = 0; cookies != null && i < cookies.length; ++i)
    {
      Cookie cookie = cookies[i];

      org.apache.commons.httpclient.Cookie reqCookie =
          new org.apache.commons.httpclient.Cookie();

      reqCookie.setName(cookie.getName());
      reqCookie.setValue(cookie.getValue());

      // patch cookie path because the HttpClient does not like null paths
      if (cookie.getPath() != null)
      {
        reqCookie.setPath(cookie.getPath());
      }
      else
      {
        reqCookie.setPath("/");
      }

      reqCookie.setSecure(cookie.getSecure());

      reqCookie.setDomain(method.getHostConfiguration().getHost());
      state.addCookie(reqCookie);
    }
    schedRequest.setClient(client);
  }

  /**
   * Handle the client request.
   * @param request
   * @param response
   * @param method
   * @throws ServletException
   * @throws IOException
   */
  protected void handleRequest(Request schedRequest) throws ServletException, IOException
  {

    boolean reschedule = false;

    try
    {
      // prepare the request
      prepareServerRequest(schedRequest);

      int tries = 0;
      // we try several times before we give up
      while (tries < MAX_RETRIES)
      {
        try
        {
          // GO
          long t1=System.currentTimeMillis();
          schedRequest.getHost().incCurrentConnections();

          schedRequest.getClient().executeMethod(schedRequest.getMethod());

          schedRequest.getHost().decCurrentConnections();

          long t2=System.currentTimeMillis();
          schedRequest.getHost().addRequest((int)(t2-t1));
          break;
        }
        catch (IOException ex)
        {
          try
          {
            schedRequest.getHost().decCurrentConnections();
            schedRequest.getMethod().recycle();
          }
          catch (Exception e)
          {
            //ignore
          }

          tries++;
          log.info("Connect retry no. " + tries, ex);
        }
      }
      // everything ok?
      if (tries < MAX_RETRIES)
      {
        // generate the response for the
        parseServerResponse(schedRequest);
      }
      else
      {
        log.error("Max retries reached - giving up. Host will be marked down");

        // Inform Scheduler of node problems
        schedRequest.getHost().markNodeDown();
        reschedule = true;
      }
    }
    finally
    {
      try
      {
        schedRequest.getMethod().recycle();
      }
      catch (Exception e)
      {
        //ignore
      }
    }
    // try again?
    if (reschedule)
    {
      String redirectURI = schedRequest.getRequest().getRequestURI();

      if (schedRequest.getRequest().getQueryString() != null)
      {
        redirectURI += "?" + schedRequest.getRequest().getQueryString();
      }
      // send redirect to force client request
      schedRequest.getResponse().sendRedirect(redirectURI);
    }
  }

  /**
   * Copy the server answer meta data to the client.
   * @param request
   * @param response
   * @param client
   * @param method
   * @throws ServletException
   * @throws IOException
   */
  protected void parseServerResponse(Request schedRequest) throws ServletException, IOException
  {
    schedRequest.getResponse().setStatus(schedRequest.getMethod().getStatusCode());

    //Cookies
    org.apache.commons.httpclient.Cookie[] respCookies =
        schedRequest.getClient().getState().getCookies();

    for (int i = 0; i < respCookies.length; ++i)
    {
      Cookie cookie =
          new Cookie(respCookies[i].getName(), respCookies[i].getValue());

      if (respCookies[i].getPath() != null)
      {
        cookie.setPath(respCookies[i].getPath());
      }
      schedRequest.getResponse().addCookie(cookie);
    }

    //Header
    Header[] header = schedRequest.getMethod().getResponseHeaders();

    for (int i = 0; i < header.length; ++i)
    {
      if (!ignorableHeader.contains(header[i].getName().toLowerCase()))
      {
        schedRequest.getResponse().setHeader(header[i].getName(), header[i].getValue());
      }
    }

    copyServerResponse(schedRequest);
  }

  /**
   * Copy content to the client.
   * @param response
   * @param method
   * @throws IOException
   */
  protected void copyServerResponse(Request schedRequest) throws IOException
  {

    InputStream bodyStream = schedRequest.getMethod().getResponseBodyAsStream();

    // any response?
    if (bodyStream == null)
    {
      log.debug("No request body");
      return;
    }

    byte[] buffer = new byte[2048];
    int numBytes;
    OutputStream out = schedRequest.getResponse().getOutputStream();

    // copy the response
    while ( (numBytes = bodyStream.read(buffer)) != -1)
    {
      out.write(buffer, 0, numBytes);
      if (log.isDebugEnabled())
      {
        log.debug("Copied " + numBytes + " bytes");
      }
    }
  }

  // MBean Interface
  /**
   * Get the currently used connection timeout to slave hosts.
   * @jmx:managed-attribute
   */
  public int getConnectionTimeout()
  {
    return this.connectionTimeout;
  }

  /**
   * Set the currently used connection timeout to slave hosts.
   * @jmx:managed-attribute
   */
  public void setConnectionTimeout(int newTimeout)
  {
    this.connectionTimeout = newTimeout;
  }

  /**
   * Get the currently used connections to slave hosts.
   * @jmx:managed-attribute
   */
  public int getConnectionsInUse()
  {
    return connectionManager.getConnectionsInUse();
  }
}