import collections
import os
import sys
from hashlib import sha1
import itertools
import pkgutil
import platform
import tempfile
import shutil
try:
    from urlparse import urlparse
except:
    from urllib.parse import urlparse


from metaflow.decorators import StepDecorator
from metaflow.environment import InvalidEnvironmentException
from metaflow.metadata import MetaDatum
from metaflow.metaflow_config import get_pinned_conda_libs, CONDA_PACKAGE_S3ROOT
from metaflow.util import get_metaflow_root
from metaflow.datatools import S3

from . import read_conda_manifest, write_to_conda_manifest
from .conda import Conda


class CondaStepDecorator(StepDecorator):
    name = 'conda'
    defaults = {'libraries': {},
                'python': None,
                'disabled': None}

    conda = None
    environments = None

    def _get_base_attributes(self):
        if 'conda_base' in self.flow._flow_decorators:
            return self.flow._flow_decorators['conda_base'].attributes
        return self.defaults

    def _python_version(self):
        return next(x for x in [
                    self.attributes['python'], 
                    self.base_attributes['python'], 
                    platform.python_version()] if x is not None)

    def is_enabled(self):
        return not next(x for x in [
                        self.attributes['disabled'],
                        self.base_attributes['disabled'],
                        False] if x is not None)

    def _lib_deps(self):
        deps = get_pinned_conda_libs()
        base_deps = self.base_attributes['libraries']
        deps.update(base_deps)
        step_deps = self.attributes['libraries']
        if isinstance(step_deps, collections.Mapping):
            deps.update(step_deps)
        return deps

    def _step_deps(self):
        deps = [b'python==%s' % self._python_version().encode()]
        deps.extend(b'%s==%s' % (name.encode('ascii'), ver.encode('ascii'))
                    for name, ver in self._lib_deps().items())
        return deps

    def _env_id(self):
        deps = self._step_deps()
        return 'metaflow_%s_%s_%s' % (self.flow.name, 
                                        self.architecture, 
                                        sha1(b' '.join(sorted(deps))).hexdigest())

    def _resolve_step_environment(self, force=False):
        env_id = self._env_id()
        cached_deps = read_conda_manifest()
        if CondaStepDecorator.conda is None:
            CondaStepDecorator.conda = Conda()
        if force or env_id not in cached_deps or 'cache_urls' not in cached_deps[env_id]:
            if force or env_id not in cached_deps:
                deps = self._step_deps()
                (exact_deps, urls) = \
                    self.conda.create(self.step, env_id, deps, architecture=self.architecture)
                payload = {
                    'explicit': exact_deps,
                    'deps': [d.decode('ascii') for d in deps],
                    'urls': urls,
                }
            else:
                payload = cached_deps[env_id]
            if self.datastore.TYPE == 's3' and 'cache_urls' not in payload:
                payload['cache_urls'] = self._cache_env()
            write_to_conda_manifest(env_id, payload)
        CondaStepDecorator.environments = CondaStepDecorator.conda.environments(self.flow.name)
        return env_id

    def _cache_env(self):
        env_id = self._env_id()
        files = []
        for package_info in self.conda.package_info(env_id):
            url = urlparse(package_info['url'])
            path = os.path.join(CONDA_PACKAGE_S3ROOT, 
                                url.netloc,
                                url.path.lstrip('/'),
                                package_info['md5'],
                                package_info['fn'])
            #The tarball maybe missing when user invokes `conda clean`!
            files.append((path, package_info['package_tarball_full_path']))
        with S3() as s3:
            s3.put_files(files, overwrite=False)
        return [files[0] for files in files]

    def _prepare_step_environment(self, step_name):
        env_id = self._resolve_step_environment()
        if env_id not in CondaStepDecorator.environments:
            cached_deps = read_conda_manifest()
            self.conda.create(self.step,
                              env_id,
                              cached_deps[env_id]['urls'],
                              architecture=self.architecture,
                              explicit=True)
            CondaStepDecorator.environments = CondaStepDecorator.conda.environments(self.flow.name)
        return env_id

    def _architecture(self, decos):
        for deco in decos:
            if deco.name == 'batch':
                # force conda resolution for linux-64 architectures
                return 'linux-64'
        bit = '32'
        if platform.machine().endswith('64'):
            bit = '64'
        if platform.system() == 'Linux':
            return 'linux-%s' % bit
        elif platform.system() == 'Darwin':
            return 'osx-%s' % bit
        else:
            raise InvalidEnvironmentException("@conda decorator not supported "
                                              "outside of Linux and Darwin platforms")

    def runtime_init(self, flow, graph, package, run_id):
        # Create a symlink to installed version of metaflow to execute user code against
        path_to_metaflow = os.path.join(get_metaflow_root(), "metaflow")
        self.metaflow_home = tempfile.mkdtemp(dir="/tmp")
        os.symlink(path_to_metaflow, os.path.join(self.metaflow_home, "metaflow"))

    def step_init(self, flow, graph, step, decos, environment, datastore, logger):
        if environment.TYPE != 'conda':
            raise InvalidEnvironmentException("@conda decorator requires "
                                              "--environment=conda")
        self.architecture = self._architecture(decos)
        self.step = step
        self.flow = flow
        self.datastore = datastore
        self.base_attributes = self._get_base_attributes()

    def package_init(self, flow, step, environment):
        if self.is_enabled():
            self._prepare_step_environment(step)

    def runtime_task_created(self, datastore, task_id, split_index, input_paths, is_cloned):
        if self.is_enabled():
            self.env_id = self._prepare_step_environment(self.step)

    def task_pre_step(
            self, step_name, ds, meta, run_id, task_id, flow, graph, retry_count, max_retries):
        meta.register_metadata(run_id, step_name, task_id,
                                   [MetaDatum(field='conda_env_id',
                                              value=self._env_id(),
                                              type='conda_env_id')])

    def runtime_step_cli(self, cli_args, retry_count, max_user_code_retries):
        if self.is_enabled() and 'batch' not in cli_args.commands:
            python_path = self.metaflow_home
            if os.environ.get('PYTHONPATH') is not None:
                python_path = os.pathsep.join([os.environ['PYTHONPATH'], python_path])
            cli_args.env['PYTHONPATH'] = python_path
            cli_args.env['_METAFLOW_CONDA_ENV'] = self.env_id
            cli_args.entrypoint[0] = self.conda.python(self.env_id)

    def runtime_finished(self, exception):
        shutil.rmtree(self.metaflow_home)