/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.hadoop.chukwa.datacollection.agent;


import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.PrintStream;
import java.net.*;
import java.nio.charset.Charset;
import java.util.Map;

import org.apache.hadoop.chukwa.datacollection.adaptor.AdaptorException;
import org.apache.hadoop.chukwa.datacollection.adaptor.AdaptorShutdownPolicy;
import org.apache.log4j.Logger;
import org.apache.hadoop.chukwa.util.ExceptionUtil;

/**
 * Class to handle the agent control protocol. This is a simple line-oriented
 * ASCII protocol, that is designed to be easy to work with both
 * programmatically and via telnet.
 * 
 * The port to bind to can be specified by setting option
 * chukwaAgent.agent.control.port. A port of 0 creates a socket on any free
 * port.
 */
public class AgentControlSocketListener extends Thread {

  static Logger log = Logger.getLogger(AgentControlSocketListener.class);

  protected ChukwaAgent agent;
  protected int portno;
  protected ServerSocket s = null;
  volatile boolean closing = false;
  static final String VERSION = "0.4.0-dev";
  public boolean ALLOW_REMOTE = true;
  public static final String REMOTE_ACCESS_OPT = "chukwaAgent.control.remote";

  private class ListenThread extends Thread {
    Socket connection;

    ListenThread(Socket conn) {
      connection = conn;
      try {
        connection.setSoTimeout(60000);
      } catch (SocketException e) {
        log.warn("Error while settin soTimeout to 60000");
        e.printStackTrace();
      }
      this.setName("listen thread for " + connection.getRemoteSocketAddress());
    }

    public void run() {
      try {
        InputStream in = connection.getInputStream();
        BufferedReader br = new BufferedReader(new InputStreamReader(in, Charset.forName("UTF-8")));
        PrintStream out = new PrintStream(new BufferedOutputStream(connection
            .getOutputStream()), true, "UTF-8");
        String cmd = null;
        while ((cmd = br.readLine()) != null) {
          processCommand(cmd, out);
        }
        connection.close();
        if (log.isDebugEnabled()) {
          log.debug("control connection closed");
        }
      } catch (SocketException e) {
        if (e.getMessage().equals("Socket Closed"))
          log.info("control socket closed");
      } catch (IOException e) {
        log.warn("a control connection broke", e);
        try {
          connection.close();
        } catch(Exception ex) {
          log.debug(ExceptionUtil.getStackTrace(ex));
        }
      }
    }

    /**
     * process a protocol command
     * 
     * @param cmd the command given by the user
     * @param out a PrintStream writing to the socket
     * @throws IOException
     */
    public void processCommand(String cmd, PrintStream out) throws IOException {
      String[] words = cmd.split("\\s+");
      if (log.isDebugEnabled()) {
        log.debug("command from " + connection.getRemoteSocketAddress() + ":"
            + cmd);
      }

      if (words[0].equalsIgnoreCase("help")) {
        out.println("you're talking to the Chukwa agent.  Commands available: ");
        out.println("add [adaptorname] [args] [offset] -- start an adaptor");
        out.println("shutdown [adaptornumber]  -- graceful stop");
        out.println("stop [adaptornumber]  -- abrupt stop");
        out.println("list -- list running adaptors");
        out.println("close -- close this connection");
        out.println("stopagent -- stop the whole agent process");
        out.println("stopall -- stop all adaptors");
        out.println("reloadCollectors -- reload the list of collectors");
        out.println("help -- print this message");
        out.println("\t Command names are case-blind.");
      } else if (words[0].equalsIgnoreCase("close")) {
        connection.close();
      } else if (words[0].equalsIgnoreCase("add")) {
        try {
          String newID = agent.processAddCommandE(cmd);
          if (newID != null)
            out.println("OK add completed; new ID is " + newID);
          else
            out.println("failed to start adaptor...check logs for details");
        } catch(AdaptorException e) {
          out.println(e);
        }
      } else if (words[0].equalsIgnoreCase("shutdown")) {
        if (words.length < 2) {
          out.println("need to specify an adaptor to shut down, by number");
        } else {
          sanitizeAdaptorName(out, words);
          long offset = agent.stopAdaptor(words[1], AdaptorShutdownPolicy.GRACEFULLY);
          if (offset != -1)
            out.println("OK adaptor " + words[1] + " stopping gracefully at "
                + offset);
          else
            out.println("FAIL: perhaps adaptor " + words[1] + " does not exist");
        }
      } else if (words[0].equalsIgnoreCase("stop")) {
        if (words.length < 2) {
          out.println("need to specify an adaptor to shut down, by number");
        } else {
          sanitizeAdaptorName(out, words);
          agent.stopAdaptor(words[1], AdaptorShutdownPolicy.HARD_STOP);
          out.println("OK adaptor " + words[1] + " stopped");
        }
      } else if (words[0].equalsIgnoreCase("reloadCollectors")) {
        agent.getConnector().reloadConfiguration();
        out.println("OK reloadCollectors done");
      } else if (words[0].equalsIgnoreCase("list")) {
        java.util.Map<String, String> adaptorList = agent.getAdaptorList();

        if (log.isDebugEnabled()) {
          log.debug("number of adaptors: " + adaptorList.size());
        }

        for (Map.Entry<String, String> a: adaptorList.entrySet()) {
            out.print(a.getKey());
            out.print(") ");
            out.print(" ");
            out.println(a.getValue());
          }
          out.println("");
        
      } else if (words[0].equalsIgnoreCase("stopagent")) {
        out.println("stopping agent process.");
        connection.close();
        agent.shutdown(true);
      } else if(words[0].equalsIgnoreCase("stopall")) {
        int stopped = 0;
        for(String id: agent.getAdaptorList().keySet()) {
          agent.stopAdaptor(id, false);
          stopped++;
         }
        out.println("stopped " + stopped + " adaptors");
      } else if (words[0].equals("")) {
        out.println(getStatusLine());
      } else {
        log.warn("unknown command " + words[0]);
        out.println("unknown command " + words[0]);
        out.println("say 'help' for a list of legal commands");
      }
      out.flush();
    }

    private void sanitizeAdaptorName(PrintStream out, String[] words) {
      if(!words[1].startsWith("adaptor_")) {
        words[1] = "adaptor_" + words[1];
        out.println("adaptor names should start with adaptor_; "
            +"assuming you meant"+ words[1] );
      }
    }

  }

  /**
   * Initializes listener, but does not bind to socket.
   * @param agent the agent to control
   */
  public AgentControlSocketListener(ChukwaAgent agent) {

    this.setDaemon(false); // to keep the local agent alive
    this.agent = agent;
    this.portno = agent.getConfiguration().getInt("chukwaAgent.control.port",
        9093);
    this.ALLOW_REMOTE = agent.getConfiguration().getBoolean(REMOTE_ACCESS_OPT, ALLOW_REMOTE);
    log.info("AgentControlSocketListerner ask for port: " + portno);
    this.setName("control socket listener");
  }


  /**
   * Binds to socket, starts looping listening for commands
   */
  public void run() {
    try {
      if (!isBound())
        tryToBind();
    } catch (IOException e) {
      return;
    }

    while (!closing) {
      try {
        Socket connection = s.accept();
        if (log.isDebugEnabled()) {
          log.debug("new connection from " + connection.getInetAddress());
        }
        ListenThread l = new ListenThread(connection);
        l.setDaemon(true);
        l.start();
      } catch (IOException e) {
        if (!closing)
          log.warn("control socket error: ", e);
        else {
          log.warn("shutting down listen thread due to shutdown() call");
          break;
        }
      }
    }// end while
  }

  /**
   * Close the control socket, and exit. Triggers graceful thread shutdown.
   */
  public void shutdown() {
    closing = true;
    try {
      if (s != null)
        s.close();
      s = null;
    } catch (IOException e) {
      log.debug(ExceptionUtil.getStackTrace(e));
    } // ignore exception on close
  }

  public boolean isBound() {
    return s != null && s.isBound();
  }

  public void tryToBind() throws IOException {
    if(ALLOW_REMOTE)
      s = new ServerSocket(portno);
    else {  //FIXME: is there a way to allow all local addresses? (including IPv6 local)
      s = new ServerSocket();
      s.bind(new InetSocketAddress(InetAddress.getByAddress(new byte[] {127,0,0,1}), portno));
    }
    s.setReuseAddress(true);
    portno = s.getLocalPort();
    if (s.isBound())
      log.info("socket bound to " + s.getLocalPort());
    else
      log.info("socket isn't bound");
  }

  public int getPort() {
    if (!s.isBound()) {
      return -1;
    } else {
      return portno;
    }
  }
  
  //FIXME: we also do this in ChunkImpl; should really do it only once
  //and make it visible everywhere?
  private static String localHostAddr;
  static {
    try {
      localHostAddr = InetAddress.getLocalHost().getHostName();
    } catch (UnknownHostException e) {
      localHostAddr = "localhost";
    }
  }
  
  public String getStatusLine() {
    int adaptorCount = agent.adaptorCount();
    
    return localHostAddr + ": Chukwa Agent running, version " + VERSION + ", with " + adaptorCount + " adaptors";
  }
}