import subprocess
import os
import re
import logging
from stimela import utils
import stimela
from shutil import which
from dataclasses import dataclass
from omegaconf import OmegaConf
from typing import Dict, List, Any, Optional, Callable
from contextlib import ExitStack
from scabha.basetypes import EmptyListDefault
import datetime
from stimela.utils.xrun_asyncio import xrun

from stimela.exceptions import BackendError

from . import native

@dataclass
class SingularityBackendOptions(object):
    enable: bool = True
    image_dir: str = os.path.expanduser("~/.singularity")
    auto_build: bool = True
    rebuild: bool = False
    executable: Optional[str] = None
    remote_only: bool = False      # if True, won't look for singularity on local system -- useful in combination with slurm wrapper

    # @dataclass
    # class EmptyVolume(object):
    #     name: str
    #     mount: str

    # tmp_dirs: List[EmptyVolume] = EmptyListDefault()

SingularityBackendSchema = OmegaConf.structured(SingularityBackendOptions)

STATUS = VERSION = BINARY = None

# images rebuilt in this run
_rebuilt_images = set()

def is_available(opts: Optional[SingularityBackendOptions] = None):
    global STATUS, VERSION, BINARY
    if STATUS is None:
        if opts and opts.remote_only:
            STATUS = VERSION = "remote"
        else:
            BINARY = (opts and opts.executable) or which("singularity")
            if BINARY:
                __version_string = subprocess.check_output([BINARY, "--version"]).decode("utf8")
                STATUS = VERSION = __version_string.strip().split()[-1]
                # if VERSION < "3.0.0":
                #     suffix = ".img"
                # else:
                #     suffix = ".sif"
            else:
                STATUS = "not installed"
                VERSION = None    
    return VERSION is not None

def get_status():
    is_available()
    return STATUS

def is_remote():
    return False

def init(backend: 'stimela.backend.StimelaBackendOptions', log: logging.Logger):
    pass

def get_image_info(cab: 'stimela.kitchen.cab.Cab', backend: 'stimela.backend.StimelaBackendOptions'):
    """returns image name/path corresponding to cab

    Args:
        cab (stimela.kitchen.cab.Cab): _description_
        backend (stimela.backend.StimelaBackendOptions): _description_

    Returns:
        name, path, enable_update: tuple of docker image name, path to singularity image on disk, and enable-updates flag
    """

    # prebuilt image
    if cab.image and cab.image.path:
        simg_path = cab.image.path
        if not os.path.exists(simg_path):
            raise BackendError(f"image {simg_path} for cab '{cab.name}' doesn't exist")
        return os.path.basename(simg_path), simg_path, False

    image_name = cab.flavour.get_image_name(cab, backend)

    if not image_name:
        raise BackendError(f"cab '{cab.name}' does not define an image")
    
    # convert to filename
    simg_name = image_name.replace("/", "-") + ".simg"
    simg_path = os.path.join(backend.singularity.image_dir, simg_name) 

    return image_name, simg_path


def build(cab: 'stimela.kitchen.cab.Cab', backend: 'stimela.backend.StimelaBackendOptions', log: logging.Logger,
          wrapper: Optional['stimela.backend.runner.BackendWrapper']=None,
          build=True, rebuild=False):
    """Builds image for cab, if necessary.

    build: if True, build missing images regardless of backend settings
    rebuild: if True, rebuild all images regardless of backend settings

    Returns:
        str: path to corresponding singularity image
    """

    image_name, simg_path = get_image_info(cab, backend)

    # this is True if we're allowed to build missing images
    build = build or rebuild or backend.singularity.auto_build   
    # this is True if we're asked to force-rebuild images
    rebuild = rebuild or backend.singularity.rebuild
    
    cached_image_exists = os.path.exists(simg_path)

    # no image? Better have builds enabled then
    if not cached_image_exists:
        log.info(f"singularity image {simg_path} does not exist")
        if not build:
            raise BackendError(f"no image, and singularity build options not enabled")
    # else we have an image
    # if rebuild is enabled, delete it
    elif rebuild:
        if simg_path in _rebuilt_images:
            log.info(f"singularity image {simg_path} was rebuilt earlier")
        else:
            log.info(f"singularity image {simg_path} exists but a rebuild was specified")
            os.unlink(simg_path)        
            cached_image_exists = False
    else:
        log.info(f"singularity image {simg_path} exists")

    ## OMS: taking this out for now, need some better auto-update logic, let's come back to it later
    ## Please retain the code for now
        
    # # else check if it need to be auto-updated
    # elif auto_update_allowed and backend.singularity.auto_update:
    #     if image_name in _auto_updated_images:
    #         log.info("image was used earlier in this run, not checking for auto-updates again")
    #     else:
    #         _auto_updated_images.add(image_name)
    #         # force check of docker binary
    #         docker.is_available()
    #         if docker.BINARY is None:
    #             log.warn("a docker runtime is required for auto-update of singularity images: forcing unconditional rebuild")
    #             build = True
    #         else:
    #             log.info("singularity auto-update: pulling and inspecting docker image")
    #             # pull image from hub
    #             retcode = xrun(docker.BINARY, ["pull", image_name], 
    #                         shell=False, log=log,
    #                             return_errcode=True, command_name="(docker pull)", 
    #                             log_command=True, 
    #                             log_result=True)
    #             if retcode != 0:
    #                 raise BackendError(f"docker pull failed with return code {retcode}") 
    #             if os.path.exists(simg_path):
    #                 # check timestamp
    #                 result = subprocess.run(
    #                         [docker.BINARY, "inspect", "-f", "{{ .Created }}", image_name],
    #                         capture_output=True)
    #                 if result.returncode != 0:
    #                     for line in result.stdout.split("\n"):
    #                         log.warn(f"docker inpect stdout: {line}")
    #                     for line in result.stderr.split("\n"):
    #                         log.error(f"docker inpect stderr: {line}")
    #                     raise BackendError(f"docker inspect failed with return code {result.returncode}")
    #                 timestamp = result.stdout.decode().strip()
    #                 log.info(f"docker inspect returns timestamp {timestamp}")
    #                 # parse timestamps like '2023-04-07T13:39:19.187572398Z'
    #                 # Pre-3.11 pythons don't do it natively so we mess around...
    #                 match = re.fullmatch("(.*)T([^.]*)(\.\d+)?Z?", timestamp)
    #                 if not match:
    #                     raise BackendError(f"docker inspect returned invalid timestamp '{timestamp}'")
    #                 try:
    #                     dt = datetime.datetime.fromisoformat(f'{match.group(1)} {match.group(2)} +00:00')
    #                 except ValueError as exc:
    #                     raise BackendError(f"docker inspect returned invalid timestamp '{timestamp}', exc")

    #                 if dt.timestamp() > os.path.getmtime(simg_path):
    #                     log.warn("docker image is newer than cached singularity image, rebuilding")
    #                     os.unlink(simg_path)        
    #                     cached_image_exists = False
    #                 else:
    #                     log.info("cached singularity image appears to be up-to-date")

    # if image doesn't exist, build it. We will have already checked for build settings
    # being enabled above
    if not cached_image_exists:
        log.info(f"(re)building image {simg_path}")

        args = [BINARY, "build", simg_path, f"docker://{image_name}"]

        if wrapper:
            args = wrapper.wrap_build_command(args, log=log)

        retcode = xrun(args[0], args[1:], shell=False, log=log,
                    return_errcode=True, command_name="(singularity build)", 
                    gentle_ctrl_c=True,
                    log_command=True, 
                    log_result=True)

        if retcode:
            raise BackendError(f"singularity build returns {retcode}")

        if not os.path.exists(simg_path):
            raise BackendError(f"singularity build did not return an error code, but the image did not appear")
        
        _rebuilt_images.add(simg_path)
        
    return simg_path



def run(cab: 'stimela.kitchen.cab.Cab', params: Dict[str, Any], fqname: str,
        backend: 'stimela.backend.StimelaBackendOptions',
        log: logging.Logger, subst: Optional[Dict[str, Any]] = None,
        wrapper: Optional['stimela.backends.runner.BackendWrapper'] = None):

    """Runs cab contents

    Args:
        cab (Cab): cab object
        log (logger): logger to use
        subst (Optional[Dict[str, Any]]): Substitution dict for commands etc., if any.

    Returns:
        Any: return value (e.g. exit code) of content
    """
    from .utils import resolve_required_mounts

    native.update_rlimits(backend.rlimits, log)

    # get path to image, rebuilding if backend options allow this
    simg_path = build(cab, backend=backend, log=log, build=False, wrapper=wrapper)

    # build up command line    
    cwd = os.getcwd()
    args = [backend.singularity.executable or BINARY, 
            "exec", 
            "--containall",
            "--pwd", cwd]
    
    # initial set of mounts has cwd as read-write
    mounts = {cwd: True}
    # get extra required filesystem bindings
    resolve_required_mounts(mounts, params, cab.inputs, cab.outputs)

    # sort mount paths before iterating -- this ensures that parent directories come first
    # (singularity doesn't like it if you specify a bind of a subdir before a bind of a parent) 
    for path, rw in sorted(mounts.items()):
        args += ["--bind", f"{path}:{path}:{'rw' if rw else 'ro'}"]

    args += [simg_path]
    args += cab.flavour.get_arguments(cab, params, subst, check_executable=False)
    log.debug(f"command line is {args}")

    cabstat = cab.reset_status()

    command_name = f"{cab.flavour.command_name}"

    # run command
    start_time = datetime.datetime.now()
    def elapsed(since=None):
        """Returns string representing elapsed time"""
        return str(datetime.datetime.now() - (since or start_time)).split('.', 1)[0]

    # log.info(f"argument lengths are {[len(a) for a in args]}")

    if wrapper:
        args = wrapper.wrap_run_command(args, fqname=fqname, log=log)

    retcode = xrun(args[0], args[1:], shell=False, log=log,
                output_wrangler=cabstat.apply_wranglers,
                return_errcode=True, command_name=command_name, 
                gentle_ctrl_c=True,
                log_command=True if cab.flavour.log_full_command else command_name, 
                log_result=False)

    # check if output marked it as a fail
    if cabstat.success is False:
        log.error(f"declaring '{command_name}' as failed based on its output")

    # if retcode != 0 and not explicitly marked as success, mark as failed
    if retcode and cabstat.success is not True:
        cabstat.declare_failure(f"{command_name} returns error code {retcode} after {elapsed()}")
    else:
        log.info(f"{command_name} returns exit code {retcode} after {elapsed()}")

    return cabstat


# class SingularityError(Exception):
#     pass

# def pull(image, name, docker=True, directory=".", force=False):
#     """ 
#         pull an image
#     """
#     if docker:
#         fp = "docker://{0:s}".format(image)
#     else:
#         fp = image
#     if not os.path.exists(directory):
#         os.mkdir(directory)

#     image_path = os.path.abspath(os.path.join(directory, name))
#     if os.path.exists(image_path) and not force:
#         stimela.logger().info(f"Singularity image already exists at '{image_path}'. To replace it, please re-run with the 'force' option")
#     else:
#         utils.xrun(f"cd {directory} && singularity", ["pull", 
#         	"--force" if force else "", "--name", 
#          	name, fp])

#     return 0

# class Container(object):
#     def __init__(self, image, name,
#                  volumes=None,
#                  logger=None,
#                  time_out=-1,
#                  runscript="/singularity",
#                  environs=None,
#                  workdir=None,
#                  execdir="."):
#         """
#         Python wrapper to singularity tools for managing containers.
#         """

#         self.image = image
#         self.volumes = volumes or []
#         self.environs = environs or []
#         self.logger = logger
#         self.status = None
#         self.WORKDIR = workdir
#         self.RUNSCRIPT = runscript
#         self.PID = os.getpid()
#         self.uptime = "00:00:00"
#         self.time_out = time_out
#         self.execdir = execdir

#         self._env = os.environ.copy()

#         hashname = hashlib.md5(name.encode('utf-8')).hexdigest()[:3]
#         self.name = hashname if version < "3.0.0" else name

#     def add_volume(self, host, container, perm="rw", noverify=False):
#         if os.path.exists(host) or noverify:
#             if self.logger:
#                 self.logger.debug("Mounting volume [{0}] in container [{1}] at [{2}]".format(
#                     host, self.name, container))
#             host = os.path.abspath(host)
#         else:
#             raise IOError(
#                 "Path {0} cannot be mounted on container: File doesn't exist".format(host))

#         self.volumes.append(":".join([host, container, perm]))

#         return 0

#     def add_environ(self, key, value):
#         self.logger.debug("Adding environ varaible [{0}={1}] "\
#                     "in container {2}".format(key, value, self.name))
#         self.environs.append("=".join([key, value]))
#         key_ = f"SINGULARITYENV_{key}"
	
#         self.logger.debug(f"Setting singularity environmental variable {key_}={value} on host")
#         self._env[key_] = value

#         return 0

#     def run(self, *args, output_wrangler=None):
#         """
#         Run a singularity container instance
#         """

#         if self.volumes:
#             volumes = " --bind " + " --bind ".join(self.volumes)
#         else:
#             volumes = ""

#         if not os.path.exists(self.image):
#             self.logger.error(f"The image, {self.image}, required to run this cab does not exist."\
#                     " Please run 'stimela pull --help' for help on how to download the image")
#             raise SystemExit from None

#         self.status = "running"
#         self._print("Starting container [{0:s}]. Timeout set to {1:d}. The container ID is printed below.".format(
#             self.name, self.time_out))
        
#         utils.xrun(f"cd {self.execdir} && singularity", ["run", "--workdir", self.execdir, "--containall"] \

#                     + list(args) + [volumes, self.image, self.RUNSCRIPT],
#                     log=self.logger, timeout=self.time_out, output_wrangler=output_wrangler,
#                     env=self._env, logfile=self.logfile)

#         self.status = "exited"

#         return 0

#     def _print(self, message):
#         if self.logger:
#             self.logger.info(message)
#         else:
#             print(message)

#         return 0
