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

import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Iterator;
import javax.management.MBeanServer;
import javax.management.Notification;
import javax.management.NotificationFilter;
import javax.management.NotificationListener;
import javax.management.ObjectName;
import javax.servlet.http.Cookie;

import org.jboss.mx.util.MBeanServerLocator;
import org.jboss.util.xml.XmlHelper;
import org.jboss.web.loadbalancer.util.Constants;
import org.jboss.web.loadbalancer.util.Request;
import org.w3c.dom.Element;
import org.jboss.system.*;

/**
 * Base-class for Scheduler
 *
 * @jmx:mbean name="jboss.web.loadbalancer: service=Scheduler"
 *         extends="org.jboss.system.ServiceMBean, org.jboss.web.loadbalancer.scheduler.SchedulerMBean"
 *
 * @author Thomas Peuss <jboss@peuss.de>
 * @version $Revision: 1.6 $
 */
public abstract class AbstractScheduler
    extends ServiceMBeanSupport
    implements AbstractSchedulerMBean, NotificationListener
{
  protected ArrayList hostsUp = new ArrayList();
  protected ArrayList hostsDown = new ArrayList();
  protected String stickyCookieName;
  protected boolean useStickySession = false;
  protected Element config;

  /**
   * Override this method to create a new scheduler.
   * @return
   */
  protected abstract Host getNextHost();

  protected void stopService() throws java.lang.Exception
  {
    // Register Hosts as MBeans
    deregisterHostMBeans();
  }

  /**
   * @jmx:managed-attribute
   * @param config
   */
  public void setConfig(Element config)
  {
    this.config = config;
  }

  protected void createService() throws java.lang.Exception
  {
    // Add all hosts to the host-up-list
    addHostsFromConfig(config);

    // do we use sticky session?
    if (XmlHelper.getUniqueChildContent(config, "sticky-session").
        equalsIgnoreCase("true"))
    {
      useStickySession = true;
      stickyCookieName = XmlHelper.getUniqueChildContent(config,
          "sticky-session-cookie-name");
    }

    if (log.isDebugEnabled())
    {
      if (useStickySession)
      {
        log.debug("Using sticky sessions with Cookie name=" + stickyCookieName);
      }
      else
      {
        log.debug("Using NO sticky sessions");
      }
    }
  }

  protected synchronized void addHostsFromConfig(Element config) throws
      Exception
  {
    Iterator hostIterator = XmlHelper.getChildrenByTagName(XmlHelper.
        getUniqueChild(config, "hosts"), "host");

    while (hostIterator.hasNext())
    {
      Element hostElement = (Element) hostIterator.next();
      String hostUrl = XmlHelper.getUniqueChildContent(hostElement, "host-url");

      // if the host has no trailing slash we add one
      if (!hostUrl.endsWith("/"))
      {
        hostUrl += "/";
      }

      Host host = addHost(hostUrl);

      try
      {
        host.setLbFactor(Integer.parseInt(XmlHelper.getUniqueChildContent(
            hostElement, "lb-factor")));
      }
      catch (Exception e)
      {
        log.debug("Could read LbFactor for Host " + host.getUrl() +
                 " - assuming LbFactor 1");
        host.setLbFactor(1);
      }
    }
  }

  protected ObjectName genObjectNameForHost(Host host) throws Exception
  {
    return new ObjectName("jboss.web.loadbalancer: type=Node, protocol=" +
                          host.getUrl().getProtocol() + ", host=" +
                          host.getUrl().getHost() + ", port=" +
                          (host.getUrl().getPort() == -1 ?
                          host.getUrl().getDefaultPort() :
                          host.getUrl().getPort()));
  }

  protected void registerHostMBean(Host host)
  {
    MBeanServer mbs = MBeanServerLocator.locateJBoss();
    try
    {
      mbs.registerMBean(host, genObjectNameForHost(host));
      host.addNotificationListener(this,
                                   new NotificationFilter()
      {
        public boolean isNotificationEnabled(Notification notification)
        {
          return (notification instanceof HostStateChangedNotification);
        }
      }

      , host);
    }
    catch (Exception ex)
    {
      log.error("Could not register HostMBean", ex);
    }
  }

  protected void deregisterHostMBean(Host host)
  {
    MBeanServer mbs = MBeanServerLocator.locateJBoss();
    try
    {
      mbs.unregisterMBean(genObjectNameForHost(host));
      host.removeNotificationListener(this);
    }
    catch (Exception ex)
    {
      log.error("Could not unregister HostMBean", ex);
    }
  }

  protected synchronized void deregisterHostMBeans() throws Exception
  {
    Iterator hostIterator = hostsUp.iterator();

    while (hostIterator.hasNext())
    {
      Host host = (Host) hostIterator.next();
      deregisterHostMBean(host);
    }

    hostIterator = hostsDown.iterator();
    while (hostIterator.hasNext())
    {
      Host host = (Host) hostIterator.next();
      deregisterHostMBean(host);
    }
  }

  public void getHost(Request schedRequest) throws
      NoHostAvailableException
  {
    Host host = null;

    // first look if we can find the sticky host for this request
    host = findStickyHost(schedRequest);

    // not found -> find a host for this request
    while (host == null)
    {
      if (hostsUp.size() == 0)
      {
        throw new NoHostAvailableException("No host to schedule request");
      }
      host = getNextHost();
    }

    schedRequest.setHost(host);

    // if we use sticky session -> set the cookie
    if (useStickySession)
    {
      setStickyCookie(schedRequest);
    }
  }

  /**
   * Set the sticky session cookie.
   */
  protected void setStickyCookie(Request schedRequest)
  {
    Cookie cookie = new Cookie(stickyCookieName,
                               Integer.toString(schedRequest.getHost().hashCode()));
    cookie.setPath("/");
    cookie.setMaxAge( -1);
    schedRequest.getResponse().addCookie(cookie);
  }

  /**
   * Find the sticky host for the given request
   * @param request The request we want to find the sticky host for
       * @return null=host not found, otherwise the sticky host URL for this request
   */
  protected Host findStickyHost(Request schedRequest)
  {
    Host host = null;
    if (useStickySession)
    {
      Cookie[] cookies = schedRequest.getRequest().getCookies();

      for (int i = 0; cookies != null && i < cookies.length; ++i)
      {
        Cookie cookie = cookies[i];

        if (cookie.getName().equals(stickyCookieName))
        {
          log.debug("Sticky Cookie found!");
          int cookieHash = Integer.parseInt(cookie.getValue());
          Iterator iter = hostsUp.iterator();

          while (iter.hasNext())
          {
            Host tempHost = (Host) iter.next();
            if (tempHost.hashCode() == cookieHash)
            {
              host = tempHost;
              if (log.isDebugEnabled())
              {
                log.debug("Sticky Cookie sticks client to host with URL " +
                          tempHost.toString());
              }
              break;
            }
          }
          break;
        }
      }
      return host;
    }
    else
    {
      return null;
    }
  }

  protected synchronized void markNodeDown(Host host)
  {
    int index = hostsUp.indexOf(host);
    if (index == -1)
    {
      return;
    }
    hostsUp.remove(index);
    hostsDown.add(host);

//    host.setState(Constants.STATE_NODE_DOWN);
  }

  protected synchronized void markNodeUp(Host host)
  {
    int index = hostsDown.indexOf(host);
    if (index == -1)
    {
      return;
    }
    hostsDown.remove(index);
    hostsUp.add(host);

//    host.setState(Constants.STATE_NODE_UP);
  }

  // MBean Interface
  /**
   * Get the list of all hosts that have been marked down.
   * @jmx:managed-attribute
   */
  public ArrayList getHostsDown()
  {
    return hostsDown;
  }

  /**
   * Get the list of all hosts that have been marked up.
   * @jmx:managed-attribute
   */
  public ArrayList getHostsUp()
  {
    return hostsUp;
  }

  /**
   * Add a host to the up list.
   * @jmx:managed-operation
   */
  public synchronized Host addHost(String hostString) throws
      MalformedURLException
  {
    if (hostString == null)
    {
      return null;
    }
    if (!hostString.endsWith("/"))
    {
      hostString += "/";
    }
    Host host = new Host(new URL(hostString));

    // Host already added?
    if (hostsUp.indexOf(host) > -1 || hostsDown.indexOf(host) > -1)
    {
      log.info("Host " + hostString + " already there. Ignored");
      return null;
    }

    hostsUp.add(host);
    registerHostMBean(host);

    return host;
  }

  /**
   * Remove a host from the up list.
   * @jmx:managed-operation
   */
  public synchronized void removeHost(URL url)
  {
    if (url == null)
    {
      return;
    }
    Host host = new Host(url);

    int index = hostsUp.indexOf(host);
    if (index > -1)
    {
      deregisterHostMBean( (Host) hostsUp.get(index));
      hostsUp.remove(index);
    }

    index = hostsDown.indexOf(host);
    if (index > -1)
    {
      deregisterHostMBean( (Host) hostsDown.get(index));
      hostsDown.remove(index);
    }
  }

  public void handleNotification(Notification notification, Object handback)
  {
    HostStateChangedNotification hscn=(HostStateChangedNotification)notification;

    Host host=(Host)hscn.getSource();

    switch (host.getState())
    {
      case Constants.STATE_NODE_UP:
        markNodeUp(host);
        break;
      case Constants.STATE_NODE_DOWN:
      case Constants.STATE_NODE_FORCED_DOWN:
        markNodeDown(host);
      break;
    }
  }
}