# Copyright 2017-2023 Posit Software, PBC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import base64
import collections
import hashlib
import logging
import os
import sys

import pkg_resources

from guild import config
from guild import entry_point_util
from guild import guildfile
from guild import namespace

log = logging.getLogger("guild")

_models = entry_point_util.EntryPointResources("guild.models", "model")

ModelRef = collections.namedtuple(
    "ModelRef", ["dist_type", "dist_name", "dist_version", "model_name"]
)


class Model:
    def __init__(self, ep):
        self.name = self._ep_model_name(ep)
        self.dist = ep.dist
        self.modeldef = self._init_modeldef()
        self._fullname = None  # lazy
        self._reference = None  # lazy

    @staticmethod
    def _ep_model_name(ep):
        if ep.name == "__anonymous__":
            return ""
        return ep.name

    def __repr__(self):
        class_module = self.__class__.__module__
        class_name = self.__class__.__name__
        return f"<{class_module}.{class_name} '{self.name}'>"

    @property
    def fullname(self):
        if self._fullname is None:
            pkg_name = namespace.apply_namespace(self.dist.project_name)
            if pkg_name and pkg_name != ".":
                self._fullname = f"{pkg_name}/{self.name}"
            else:
                self._fullname = self.name
        return self._fullname

    @property
    def reference(self):
        if self._reference is None:
            self._reference = self._init_reference()
        return self._reference

    def _init_modeldef(self):
        raise NotImplementedError()

    def _init_reference(self):
        raise NotImplementedError()


class GuildfileModel(Model):
    """A model associated with a guildfile.

    These are generated from GuildfileDistribution instances
    (i.e. distributions that are derrived from modefiles).
    """
    def _init_modeldef(self):
        assert isinstance(self.dist, GuildfileDistribution), self.dist
        return self.dist.get_modeldef(self.name)

    def _init_reference(self):
        src = self.dist.guildfile.src
        if src and os.path.isfile(src):
            version = file_hash(src)
        else:
            version = "unknown"
        src = self.dist.guildfile.src
        if src is not None:
            src = os.path.abspath(src)
        return ModelRef("guildfile", src, version, self.name)


class PackageModel(Model):
    """A model associated with a package.

    These are generated by Guild packages.
    """
    def _init_modeldef(self):
        modeldef = _find_dist_modeldef(self.name, self.dist)
        if modeldef is None:
            raise ValueError(f"undefined model '{self.name}' in {self.dist}")
        return modeldef

    def _init_reference(self):
        pkg_name = namespace.apply_namespace(self.dist.project_name)
        return ModelRef("package", pkg_name, self.dist.version, self.name)


def _find_dist_modeldef(name, dist):
    for modeldef in _ensure_dist_modeldefs(dist):
        if modeldef.name == name:
            return modeldef
    return None


def _ensure_dist_modeldefs(dist):
    if not hasattr(dist, "_modelefs"):
        dist._modeldefs = _load_dist_modeldefs(dist)
    return dist._modeldefs


def _load_dist_modeldefs(dist):
    modeldefs = []
    try:
        record = dist.get_metadata_lines("RECORD")
    except IOError:
        log.warning(
            "distribution %s missing RECORD metadata - unable to find models", dist
        )
    else:
        for line in record:
            path = line.split(",", 1)[0]
            if os.path.basename(path) == guildfile.NAME:
                fullpath = os.path.join(dist.location, path)
                _try_acc_modeldefs(fullpath, modeldefs)
    return modeldefs


def _try_acc_modeldefs(path, acc):
    try:
        models = guildfile.for_file(path)
    except Exception as e:
        log.error("unable to load models from %s: %s", path, e)
    else:
        for modeldef in models.models.values():
            acc.append(modeldef)


class GuildfileDistribution(pkg_resources.Distribution):
    def __init__(self, guildfile):
        super().__init__(guildfile.dir, project_name=self._init_project_name(guildfile))
        self.guildfile = guildfile
        self._entry_map = self._init_entry_map()

    def __repr__(self):
        return f"<guild.model.GuildfileDistribution '{self.guildfile.dir}'>"

    def get_entry_map(self, group=None):
        if group is None:
            return self._entry_map
        return self._entry_map.get(group, {})

    def get_modeldef(self, name):
        for modeldef_name, modeldef in self.guildfile.models.items():
            if modeldef_name == name:
                return modeldef
        raise ValueError(name)

    @staticmethod
    def _init_project_name(guildfile):
        """Returns a project name for a guildfile distribution.

        Guildfile distribution project names are of the format:

            '.guildfile.' + ESCAPED_GUILDFILE_PATH

        ESCAPED_GUILDFILE_PATH is a 'safe' project name (i.e. will not be
        modified in a call to `pkg_resources.safe_name`) that, when
        unescaped using `unescape_project_name`, is the relative path of
        the directory containing the guildfile. The modefile name itself
        (e.g. 'guild.yml') is not contained in the path.

        Guildfile paths are relative to the current working directory
        (i.e. the value of os.getcwd() at the time they are generated) and
        always start with '.'.
        """
        pkg_path = os.path.relpath(guildfile.dir, config.cwd())
        if pkg_path[0] != ".":
            pkg_path = os.path.join(".", pkg_path)
        safe_path = escape_project_name(pkg_path)
        return f".guildfile.{safe_path}"

    def _init_entry_map(self):
        return {
            "guild.models": {
                name: self._model_entry_point(model)
                for name, model in self.guildfile.models.items()
            },
            "guild.resources": {
                res.fullname: self._resource_entry_point(res.fullname)
                for res in self._guildfile_resources()
            },
        }

    def _model_entry_point(self, model):
        return pkg_resources.EntryPoint(
            name=model.name,
            module_name="guild.model",
            attrs=("GuildfileModel",),
            dist=self,
        )

    def _guildfile_resources(self):
        for modeldef in self.guildfile.models.values():
            for res in modeldef.resources:
                yield res

    def _resource_entry_point(self, name):
        return pkg_resources.EntryPoint(
            name=name,
            module_name="guild.model",
            attrs=("GuildfileResource",),
            dist=self,
        )


def escape_project_name(name):
    """Escapes name for use as a valie pkg_resources project name."""
    return str(base64.b16encode(name.encode("utf-8")).decode("utf-8"))


def unescape_project_name(escaped_name):
    """Unescapes names escaped with `escape_project_name`."""
    return str(base64.b16decode(escaped_name).decode("utf-8"))


class ModelImportError(ImportError):
    pass


class BadGuildfileDistribution(pkg_resources.Distribution):
    """Distribution for a guildfile that can't be read."""
    def __repr__(self):
        return f"<guild.model.BadGuildfileDistribution '{self.location}'>"

    def get_entry_map(self, group=None):
        return {}


class ModelImporter:
    undef = object()

    def __init__(self, path):
        if not _is_guildfile_dir(path):
            raise ModelImportError(path)
        self.path = path
        self._dist = self.undef  # lazy

    @property
    def dist(self):
        if self._dist is self.undef:
            self._dist = _init_guildfile_dist(self.path)
        return self._dist

    @staticmethod
    def find_module(_fullname, _path=None):
        return None


def _init_guildfile_dist(path):
    if not os.path.isdir(path):
        return None
    try:
        gf = guildfile.for_dir(path)
    except guildfile.NoModels:
        return None
    except Exception as e:
        if log.getEffectiveLevel() <= logging.DEBUG:
            log.exception(path)
        log.error("error loading guildfile from %s: %s", path, e)
        return BadGuildfileDistribution(path)
    else:
        return GuildfileDistribution(gf)


def _is_guildfile_dir(path):
    return guildfile.is_guildfile_dir(path) or _is_cwd(path)


def _is_cwd(path):
    return os.path.abspath(path) == os.path.abspath(config.cwd())


def _model_finder(importer, path, _only=False):
    assert isinstance(importer, ModelImporter)
    assert importer.path == path, (importer.path, path)
    if importer.dist:
        yield importer.dist


class GuildfileNamespace(namespace.PrefixNamespace):
    prefix = ".guildfile."

    @staticmethod
    def pip_info(*_):
        raise TypeError("guildfiles cannot be installed using pip")

    def package_name(self, project_name):
        pkg = super().package_name(project_name)
        parts = pkg.split("/", 1)
        decoded_project_name = unescape_project_name(parts[0])
        rest = "/" + parts[1] if len(parts) == 2 else ""
        return decoded_project_name + rest


def file_hash(path):
    try:
        path_bytes = open(path, "rb").read()
    except IOError:
        log.warning("unable to read %s to calculate guildfile hash", path)
        return "unknown"
    else:
        return hashlib.md5(path_bytes).hexdigest()


def script_model_ref(model_name, script_base):
    return ModelRef("script", os.path.abspath(script_base), "", model_name)


def get_path():
    return _models.path()


def set_path(path, clear_cache=False):
    _models.set_path(path, clear_cache)


class SetPath:
    def __init__(self, path, clear_cache=False):
        self._path = path
        self._clear_cache = clear_cache
        self._save = None

    def __enter__(self):
        assert self._save is None, self._save
        self._save = get_path()
        set_path(self._path, clear_cache=self._clear_cache)

    def __exit__(self, *_exc):
        assert self._save is not None
        set_path(self._save)
        self._save = None


def insert_path(item, clear_cache=False):
    path = _models.path()
    try:
        path.remove(item)
    except ValueError:
        pass
    path.insert(0, item)
    _models.set_path(path, clear_cache)


def iter_models():
    for _name, model in _models:
        yield model


def for_name(name):
    return _models.for_name(name)


def _register_model_finder():
    sys.path_hooks.insert(0, ModelImporter)
    pkg_resources.register_finder(ModelImporter, _model_finder)


# TODO: This needs to be removed as we don't need to customize module
# loading. The original intention was use Python's module import
# mechanism (PEP 302) as a facility to load Guild models. This is tied
# to Guild's original vision as a software packaging system, where
# models behave like distributions. It's not clear there's any benefit
# to this. Note that removing this will trigger a consider refactoring
# effort throughout the code base.

if os.getenv("NO_GUILD_MODEL_FINDER") != "1":
    _register_model_finder()
