# This file is based on conda_env source
# because we only read conda.yaml files generated by mlflow we can be more relaxed
# and we dont have to add a dependency on conda itself

import re
from collections import OrderedDict

import pkg_resources
import ruamel.yaml as yaml
from pip._vendor.packaging.utils import canonicalize_name, canonicalize_version


def yaml_safe_load(string):
    """
    Examples:
        >>> yaml_safe_load("key: value")
        {'key': 'value'}
    """
    # Version 1.2 is the same that conda supports
    return yaml.safe_load(string, version="1.2")


class Environment(object):
    def __init__(
        self,
        name=None,
        channels=None,
        dependencies=None,
    ):
        self.name = name
        self.dependencies = Dependencies(dependencies)

        if channels is None:
            channels = []
        self.channels = channels

    @classmethod
    def from_file(cls, fpath):
        with open(fpath, "r") as file:
            yaml_str = file.read()
            return cls.from_yamlstr(yaml_str)

    @classmethod
    def from_yamlstr(cls, yamlstr, **kwargs):
        """
        Load and return a ``Environment`` from a given ``yaml string``
        """
        data = yaml_safe_load(yamlstr)
        # data = validate_keys(data, kwargs)

        if kwargs is not None:
            for key, value in kwargs.items():
                data[key] = value

        return cls(**data)

    def to_dict(self):
        d = OrderedDict([("name", self.name)])
        if self.channels:
            d["channels"] = self.channels
        if self.dependencies:
            d["dependencies"] = self.dependencies
        return d

    def list_deps(self):
        """
        Returns a single list with the dependencies we need to deploy the model
        1. Python and pip
        """
        ret = []
        ignore_list = ("python", "pip")
        for dep in self.dependencies["conda"]:
            if dep.name not in ignore_list:
                ret.append(str(dep))

        for dep in self.dependencies["pip"]:
            ret.append(str(dep))

        return ret

    def __repr__(self):
        return str(self.to_dict())


class Dependencies(OrderedDict):
    def __init__(self, raw, *args, **kwargs):
        super(Dependencies, self).__init__(*args, **kwargs)
        self.raw = raw
        self.parse()

    def parse(self):
        if not self.raw:
            return

        self.update({"conda": [], "pip": []})

        for dep in self.raw:
            if isinstance(dep, str):
                # conda deps

                # 1. Convert the conda single = to pip double ==
                dep = re.sub(r"(?<!=)=(?!=)", "==", dep)

                self["conda"].append(MatchSpec.from_str(dep))
            elif isinstance(dep, dict):
                # pip dependencies
                for dep in dep["pip"]:
                    self["pip"].append(MatchSpec.from_str(dep))


class MatchSpec(object):
    def __init__(self, requirement):
        self.requirement = requirement
        self.name = requirement.name
        self.specs = sorted(requirement.specs)

    @classmethod
    def from_str(cls, string):
        requirement = pkg_resources.parse_requirements(string)
        requirement = list(requirement)[0]
        return cls(requirement=requirement)

    def __repr__(self):
        specs = []
        for op, ver in self.specs:
            specs.append(f"{op}{ver}")

        if len(specs) > 0:
            specs = ",".join(specs)
            return f"{self.name}{specs}"
        else:
            return f"{self.name}"

    def __str__(self):
        return self.__repr__()
