import os
import sys
import time
import traceback
import tempfile
import json

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

from plugins.container import api

import config

_ = api._

kInstanceLaunchTimeout = 10.0
kDockerCmd = 'docker'

gDockerInstances = {}
class CContainerInstance(destroyable.CDestroyable, api.CContainerInstance):
  """Support for launching a Docker 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.container_id = None
    self.__did_launch = False
    gDockerInstances[self.instance_id] = self
    
    self.__docker_tag = '{}_wing_{}'.format(self.instance_id, textutils.GeneratePW(pwlen=4, mix_case=False,
                                                                                   chars='abcdefghmnopqrstuvwxyz023456789'))
    self._valid = True
    self.__status = self.kTerminated
    self.__starting = False
    self.__retries_allowed = 2
    self.__launch_cbs = set()
    self.__image_time = None
    self.__launch_time = None
    self._terminating = False
    self.__terminate_cbs = set()
    
    self.pyexec = pyexec

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

    return not self.destroyed() and (self.container_id is not None or self._terminating)
  
  def GetStatus(self):
    
    return self.__status
  
  def _UpdateStatus(self, status):
    
    self.__status = status
    CContainerSystem._CallStatusCallbacks(self.instance_id, status)
    
  @classmethod
  def _DockerPS(cls, result_cb):
    """Run 'docker ps' and return a list of running image_ids"""
    
    cmd = [kDockerCmd, 'ps']
    
    def parse_docker_ps(txt):
      retval = []
      for line in txt.splitlines()[1:]:
        line = line.strip()
        image_id = line.split()[-1]
        retval.append(image_id)
      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 'docker ps'"))
        return
      try:
        retval = parse_docker_ps(proc.GetOutput())
      except:
        log_error(_("Failed to parse 'docker ps' output"))
        return
      
      result_cb(None, retval)
      
    def start_failed(p, exc_text):
      log_error(_("Failed to start 'docker ps'"))
      
    proc = CContainerSystem._CreateChildProcess(cmd)
    proc.connect('terminated', terminated)
    proc.connect('start-failed', start_failed)
  
    proc.Start()
    
  @classmethod
  def _DockerKillInstances(cls, image_id, complete_cb=None):
    """Stop any instances of given image, to get rid of any left over from previous
    runs.  In general there should be none but can be if Wing crashes or is terminated
    from the outside."""
    
    def ps_result(err, instances):
      assert not err
      for instance_id in instances:
        if instance_id.startswith('{}_wing_'.format(image_id)):
          print("Killing docker instance", instance_id)
          os.system('docker kill {}'.format(instance_id))
          
      if callable(complete_cb):
        cls.__WaitForInstanceExit(image_id, complete_cb)
        
    print("Killing docker instances for image_id", image_id)
    cls._DockerPS(ps_result)
    
  @classmethod
  def __WaitForInstanceExit(cls, image_id, complete_cb):
    
    end_time = time.time() + 3.0
    
    def ps_result(err, instances):
      if err:
        print("Failed to wait for instance exit - no docker ps results")
        complete_cb(False)
        return
      
      for instance_id in instances:
        if instance_id.startswith('{}_wing_'.format(image_id)):
          if time.time() > end_time:
            print("Waiting for instance exit timed out; instance_id", instance_id)
            complete_cb(False)
          else:
            def try_again():
              cls.__WaitForInstanceExit(image_id, complete_cb)
            timeout.timeout_add(100, try_again)
          return
        
      complete_cb(True)
      
    print("Waiting for instance exit; image id", image_id)
    cls._DockerPS(ps_result)
      
  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.__ContinueLaunch()
      CContainerSystem._GetContainerImageCreationTime(self.image_id, inspect_complete)
      return
    
    else:
      self.__ContinueLaunch()
      
  def __ContinueLaunch(self):
        
    self.cidfile = tempfile.mktemp()
    
    # Base command line:
    # We specify the instance name, cidfile, and also override ENTRYPOINT
    # in case use set it in Dockerfile so our following commands are run
    cmd = [kDockerCmd, 'run', '--name', self.__docker_tag, '--cidfile', self.cidfile, '--entrypoint', '/usr/bin/env']
    
    # XXX Hack to get networking to work on Linux (in conjunction with using hostname
    # XXX '127.0.0.1' instead of 'host.docker.internal') so the container can connect 
    # XXX to the IDE; This likely needs more work or a different approach., and it does
    # XXX not work when running under WSL2.
    if sys.platform.startswith('linux') and not os.environ.get('WSL_DISTRO_NAME'):
      cmd.extend(['--network=host'])      
    
    self._AddEnvToCmdLine(cmd, {})
    self._AddMappingsToCmdLine(cmd)
    self._AddPortForwardsToCmdLine(cmd)
    self._AddExtraRunArgsToCmdLine(cmd)
    cmd.append(self.image_id)
    
    cmd.extend([self.pyexec, '/wingpro{}/plugins/container/_docker_sleep.py'.format(config.kVersionTuple[0]), self.image_id])
    
    def terminated(p):
      if not self._terminating:

        self.io_cb('error', _("Docker instance {} terminated unexpectedly; err={}").format(self.__docker_tag,
                                                                                           proc.GetExitCode()))

        # Note that 'docker run' actually writes the instance cidfile file and looks like it's working
        # but then terminates on error later, so we may already have called the launch_cbs
        # and set __starting back to False.  If this has not yet happened, we retry launch up to
        # the allowed limit.  Otherwise it's too late and the launch_cbs from the first call
        # would have failed.  We try to reduce the chances of that by waiting briefly
        # before setting instance_id but that is imperfect and there's unavoidably a race condition.
        if self.__starting and self.__retries_allowed:
          self.io_cb('info', _("Docker container launch failed; trying to terminate stale instances"))
          self.__retries_allowed -= 1
          def retry():
            self._TryTerminatingStaleInstances()
          self._Remove(retry)
          return
        else:
          self.io_cb('error', _("Docker container launch failed - docker run terminated unexpectedly"))
          self.__CallLaunchCBs(self.kUnexpectedExit)
          
        self._UpdateStatus(self.kStartFailed)

      self.__CallTerminateCBs()      
      self._Remove()

    def start_failed(p, exc_text):
      self.__CallLaunchCBs(self.kProcessFailedToStart)

    def output(s, txt):        
      if 'CONTAINER-READY:{}'.format(self.image_id) in txt:
        self.__SetContainerIDFromCidFile()
        txt = txt.replace('CONTAINER-READY:{}'.format(self.image_id), _("Container {} is ready").format(self.image_id))
      
    proc = self._CreateChildProcess(cmd)
    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(retry_number=self.__retries_allowed):

      # Stop waiting if did launch successfully or it's a late call from a previous retry
      if self.__did_launch or retry_number != self.__retries_allowed:
        return False

      # Call launch cbs with error if timed out
      if time.time() > max_time:
        self.io_cb('error', _("Docker run 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 __SetContainerIDFromCidFile(self):
    """Read the container ID from the cidfile.  Sets the container_id attribute and
    calls launch cbs with appropriate status/error."""
    
    assert self.cidfile is not None
    
    # Container id file doesn't exist
    if not os.path.exists(self.cidfile):
      self.io_cb('error', _("Docker instance cidfile {} was not created").format(self.cidfile))
      self.__CallLaunchCBs(self.kProcessFailedToStart)
      return
      
    # Get the container ID from the file written by 'docker run'
    container_id = None
    try:
      with open(self.cidfile, 'rb') as f:
        container_id = f.read().strip().decode('utf-8')
    except:
      self.io_cb('error', _("Failed to read Docker instance cidfile {}").format(self.cidfile))
      self.io_cb('error', traceback.format_exc())
      self.__CallLaunchCBs(self.kUnexpectedExit)
      return

    if not container_id:
      self.__CallLaunchCBs(self.kProcessFailedToStart)
      return
    
    # Success
    os.unlink(self.cidfile)
    self.cidfile = None
    self.container_id = container_id
    self.__CallLaunchCBs(self.kSuccess)
    
  def _TryTerminatingStaleInstances(self):
    """Try terminating stale instances for this container and then attempt Launch again"""
    
    def retry_launch(success):
      if not success:
        self.io_cb('warn', ("Failed to terminate stale instances for image id {}").format(self.image_id))
        self.__CallLaunchCBs(self.kProcessFailedToStart)
      else:
        self.io_cb('info', _("Retrying Docker launch"))
        self.__ContinueLaunch()
        
    self.io_cb('info', _("Terminating stale instances for image_id").format(self.image_id))
    self._DockerKillInstances(self.image_id, complete_cb=retry_launch)
    
  def Terminate(self, completed_cb=None):
    
    if self.destroyed() or self.container_id is None:
      return
    
    if completed_cb is not None:
      self.__terminate_cbs.add(completed_cb)

    if self._terminating:
      return
    
    self._UpdateStatus(self.kTerminating)
    
    cmd = [kDockerCmd, 'kill', self.__docker_tag]

    def terminated(p):
      def remove_complete():
        self._valid = False
        self.container_id = None
        self.destroy()
      self._Remove(remove_complete)
      
    def start_failed(p, exc_text):
      self._terminating = False
      self.__terminate_cbs.remove(completed_cb)
      
    proc = self._CreateChildProcess(cmd)
    proc.connect_while_alive('terminated', terminated, self)
    proc.connect_while_alive('start-failed', start_failed, self)

    self._terminating = True
      
    self.io_cb('info', _("Running: {}").format(spawn._win32_cmd_line_from_argv(cmd)))
    proc.Start()
    
  def _Remove(self, remove_complete_cb=None):

    print("Removing container for instance_id {}".format(self.__docker_tag))

    if remove_complete_cb is None:
      def remove_complete_cb():
        print("Container for instance_id {} removed".format(self.__docker_tag))
      
    cmd = [kDockerCmd, 'rm', self.__docker_tag]

    def terminated(p):
      remove_complete_cb()
      
    def start_failed(p, exc_text):
      remove_complete_cb()
            
    proc = self._CreateChildProcess(cmd)
    proc.connect_while_alive('terminated', terminated, self)
    proc.connect_while_alive('start-failed', start_failed, self)

    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=None, pseudo_tty=False):

    # If this is raised, caller needs to start the container first and wait for that
    assert self._valid and self.container_id is not None
    
    cmd = [kDockerCmd, 'exec']

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

    return None, cmd

  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._Remove()
      
  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.container_id = None
    
    self._UpdateStatus(self.kTerminated)
    
  def _AddPtyToCmdLine(self, cmd, pseudo_tty):
    if pseudo_tty:
      cmd.append('-it')
        
  def _AddStartDirToCmdLine(self, cmd, startdir):
    if startdir:
      cmd.append('-w')
      cmd.append(startdir)
        
  def _AddEnvToCmdLine(self, cmd, env):
    combined_env = self.env.copy()
    if env:
      combined_env.update(env)
      
    for key, value in combined_env.items():
      cmd.append('-e')
      cmd.append('{}={}'.format(key, value))
    
  def _AddMappingsToCmdLine(self, cmd):
    for local, container in self.mappings:
      cmd.append('-v')
      cmd.append('{}:{}'.format(local, container))
  
  def _AddPortForwardsToCmdLine(self, cmd):
    for local, remote in self.port_forwards:
      cmd.append('-p')
      cmd.append('{}:{}'.format(local, remote))

  def _AddExtraRunArgsToCmdLine(self, cmd):
    if self.extra_run_args:
      cmd.extend(spawn.ParseCmdArgs(self.extra_run_args))

  def _CreateChildProcess(self, cmdv, *, env=None, use_std_handlers=True):
    
    proc = CContainerSystem._CreateChildProcess(cmdv, env=env)
    
    cmd_line_text = spawn._win32_cmd_line_from_argv(cmdv)
    
    if use_std_handlers:
      def terminated(p):
        self.io_cb('info', _("Process terminated: {}").format(cmd_line_text))
        
      def start_failed(p, exc_text):
        self.io_cb('error', _("Failed to run '{}'").format(cmd_line_text))
        self.io_cb('error', exc_text)
        
      def output(p, txt):
        self.io_cb('output', txt)

      proc.connect_while_alive('terminated', terminated, self)
      proc.connect_while_alive('start-failed', start_failed, self)
      proc.connect_while_alive('received-output', output, self)
      
    return proc
      
class CContainerSystem:
  
  gStatusCBs = []
  
  @classmethod
  def _CreateChildProcess(self, cmdv, env=None):
    
    def start_failed(p, exc_text):
      print(_("Failed to run '{}'").format(spawn._win32_cmd_line_from_argv(cmdv)))
      print(exc_text)
      
    if env is None:
      env = config.gStartupEnv
    proc = spawn.CChildProcess(cmdv, env=env)
    proc.connect('start-failed', start_failed)
    return proc
    
  
  def GetContainerSystemInfo(self):

    # XXX Hack to get networking working on Linux in conjunction with using 'host' style
    # XXX networking - this probably needs more work or a different approach, and it
    # XXX isn't needed when running under WSL2.
    if sys.platform.startswith('linux') and not os.environ.get('WSL_DISTRO_NAME'):
      connect_hostname = '127.0.0.1'
    else:
      connect_hostname = 'host.docker.internal'
      
    return {
      'id': 'docker',
      'name': 'Docker',
      'default-host-ip': connect_hostname,
      'config-file-spec': ('file', ['*.*']),
    }
    
  def GetAllContainerInstances(self):
    """Return a list of all active container instances"""

    return [inst for inst in gDockerInstances.values() if inst._valid and not inst._terminating]
  
  def Shutdown(self):

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

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

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

      output = proc.GetOutput()
      try:
        jdata = json.loads(output)
      except:
        print(_("Failed to run 'docker image inspect {}'").format(image_id))
        print(output)
        result_cb(None)
      else:
        result_cb(jdata[0]['Created'])
    
    def start_failed(p, exc_text):
      result_cb(None)
      
    proc = cls._CreateChildProcess(cmd)
    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):
    
    try:
      with open(config_path, 'r') as f:
        lines = f.readlines()
        f.close()
    except Exception:
      print("Failed to read Docker configuration", config_path)
      return None
    
    dirname = os.path.dirname(config_path)
    
    retval = []
    container_dir = [None]
    curr_item = {}
    
    def expand_mappings():
      
      curr_item['image'] = curr_item.get('image', None)
      curr_item['source_image'] = curr_item.get('source_image', '')
      curr_item['ports'] = curr_item.get('ports', [])
      curr_item['env'] = curr_item.get('env', [])      
      curr_item['mappings'] = curr_item.get('mappings', [])

      if container_dir[0] is not None:
        cvt_mappings = []
        for hfile, cfile in curr_item['mappings']:
          hfile = os.path.join(dirname, hfile)
          if not cfile.startswith('/'):
            cfile = container_dir[0] + '/' + cfile
          cvt_mappings.append((hfile, cfile))
        curr_item['mappings'] = cvt_mappings
        
    def append_item():
      
      if curr_item:
        expand_mappings()      
        retval.append(curr_item.copy())
      
      container_dir[0] = None
      curr_item.clear()
    
    try:
      for line in lines:
        parts = textutils.QuotedSplit(line, allow_esc=True)
        if not parts:
          continue
        cmd = parts[0].lower()
        parts = parts[1:]
      
        if cmd == 'from' and parts:
  
          append_item()
            
          if parts[0].startswith('--platform'):
            parts = parts[1:]
            
          source_image = parts[0]
          if '@' in source_image:
            source_image = source_image.split('@')[0]
          if len(parts) == 3 and parts[1].lower() == 'as':
            image = parts[2]
          else:
            image = None
            
          curr_item['image'] = image
          curr_item['source_image'] = source_image
            
        elif cmd == 'copy' and parts:
          if parts[0].startswith('--chown'):
            parts = parts[1:]
          if len(parts) == 2:
            hfile, cfile = parts
            curr_item['mappings'] = curr_item.get('mappings', [])
            curr_item['mappings'].append((hfile, cfile))
            
        elif cmd == 'workdir' and parts:
          container_dir[0] = parts[0]
          if not container_dir[0].endswith('/'):
            container_dir[0] += '/'
            
        elif cmd == 'expose':
          if len(parts) == 1:
            curr_item['ports'] = curr_item.get('ports', [])
            curr_item['ports'].append((parts[0], parts[0]))
          elif len(parts) == 2:
            curr_item['ports'].append(tuple(parts))
            
        elif cmd == 'env' and parts:
          curr_item['env'] = curr_item.get('env', [])
          for pair in parts:
            key, value = pair.split('=')
            if value[0] == '"' and value[-1] == value[0]:
              value = value[1:-1]
            curr_item['env'].append((key, value))
        
      append_item()
      
    except:
      print("Failed to parse configuration", config_path)
      traceback.print_exc()
      return None

    return retval
        
  def GetBuildCommandLines(self, config_path, image_id):

    parse = self.ParseConfiguration(config_path)
    if parse is None:
      return None
    
    images = []
    for info in parse:
      if info['image']:
        images.append(info['image'])
        
    if image_id in images:
      if len(parse) > 1:
        target_args = ['--target', image_id]
      else:
        target_args = []
    else:
      target_args = ['--tag={}'.format(image_id)]
      
    cmd = [kDockerCmd, 'build', '-f', config_path] + target_args + ['.']
    
    return [(os.path.dirname(config_path), cmd)]
    
  @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()
    
    proc = self._CreateChildProcess([kDockerCmd, 'images'])
    
    def start_failed(pobj, exc_text):
      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()[1:]
      for line in lines:
        name = None
        tag = None
        image_id = None
        for word_count, word in enumerate(line.split()):
          if word_count == 0:
            name = word
          elif word_count == 1:
            tag = word
          elif word_count == 2:
            image_id = word
        if name is None or tag is None or image_id is None:
          continue
        if name.startswith('<') or name.startswith('tmp') or name.startswith('k8s.'):
          continue
        found_image_ids.add((name + ':' + tag, name + ':' + tag))
          
      result_cb(sorted(found_image_ids))
        
    proc.connect('terminated', terminated)
    
    proc.Start()

