from typing import Text

from furl import furl
from pathlib2 import Path

from clearml_agent.config import Config
from .pip_api.system import SystemPip


class RequirementsTranslator(object):

    """
    Translate explicit URLs to local URLs after downloading them to cache
    """

    SUPPORTED_SCHEMES = ["http", "https", "ftp"]

    def __init__(self, session, interpreter=None, cache_dir=None):
        self._session = session
        config = session.config
        self.cache_dir = cache_dir or Path(config["agent.pip_download_cache.path"]).expanduser().as_posix()
        self.enabled = config["agent.pip_download_cache.enabled"]
        Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
        self.config = Config()
        self.pip = SystemPip(interpreter=interpreter, session=self._session)
        self._translate_back = {}

    def download(self, url):
        self.pip.download_package(url, cache_dir=self.cache_dir)

    @classmethod
    def is_supported_link(cls, line):
        # type: (Text) -> bool
        """
        Return whether requirement is a link that should be downloaded to cache
        """
        url = furl(line)
        return (
            url.scheme
            and url.scheme.lower() in cls.SUPPORTED_SCHEMES
            and line.lstrip().lower().startswith(url.scheme.lower())
        )

    def translate(self, line):
        """
        If requirement is supported, download it to cache and return the download path
        """
        if not (self.enabled and self.is_supported_link(line)):
            return line
        command = self.config.command
        command.log('Downloading "{}" to pip cache'.format(line))
        url = furl(line)
        try:
            wheel_name = url.path.segments[-1]
        except IndexError:
            command.error('Could not parse wheel name of "{}"'.format(line))
            return line
        try:
            self.download(line)
            downloaded = Path(self.cache_dir, wheel_name).expanduser().as_uri()
        except Exception:
            command.error('Could not download wheel name of "{}"'.format(line))
            return line

        self._translate_back[str(downloaded)] = line

        return downloaded

    def replace_back(self, requirements):
        if not requirements:
            return requirements

        for k in requirements:
            # k is either pip/conda
            if k not in ('pip', 'conda'):
                continue

            original_requirements = requirements[k]
            new_requirements = []
            for line in original_requirements:
                local_file = [d for d in self._translate_back.keys() if d in line]
                if local_file:
                    local_file = local_file[0]
                    new_requirements.append(line.replace(local_file, self._translate_back[local_file]))
                else:
                    new_requirements.append(line)

            requirements[k] = new_requirements

        return requirements

from typing import Text

from furl import furl
from pathlib2 import Path

from clearml_agent.config import Config
from .pip_api.system import SystemPip


class RequirementsTranslator(object):

    """
    Translate explicit URLs to local URLs after downloading them to cache
    """

    SUPPORTED_SCHEMES = ["http", "https", "ftp"]

    def __init__(self, session, interpreter=None, cache_dir=None):
        self._session = session
        config = session.config
        self.cache_dir = cache_dir or Path(config["agent.pip_download_cache.path"]).expanduser().as_posix()
        self.enabled = config["agent.pip_download_cache.enabled"]
        Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
        self.config = Config()
        self.pip = SystemPip(interpreter=interpreter, session=self._session)
        self._translate_back = {}

    def download(self, url):
        self.pip.download_package(url, cache_dir=self.cache_dir)

    @classmethod
    def is_supported_link(cls, line):
        # type: (Text) -> bool
        """
        Return whether requirement is a link that should be downloaded to cache
        """
        url = furl(line)
        return (
            url.scheme
            and url.scheme.lower() in cls.SUPPORTED_SCHEMES
            and line.lstrip().lower().startswith(url.scheme.lower())
        )

    def translate(self, line):
        """
        If requirement is supported, download it to cache and return the download path
        """
        if not (self.enabled and self.is_supported_link(line)):
            return line
        command = self.config.command
        command.log('Downloading "{}" to pip cache'.format(line))
        url = furl(line)
        try:
            wheel_name = url.path.segments[-1]
        except IndexError:
            command.error('Could not parse wheel name of "{}"'.format(line))
            return line
        try:
            self.download(line)
            downloaded = Path(self.cache_dir, wheel_name).expanduser().as_uri()
        except Exception:
            command.error('Could not download wheel name of "{}"'.format(line))
            return line

        self._translate_back[str(downloaded)] = line

        return downloaded

    def replace_back(self, requirements):
        if not requirements:
            return requirements

        for k in requirements:
            # k is either pip/conda
            if k not in ('pip', 'conda'):
                continue

            original_requirements = requirements[k]
            new_requirements = []
            for line in original_requirements:
                local_file = [d for d in self._translate_back.keys() if d in line]
                if local_file:
                    local_file = local_file[0]
                    new_requirements.append(line.replace(local_file, self._translate_back[local_file]))
                else:
                    new_requirements.append(line)

            requirements[k] = new_requirements

        return requirements

Translate explicit URLs to local URLs after downloading them to cache # type: (Text) -> bool Return whether requirement is a link that should be downloaded to cache If requirement is supported, download it to cache and return the download path # k is either pip/conda