/*
 * JBoss, Home of Professional Open Source
 * Copyright 2005, JBoss Inc., and individual contributors as indicated
 * by the @authors tag. See the copyright.txt in the distribution for a
 * full listing of individual contributors.
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2.1 of
 * the License, or (at your option) any later version.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, write to the Free
 * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301 USA, or see the FSF site: http://www.fsf.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.metadata.MetaData; 
import org.jboss.mx.util.MBeanServerLocator;
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.5.2.3 $
 */
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 (MetaData.getUniqueChildContent(config, "sticky-session").
        equalsIgnoreCase("true"))
    {
      useStickySession = true;
      stickyCookieName = MetaData.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 = MetaData.getChildrenByTagName(MetaData.
        getUniqueChild(config, "hosts"), "host");

    while (hostIterator.hasNext())
    {
      Element hostElement = (Element) hostIterator.next();
      String hostUrl = MetaData.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(MetaData.getUniqueChildContent(
            hostElement, "lb-factor")));
      }
      catch (Exception e)
      {
        log.info("Could read LbFactor for Host " + host.getUrl() +
                 " - assuming LbFactor 1");
        host.setLbFactor(1);
      }
    }
  }

  protected ObjectName genObjectNameForHost(Host host) throws Exception
  {
    String port;

    // try to defer port for given URL
    if (host.getUrl().getPort() == -1)
    {
       if (host.getUrl().getProtocol().equals("http"))
       {
          port="80";
       }
       else if (host.getUrl().getProtocol().equals("https"))
       {
          port="443";
       }
       else
       {
          port="999999";
       }
    }
    else
    {
       port=Integer.toString(host.getUrl().getPort());
    }

    return new ObjectName("jboss.web.loadbalancer: type=Node, protocol=" +
                          host.getUrl().getProtocol() + ", host=" +
                          host.getUrl().getHost() + ", port=" +
                          port);
  }

  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;
    }
  }
}
