#!/usr/bin/env python
"""Manage data servers."""


import atexit
import os
import readline
import time
import urlparse

import urllib3
from urllib3 import connectionpool

# pylint: disable=unused-import,g-bad-import-order
from grr.client import client_plugins
# pylint: enable=unused-import,g-bad-import-order

from grr.lib import config_lib
from grr.lib import flags
from grr.lib import rdfvalue
from grr.lib import startup
from grr.lib import utils

from grr.server.data_server import constants
from grr.server.data_server import errors
from grr.server.data_server import utils as sutils


class Manager(object):
  """Manage a data server group using a connection to the master."""

  def __init__(self):
    servers = config_lib.CONFIG["Dataserver.server_list"]
    if not servers:
      raise errors.DataServerError("List of data servers not available.")
    master_location = servers[0]
    loc = urlparse.urlparse(master_location, scheme="http")
    self.addr = loc.hostname
    self.port = int(loc.port)
    self.pool = connectionpool.HTTPConnectionPool(self.addr, port=self.port)
    self.history_path = os.path.expanduser("~/.grr-data-store-manager")
    if os.path.exists(self.history_path):
      readline.read_history_file(self.history_path)
    self.periodic_thread = None
    self.mapping = None
    self.mapping_time = 0

  def Start(self):
    self._PeriodicThread()
    self.periodic_thread = utils.InterruptableThread(
        target=self._PeriodicThread, sleep_time=10)
    self.periodic_thread.start()
    return True

  def _PeriodicThread(self):
    body = ""
    headers = {"Content-Length": len(body)}
    try:
      res = self.pool.urlopen("POST", "/manage", headers=headers, body=body)
      if res.status != constants.RESPONSE_OK:
        return False
      self.mapping = rdfvalue.DataServerMapping(res.data)
      self.mapping_time = time.time()
    except urllib3.exceptions.MaxRetryError:
      pass

  def SaveHistory(self):
    readline.write_history_file(self.history_path)

  def _ShowServers(self):
    if not self.mapping:
      print "Server information not available"
      return
    last = time.asctime(time.localtime(self.mapping_time))
    print "Last refresh:", last
    for i, serv in enumerate(list(self.mapping.servers)):
      addr = serv.address
      port = serv.port
      size = serv.state.size
      load = serv.state.load
      ncomp = serv.state.num_components
      avgcomp = serv.state.avg_component
      print "Server %d %s:%d (Size: %dKB, Load: %d)" % (i, addr, port,
                                                        size / 1024, load)
      print "\t\t%d components %dKB average size" % (ncomp, avgcomp / 1024)

  def _ShowRanges(self):
    if not self.mapping:
      print "Server information not available"
      return
    self._ShowRange(self.mapping)

  def _ShowRange(self, mapping):
    for i, serv in enumerate(list(mapping.servers)):
      addr = serv.address
      port = serv.port
      start = serv.interval.start
      end = serv.interval.end
      perc = float(end - start) / float(2 ** 64)
      perc *= 100
      print "Server %d %s:%d %d%% [%s, %s[" % (i, addr, port, perc,
                                               str(start).zfill(20),
                                               str(end).zfill(20))

  def _ComputeMappingSize(self, mapping):
    totalsize = 0
    servers = list(mapping.servers)
    for serv in servers:
      totalsize += serv.state.size
    return totalsize

  def _ComputeMappingFromPercentages(self, mapping, newperc):
    """Builds a new mapping based on the new server range percentages."""
    newstart = 0
    n_servers = self.mapping.num_servers
    servers = list(mapping.servers)
    new_mapping = rdfvalue.DataServerMapping(version=self.mapping.version + 1,
                                             num_servers=n_servers,
                                             pathing=self.mapping.pathing)
    for i, perc in enumerate(newperc):
      quant = int(perc * constants.MAX_RANGE)
      interval = rdfvalue.DataServerInterval(start=newstart)
      end = newstart + quant
      if i == len(newperc) - 1:
        end = constants.MAX_RANGE
      interval.end = end
      old_server = servers[i]
      newstart = end
      new_mapping.servers.Append(index=old_server.index,
                                 address=old_server.address,
                                 port=old_server.port,
                                 state=old_server.state,
                                 interval=interval)
    return new_mapping

  def _Rebalance(self):
    """Starts the rebalance process."""
    if not self.mapping:
      print "Server information not available"
      return
    # Compute total size of database.
    servers = list(self.mapping.servers)
    num_servers = len(servers)
    target = 1.0 / float(num_servers)
    perc = [target] * num_servers
    new_mapping = self._ComputeMappingFromPercentages(self.mapping, perc)
    print "The new ranges will be:"
    self._ShowRange(new_mapping)
    print
    self._DoRebalance(new_mapping)

  def _DoRebalance(self, new_mapping):
    """Performs a new rebalancing operation with the master server."""
    print "Contacting master server to start re-sharding...",
    # Send mapping information to master.
    pool = None
    try:
      pool = connectionpool.HTTPConnectionPool(self.addr, port=self.port)
    except urllib3.exceptions.MaxRetryError:
      print "Unable to contact master..."
      return
    body = new_mapping.SerializeToString()
    headers = {"Content-Length": len(body)}
    res = None
    try:
      res = pool.urlopen("POST", "/rebalance/phase1", headers=headers,
                         body=body)
    except urllib3.exceptions.MaxRetryError:
      print "Unable to talk with master..."
      pool.close()
      return
    if res.status != constants.RESPONSE_OK:
      print "Re-sharding cannot be done!"
      return
    rebalance = rdfvalue.DataServerRebalance(res.data)
    print "OK"
    print
    print "The following servers will need to move data:"
    for i, move in enumerate(list(rebalance.moving)):
      print "Server %d moves %dKB" % (i, move / 1024)
    answer = raw_input("Proceed with re-sharding? (y/n) ")
    if answer != "y":
      return
    body = rebalance.SerializeToString()
    headers = {"Content-Length": len(body)}
    try:
      res = pool.urlopen("POST", "/rebalance/phase2", headers=headers,
                         body=body)
    except urllib3.exceptions.MaxRetryError:
      print "Unable to contact server for re-sharding."
      print "Make sure the data servers are up and try again."
      return
    if res.status != constants.RESPONSE_OK:
      print "Could not start copying files for re-sharding"
      print "Make sure the data servers are up and try again."
      return

    try:
      res = pool.urlopen("POST", "/rebalance/commit", headers=headers,
                         body=body)
    except urllib3.exceptions.MaxRetryError:
      print ("Could not commit the re-sharding transaction with id "
             "%s") % rebalance.id
      print "Make sure the data servers are up and then run:"
      print "'recover %s' in order to re-run transaction" % rebalance.id
      return

    if res.status != constants.RESPONSE_OK:
      print "Could not commit the transaction %s" % rebalance.id
      print "Make sure the data servers are up and then run:"
      print "'recover %s' in order to re-run transaction" % rebalance.id
      return

    self.mapping = rdfvalue.DataServerMapping(res.data)

    print "Rebalance with id %s fully performed." % rebalance.id

  def _Recover(self, transid):
    """Completes a rebalancing transaction that was unsuccessful."""
    print "Contacting master about transaction %s..." % transid,
    pool = None
    try:
      pool = connectionpool.HTTPConnectionPool(self.addr, port=self.port)
    except urllib3.exceptions.MaxRetryError:
      print "Unable to contact master..."
      return

    print "OK."

    try:
      body = transid
      headers = {"Content-Length": len(body)}
      res = pool.urlopen("POST", "/rebalance/recover", headers=headers,
                         body=body)
    except urllib3.exceptions.MaxRetryError:
      print "Unable to contact master..."
      return

    if res.status == constants.RESPONSE_TRANSACTION_NOT_FOUND:
      print "Transaction %s was not found" % transid
      return
    if res.status != constants.RESPONSE_OK:
      print "Potential data master error. Giving up..."
      return
    rebalance = rdfvalue.DataServerRebalance(res.data)
    print "Got transaction object %s" % rebalance.id
    answer = raw_input("Proceed with the recover process? (y/n) ")
    if answer != "y":
      return

    body = rebalance.SerializeToString()
    headers = {"Content-Length": len(body)}

    try:
      res = pool.urlopen("POST", "/rebalance/commit", headers=headers,
                         body=body)
    except urllib3.exceptions.MaxRetryError:
      print "Could not commit re-sharding transaction with id %s" % rebalance.id
      print "Make sure the data servers are up and then run:"
      print "'recover %s' in order to re-run transaction" % rebalance.id
      return