import os
import sys
import time
import traceback
import csv
import collections

from proj import project

from wingbase import spawn
from wingbase import destroyable
from wingbase import timeout

from plugins.container import api

import config

_ = api._

# This plugin is only available on Linux
def PluginMetaData():
  return {
    'enabled': sys.platform.startswith('linux'),
  }

kInstanceLaunchTimeout = 10.0

gInstances = {}
class CContainerInstance(destroyable.CDestroyable, api.CContainerInstance):
  """Support for launching an LXC container instance and commands."""
  
  def __init__(self, instance_id, image_id, io_cb, pyexec, env={}, mappings=[], port_forwards=None,
               extra_run_args=''):

    destroyable.CDestroyable.__init__(self)
    api.CContainerInstance.__init__(self, instance_id, image_id, io_cb, env, mappings, port_forwards,
                                    extra_run_args)
    
    self.__did_launch = False
    gInstances[self.instance_id] = self
    
    self._valid = True
    self.__status = self.kTerminated
    self.__starting = False
    self.__launch_cbs = set()
    self.__image_time = None
    self.__launch_time = None
    self._terminating = False
    self.__terminate_cbs = set()
    
    # Ignored / not needed in LXC
    # self.pyexec = pyexec

  def _destroy_impl(self):
    
    if gInstances.get(self.instance_id, None) is self:
      del gInstances[self.instance_id]
      
    self.Terminate()
    
  def IsRunning(self):

    return not self.destroyed() and (self.__status == self.kRunning or self._terminating)
  
  def GetStatus(self):
    
    return self.__status
  
  def _UpdateStatus(self, status):
    
    self.__status = status
    CContainerSystem._CallStatusCallbacks(self.instance_id, status)
    
  @classmethod
  def _InstanceList(cls, result_cb):
    """Get a list of running image_ids"""
    
    cmd = ['lxc', 'list', 'c', 'ns', '--format', 'csv']
    
    def parse_result(txt):
      retval = []
      for line in txt.splitlines()[1:]:
        line = line.strip()
        image_id, status = line.split(',')
        if status.lower().strip() == 'running':
          retval.append(image_id.strip())
      return retval
    
    def log_error(err):
      print(err)
      print(_("Output:"))
      print(proc.GetOutput())
      result_cb(err, None)
      
    def terminated(p):
      if proc.GetExitCode() != 0:
        log_error(_("Failed to run 'lxc list'"))
        return
      try:
        retval = parse_result(proc.GetOutput())
      except:
        log_error(_("Failed to parse 'lxc list' output"))
        return
      
      result_cb(None, retval)
      
    def start_failed(p, exc):
      log_error(_("Failed to start 'lxc list'"))
      
    proc = spawn.CChildProcess(cmd, env=config.gStartupEnv)
    proc.connect('terminated', terminated)
    proc.connect('start-failed', start_failed)

    proc.Start()

  def Launch(self, result_cb):

    assert not self._terminating and self._valid
    
    if self.IsRunning():
      result_cb(self.kSuccess, self)
      return

    self.__launch_cbs.add(result_cb)
    
    if self.__starting:
      return
    self.__starting = True
    
    self._UpdateStatus(self.kStarting)
    
    if self.__image_time is None:
      def inspect_complete(modtime):
        self.__image_time = modtime
        self.__launch_time = modtime
        self.__DeconfigureAndLaunch()
      CContainerSystem._GetContainerImageCreationTime(self.image_id, inspect_complete)
      return
    
    else:
      self.__DeconfigureAndLaunch()
      
  def __LaunchInstance(self):
    
    cmd = ['lxc', 'start', self.image_id]
    
    def terminated(p):
      
      if proc.GetExitCode() != 0:
        self.io_cb('error', _("LXC instance {} failed to start; err={}").format(self.image_id,
                                                                                proc.GetExitCode()))
        self.__CallLaunchCBs(self.kStartFailed)
        self._UpdateStatus(self.kStartFailed)
        
      else:
        self.__WaitForIPAfterLaunch()

    def start_failed(p, exc_text):
      self.io_cb('error', _("Failed to start LXC instance: {}").format(cmd))
      self.io_cb('error', exc_text)
      self.__CallLaunchCBs(self.kProcessFailedToStart)

    def output(s, txt):        
      self.io_cb('output', txt)
      
    proc = spawn.CChildProcess(cmd, env=config.gStartupEnv)
    proc.connect_while_alive('terminated', terminated, self)
    proc.connect_while_alive('start-failed', start_failed, self)
    proc.connect_while_alive('received-output', output, self)

    self.io_cb('info', _("Running: {}").format(spawn._win32_cmd_line_from_argv(cmd)))
    proc.Start()
    
    # Wait for launch to complete so we timeout if nothing happens
    max_time = time.time() + kInstanceLaunchTimeout
    
    def check_for_launch():

      # Stop waiting if did launch successfully
      if self.__did_launch:
        return False

      # Call launch cbs with error if timed out
      if time.time() > max_time:
        self.io_cb('error', _("lxc start failed to start - timeout"))
        self.__CallLaunchCBs(self.kProcessFailedToStart)
        return False

      # Keep waiting
      else:
        return True
    
    timeout.timeout_add_while_alive(100, check_for_launch, self)
    
  def __WaitForIPAfterLaunch(self):
    
    end_time = time.time() + 10
    
    def check_ip():
    
      if self.GetContainerIP():
        # XXX The container may actually not be internally ready, but hard to detect this
        # XXX Could try 'lxc exec <container> echo something' although this seems to
        # XXX work instantly even if container is still starting when console is viewed.
        # XXX Waiting for DHCP to get an IP is probably usually enough, though.
        # XXX See also https://blog.simos.info/how-to-know-when-a-lxd-container-has-finished-starting-up/
        self.__CallLaunchCBs(self.kSuccess)
        self._UpdateStatus(self.kRunning)
        return False
        
      elif time.time() > end_time:
        self.__CallLaunchCBs(self.kConfigurationError)
        self._UpdateStatus(self.kStartFailed)
        return True
      
      else:
        return True
      
    # Call immediately for cases where there's no DHCP involved; add timeout to poll
    # if does not yet have IP
    if check_ip():
      timeout.timeout_add_while_alive(100, check_ip, self)
    
  def Terminate(self, completed_cb=None):
    
    if self.destroyed() or self.__status != self.kRunning:
      return
    
    if completed_cb is not None:
      self.__terminate_cbs.add(completed_cb)

    if self._terminating:
      return
    
    self._UpdateStatus(self.kTerminating)
    
    cmd = ['lxc', 'stop', self.image_id]

    def terminated(p):
      self.io_cb('info', _("Terminated LXC container {}").format(self.image_id))
      def remove_complete():
        self._valid = False
        self.__CallTerminateCBs()
        self.destroy()
      self._DeconfigureInstance(remove_complete)
      
    def start_failed(p, exc_text):
      self._terminating = False
      self.__terminate_cbs.remove(completed_cb)
      self.io_cb('error', _("Failed to run 'lxc stop {}'").format(self.image_id))
      self.io_cb('error', exc_text)
      
    def output(p, txt):
      self.io_cb('output', txt)
      
    proc = spawn.CChildProcess(cmd, env=config.gStartupEnv)
    proc.connect_while_alive('terminated', terminated, self)
    proc.connect_while_alive('start-failed', start_failed, self)
    proc.connect_while_alive('received-output', output, self)

    self._terminating = True
      
    self.io_cb('info', _("Running: {}").format(spawn._win32_cmd_line_from_argv(cmd)))
    proc.Start()
    
  # Keeps a class registry for deconfiguring so multiple calls to different
  # instances are handled correctly (as occurs when restarting the container;
  # the old instance deconfigures and the new one also tries to)
  gDeconfiguring = collections.defaultdict(list)
  def _DeconfigureInstance(self, remove_complete_cb=None):
    """Undo configuration put into place by __ConfigureAndLaunch"""
    
    if remove_complete_cb is None:
      def remove_complete_cb():
        print("Configuration for instance_id {} removed".format(self.image_id))

    if self.image_id in self.gDeconfiguring:
      self.gDeconfiguring[self.image_id].append(remove_complete_cb)
      return
    self.gDeconfiguring[self.image_id].append(remove_complete_cb)
    
    print("Removing configuration for instance_id {}".format(self.image_id))
      
    # Make sure the instance is stopped, in case it was accidentally left running
    cmds = []
    cmds.append(['lxc', 'stop', self.image_id])
    
    # Unset environment variables added during launch
    for key in self.env:
      cmds.append(['lxc', 'config', 'unset', 'environment.{}'.format(key)])
      
    # Remove file mounts added during launch
    for i in range(0, len(self.mappings)):
      cmds.append(['lxc', 'config', 'device', 'remove', self.image_id, 'wingmt{}'.format(i)])
      
    # Remove network and nic added during launch
    # XXX Hopefully not needed after all
    #cmds.append(['lxc', 'config', 'device', 'remove', self.image_id, 'wing-eth0'])
    #cmds.append(['lxc', 'network', 'delete', '{}-br0'.format(self.image_id)])
    
    # Add directories
    cmds = [(None, c) for c in cmds]
    
    # Run all processes in a series, logging any failures along the way
    proc = spawn.CChildProcessSequence(cmds, failures='continue',
                                       env=project.GetEffectiveEnvironment())
      
    def starting(dirname, cmd):
      print(_("Starting command {} in {}").format(cmd, dirname))
      
    def terminated(dirname, cmd, errno, exc_text):
      if errno == 'start-failed':
        print(_("Process failed to start"))
        print(exc_text)
      elif errno != 0:
        print(_("Process failed with errno={}").format(errno))

    def output(dirname, cmd, txt):
      print(txt.rstrip())

    def completed(success):
      if not success:
        print(_("Deconfigure failed"))
      else:
        print(_("Deconfigure complete"))
      cbs = self.gDeconfiguring[self.image_id]
      del self.gDeconfiguring[self.image_id]
      for cb in cbs:
        try:
          cb()
        except:
          traceback.print_exc()
            
    proc.connect_while_alive('starting', starting, self)
    proc.connect_while_alive('terminated', terminated, self)
    proc.connect_while_alive('received-output', output, self)
    proc.connect_while_alive('completed', completed, self)

    print(_("Deconfiguring container").format(self.instance_id))
    proc.Start()
    
  def GetImageVersions(self, result_cb):

    def inspect_complete(modtime):
      result_cb(modtime, self.__launch_time)
      self.__image_time = modtime
      
    CContainerSystem._GetContainerImageCreationTime(self.image_id, inspect_complete)
    
  def GetCommandLine(self, cmdline, startdir=None, env={}, pseudo_tty=False):

    # If this is raised, caller needs to start the container first and wait for that
    assert self._valid and self.__status == self.kRunning
    
    cmd = ['lxc', 'exec', self.image_id]

    self._AddPtyToCmdLine(cmd, pseudo_tty)
    self._AddStartDirToCmdLine(cmd, startdir)
    self._AddEnvToCmdLine(cmd, env)
    
    cmd.append('--')
    cmd.extend(cmdline)

    return None, cmd

  def GetHostSystemAddress(self):
    
    ip = self.GetContainerIP()
    if not ip:
      return ''

    # XXX This is a bit dicey but should always work for ip4
    ip_parts = ip.split('.')
    ip = '.'.join(ip_parts[:-1]) + '.1'
    return ip
  
  def GetContainerIP(self):
    
    cmd = ['lxc', 'list', '--columns', 'n4', '--format', 'csv']
    output = spawn.ReadWaitOnPipe(cmd[0], None, 10.0, *cmd[1:])
    for parts in csv.reader(output.splitlines()):
      try:
        name, ip = parts
      except:
        continue
      if name == self.image_id:
        if ' ' in ip:
          ip = ip[:ip.find(' ')]
          ip_parts = ip.split('.')
          if len(ip_parts) != 4:
            continue
          return ip

    return ''
    
  def __CallLaunchCBs(self, err):
    
    launch_cbs = self.__launch_cbs

    self.__starting = False
    self.__launch_cbs = set()
    self.__launch_time = self.__image_time
    
    if err != self.kSuccess:
      self._valid = False
      self._UpdateStatus(self.kStartFailed)
    else:
      self.__did_launch = True      
      self._UpdateStatus(self.kRunning)
    
    for cb in launch_cbs:
      try:
        cb(err, self)
      except:
        if __debug__:
          raise
        print("Launch callback failed", cb)
        traceback.print_exc(limit=10)
        
    if err != self.kSuccess:
      self._DeconfigureInstance()
      
  def __CallTerminateCBs(self):
    
    terminate_cbs = self.__terminate_cbs

    self._terminating = False
    self.__terminate_cbs = set()
    self.__launch_time = None
    
    for cb in terminate_cbs:
      try:
        cb()
      except:
        if __debug__:
          raise
        print("Termination callback failed", cb)
        traceback.print_exc(limit=10)
    
    self.__valid = False
    
    self._UpdateStatus(self.kTerminated)
    
  def _AddPtyToCmdLine(self, cmd, pseudo_tty):
    if pseudo_tty:
      cmd.append('-t')
    else:
      cmd.append('-T')
        
  def _AddStartDirToCmdLine(self, cmd, startdir):
    if startdir:
      cmd.append('--cwd')
      cmd.append(startdir)
        
  def _AddEnvToCmdLine(self, cmd, env):
    combined_env = self.env.copy()
    combined_env.update(env)
    for key, value in combined_env.items():
      cmd.append('--env')
      cmd.append('{}={}'.format(key, value))
    
  def __DeconfigureAndLaunch(self):
    """Stop container if running, remove any stale configuration added previously
    in __ConfigureAndLaunch and then start reconfiguration and launch"""
    
    self._DeconfigureInstance(self.__ConfigureAndLaunch)
    
  def __ConfigureAndLaunch(self):
    """Set up environment variables, file mounts, port forwards, and other networking
    used by container, and then launch the container instance."""
    
    cmds = []
    
    # Commands to set up envs
    for key, value in self.env.items():
      cmds.append(['lxc', 'config', 'set', self.image_id, 'environment.{}={}'.format(key, value)])

    # Commands to set up file mounts
    for i, (local, container) in enumerate(self.mappings):
      cmds.append(['lxc', 'config', 'device', 'add', self.image_id, 'wingmt{}'.format(i),
                   'disk', 'source={}'.format(local), 'path={}'.format(container)])

    # XXX Hopefully not needed after all
    # Commands to create a bridge network and nic on the container
    # XXX Setting ipv6.address=none is required for ipv4 DHCP to work (likely an LXD bug?)
    #cmds.append(['lxc', 'network', 'create', '{}-br0'.format(self.image_id), 'ipv6.address=none'])
    #cmds.append(['lxc', 'config', 'device', 'add', self.image_id, 'wing-eth0', 'nic', 'nictype=bridged',
                 #'parent={}-br0'.format(self.image_id), 'name=wing-eth0'])
    
    # Commands to set up port mappings
    for i, (local, remote) in enumerate(self.port_forwards):
      cmds.append(['lxc', 'config', 'device', 'add', self.image_id, 'wingport{}'.format(i),
                   'proxy', 'listen=tcp:0.0.0.0:{}'.format(local), 'connect=tcp:127.0.0.1:{}'.format(remote)])

    # Add directories
    cmds = [(None, c) for c in cmds]      
    
    # Run all the commands in sequence and continue when all are successful  
    proc = spawn.CChildProcessSequence(cmds, env=project.GetEffectiveEnvironment())
      
    def starting(dirname, cmd):
      print(_("Starting command {} in {}").format(cmd, dirname))
      
    def terminated(dirname, cmd, errno, exc_text):
      if errno == 'start-failed':
        print(_("Process failed to start:"))
        print(exc_text)
        self._DeconfigureInstance()
      elif errno != 0:
        print(_("Process failed with errno={}").format(errno))
      
    def output(dirname, cmd, txt):
      print(txt)
      
    def completed(success):
      if success:
        print(_("LXC configuration succeeded"))
        self.__LaunchInstance()
      else:
        print(_("LXC configuration failed"))
        self.__CallLaunchCBs(self.kProcessFailedToStart)

    proc.connect_while_alive('starting', starting, self)
    proc.connect_while_alive('terminated', terminated, self)
    proc.connect_while_alive('received-output', output, self)
    proc.connect_while_alive('completed', completed, self)

    print(_("Configuring LXC instance {}").format(self.instance_id))
    proc.Start()
      
class CContainerSystem:
  
  gStatusCBs = []
  
  def GetContainerSystemInfo(self):

    return {
      'id': 'lxc',
      'name': 'LXC/LXD',
      'default-host-ip': None,
      'config-file-spec': (None, []),
    }
    
  def GetAllContainerInstances(self):
    """Return a list of all active container instances"""

    return [inst for inst in gInstances.values() if inst._valid]
  
  def Shutdown(self):

    for instance_id, instance in list(gInstances.items()):
      instance.Terminate()

  @classmethod
  def _GetContainerImageCreationTime(cls, image_id, result_cb):

    cmd = ['lxc', 'info', image_id]
    
    def terminated(p):
      
      if proc.fTimedOut:
        print(_("Timed out in 'lxc info {}'").format(image_id))
        result_cb(None)
        return

      output = proc.GetOutput()
      try:
        created = None
        for line in output.splitlines():
          line = line.strip()
          if line.lower().startswith('created:'):
            created = line[len('created:'):].strip()
            break
        if created is None:
          raise RuntimeError('notfound')
      except:
        print(_("Failed to parse output from 'lxc info {}'").format(image_id))
        print(output)
        result_cb(None)
      else:
        result_cb(created)
    
    def start_failed(p, exc_text):
      print(_("Failed to run 'lxc info {}'").format(image_id))
      print(exc_text)
      result_cb(None)
      
    proc = spawn.CChildProcess(cmd, env=config.gStartupEnv)
    proc.connect('terminated', terminated)
    proc.connect('start-failed', start_failed)

    proc.Start()

  def RegisterStatusCallback(self, cb):
    """Register a callback to invoke when the runtime status of containers change. The
    callback is invoked with (uid, status) where status is one of the container status
    class attributes for CContainerInstance: kRunning, kTerminated, kStarting,
    kStartFailed, kTerminating."""
    
    self.gStatusCBs.append(cb)
    
  def RemoveStatusCallback(self, cb):
    """Remove a callback previously registered with RegisterStatusCallback()."""

    self.gStatusCBs.remove(cb)
  
  def ParseConfiguration(self, config_path):

    # LXC does not use configuration files/dirs so this is never called
    return None

  def GetBuildCommandLines(self, config_path, image_id):

    # LXC does not really have a build process:  Images may be set up
    # using profiles, but then automatically update if the profile is
    # changed or other configuration options are set on the instance.
    return []
  
  @classmethod
  def _CallStatusCallbacks(cls, image_id, status):
    
    for cb in cls.gStatusCBs:
      cb(image_id, status)
      
  def FindAllContainerImages(self, result_cb):
  
    found_image_ids = set()
    
    cmd = ['lxc', 'list', '--columns', 'n', '--format', 'csv']
    
    proc = spawn.CChildProcess(cmd, env=config.gStartupEnv)

    def start_failed(pobj, exc_text):
      print("Failed to run 'lxc list'")
      print(exc_text)
      print(pobj.GetOutput())
      result_cb(None)
      
    proc.connect('start-failed', start_failed)

    def terminated(pobj):
      errno = pobj.GetExitCode()
      if errno != 0:
        start_failed(pobj, "Non-zero exit code: {}".format(errno))
        return
      
      output = pobj.GetOutput()
      lines = output.splitlines()
      for line in lines:
        line = line.strip()
        found_image_ids.add((line, line))
          
      result_cb(sorted(found_image_ids))
        
    proc.connect('terminated', terminated)
    
    proc.Start()

