"""This module contains tools for describing objects being published to DLHub."""
import importlib
from typing import List, Sequence, Optional, Iterable, Union, Dict
import pkg_resources
import requests
from enum import Enum

from pydantic import BaseModel, Field
from six import string_types
from zipfile import ZipFile
from glob import glob
import fnmatch
import json
import sys
import os
import re

from dlhub_sdk.models.datacite import Datacite, DataciteRelatedIdentifierType, DataciteRelationType
from dlhub_sdk.version import __version__

name_re = re.compile(r'^\S+$')


class DLHubType(Enum):
    """Type supported by DLHub"""

    servable = "servable"
    dataset = "dataset"
    pipeline = "piepline"


class DLHubMetadata(BaseModel):
    """Basic metadata for a DLHub artefact

    Includes information used by the DLHub web service to recognize an object,
    define how to build the computational environment,
    control access to it, et cetera"""

    version: str = Field(__version__, description="Version of the SDK used to create this object")
    domains: List[str] = Field(default_factory=list, description="List of scientific domains that describe this object")
    visible_to: List[str] = Field(['public'],
                                  description="List of groups or identities for which this object is visible")
    name: Optional[str] = Field(None, description="Name of the object", regex=r'^\S+$')
    files: Dict[str, Union[str, List[str]]] = Field(default_factory=dict,
                                                    description="List of files associated with this object")
    dependencies: Dict[str, Dict[str, str]] = Field(default_factory=dict,
                                                    description="Library dependencies for an object")
    test: Optional[bool] = Field(None, description="Internal use only. Used for testing the DLHub service")
    type: DLHubType = Field(DLHubType.servable, description="Type of this object")


class BaseMetadataModel(BaseModel):
    """Base class for models describing objects published via DLHub

    Covers information that goes in the :code:`datacite` block of the metadata file and
    some of the DLHub block.

    Using a MetadataModel
    ---------------------

    There are many kinds of MetadataModel classes that each describe a different kind of object.
    Each of these different types are created using the :code:`create_model` operation
    (e.g., :code:`KerasModel.create_model('model.h5')`), but have different arguments depending
    on the type of object. For example, TensorFlow models only require the directory created
    when saving the model for serving but scikit-learn models require the pickle file, how
    the pickle was created (e.g., with joblib), and how many input features it requires.

    Once created, you will need to fill in additional details about the object to make it reusable.
    The MetadataModel classes attempt to learn as much about an object as possible automatically,
    but there is some information that must be provided by a human. To start, you must define a
    title and name for the object and are encouraged to provide an abstract describing the model
    and list any associated papers/websites that describe the model. You will fill plenty of
    examples for how to describe the models in the DLHub_containers repostiory. Some types of
    objects require data specific to their type (e.g., Python servables need a list of required
    packages). We encourage you to find examples for your specific type of object in the containers
    repository for inspiration and to see the Python documentation for each Metadata Model.

    The MetadataModel object can be saved using the `to_dict` operation and read back into
    memory using the `from_dict` method. We recommend you save your dictionary to disk in the
    JSON or yaml format, which will allow for manual edits to be made before submitting or
    resubmitting a object description.
    """

    dlhub: DLHubMetadata = Field(default_factory=DLHubMetadata, description="Metadata used by the DLHub web service")
    datacite: Datacite = Field(default_factory=Datacite)

    @property
    def name(self) -> str:
        """
        Get the name of the servable

        Returns:
            (string) Name of the servable
        """
        return self.dlhub.name

    @classmethod
    def create_model(cls, **kwargs) -> 'BaseMetadataModel':
        """Instantiate the metadata model.

        Takes in arguments that allow metadata describing a dataset to be autogenerated. For
        example, these could include options describing how to read a dataset from a CSV file or
        which class method to invoke on a Python pickle object.
        """
        return cls()

    def read_codemeta_file(self, directory: Optional[str] = None):
        """Read in metadata from a codemeta.json file

        Args:
            directory (string): Path to directory contain the codemeta.json file (default: current working directory)
        """

        # If no directory, use the current
        if directory is None:
            directory = os.getcwd()

        # Load in the codemeta
        with open(os.path.join(directory, 'codemeta.json')) as fp:
            codemeta = json.load(fp)

        # Convert it to datacite and store it
        self.datacite = Datacite.from_codemeta(codemeta)

        return self

    def set_creators(self, authors: List[str], affiliations: List[Sequence[str]] = ()) -> 'BaseMetadataModel':
        """Add authors to this object

        Args:
            authors: List of authors for the dataset. In format: "<Family Name>, <Given Name>"
            affiliations: List of affiliations for each author.
        """
        self.datacite.set_creators(authors, affiliations)
        return self

    def set_title(self, title: str) -> 'BaseMetadataModel':
        """Set the title for this object

        Args:
            title: Desired title
        """
        self.datacite.set_title(title)
        return self

    def set_version(self, version):
        """Set the version of this resource

        Args:
            version (string): Version number
        """
        self.datacite.version = str(version)
        return self

    def set_abstract(self, abstract: str) -> 'BaseMetadataModel':
        """Define an abstract for this object. Use for a high-level summary

        Args:
            abstract: Description of this artifact
        """
        self.datacite.set_abstract(abstract)
        return self

    def set_methods(self, methods: str) -> 'BaseMetadataModel':
        """Define a methods section for this object. Use to describe any specific details
        about how the dataset, model, etc was generated.

        Args:
            methods (str): Detailed method descriptions
        """
        self.datacite.set_methods(methods)
        return self

    def set_domains(self, domains: List[str]):
        """Set the field of science that is associated with this object

        Args:
            domains: Names of fields of science (e.g., "materials science")
        """
        self.dlhub.domains = domains.copy()
        return self

    def set_visibility(self, users: Optional[Iterable[str]] = None,
                       groups: Optional[Iterable[str]] = None):
        """Set the list of people this object should be visible to.

        By default, it will be visible to anyone (["public"]).

        Args:
            users: GlobusAuth UUIDs of allowed users
            groups: GlobusAuth UUIDs of allowed Globus groups
        """
        visibilities = []
        if not (users or groups):
            visibilities.append('public')
        else:
            if users:
                for user in users:
                    visibilities.append('urn:globus:auth:identity:' + user)
            if groups:
                for group in groups:
                    visibilities.append('urn:globus:groups:id:' + group)
        self.dlhub.visible_to = visibilities
        return self

    def set_doi(self, doi):
        """Set the DOI of this object, if available

        This function is only for advanced usage. Most users of the toolbox will not
        know the DOI before sending the object in to DLHub.

        Args:
            doi (string): DOI of the object
        """
        self.datacite.set_identifier(doi)
        return self

    def set_name(self, name):
        """Set the name of the object.

        Should be something short, descriptive, and memorable

        Args:
            name (string): Name of artifact
        """
        if name_re.match(name) is None:
            raise ValueError('Name cannot contain any whitespace')
        self.dlhub.name = name
        return self

    def set_publication_year(self, year):
        """Define the publication year

        This function is only for advanced usage. Normally, this will be assigned automatically

        Args:
            year (string): Publication year
        """
        self.datacite.publicationYear = str(year)
        return self

    def add_requirement(self, library, version=None):
        """Add a required Python library.

        The name of the library should be either the name on PyPI, or a URL for the git repository
        holding the code (e.g., ``git+https://github.com/DLHub-Argonne/dlhub_sdk.git``)

        Args:
            library (string): Name of library
            version (string): Required version. 'latest' to use the most recent version on PyPi (if
                available). 'detect' will attempt to find the version of the library installed on
                the computer running this software. Default is ``None``
        """

        # Attempt to determine the version automatically
        if version == "detect":
            try:
                module = importlib.import_module(library)
                version = module.__version__
            except (AttributeError, ModuleNotFoundError):
                version = pkg_resources.get_distribution(library).version
        elif version == "latest":
            pypi_req = requests.get('https://pypi.org/pypi/{}/json'.format(library))
            version = pypi_req.json()['info']['version']

        # Set the requirements
        if "python" not in self.dlhub.dependencies:
            self.dlhub.dependencies["python"] = {}
        self.dlhub.dependencies["python"][library] = version
        return self

    def add_requirements(self, requirements):
        """Add several Python library requirements

        Args:
            requirements (dict): Keys are names of library (str), values are the version
        """
        for p, v in requirements.items():
            self.add_requirement(p, v)
        return self

    def parse_repo2docker_configuration(self, directory=None):
        """Gathers information about required environment from repo2docker configuration files.

        See https://repo2docker.readthedocs.io/en/latest/config_files.html for more details

        Args:
            directory (str): Path to directory containing configuration files
                (default: current working directory)
        """

        # Get a list of all files
        config_files = ['environment.yml', 'requirements.txt', 'setup.py', 'REQUIRE', 'install.R',
                        'apt.txt', 'DESCRIPTION', 'manifest.xml', 'postBuild', 'start',
                        'runtime.txt', 'default.nix', 'Dockerfile']

        # Get the directory name if `None`
        if directory is None:
            directory = os.getcwd()

        # Add every file we can find
        for file in config_files:
            path = os.path.join(directory, file)
            if os.path.isfile(path):
                self.add_file(path)

        return self

    def add_related_resource(self, identifier: str, identifier_type: Union[str, DataciteRelatedIdentifierType],
                             relation_type: Union[str, DataciteRelationType]) -> 'BaseMetadataModel':
        """Add a resource that is related to this object.

        We use the DataCite to describe the relations. Common relation types for DLHub objects are:
            - "IsDescribedBy": For a paper that describes a dataset or model
            - "IsDocumentedBy": For the software documentation for a model
            - "IsDerviedFrom": For the database a training set was pulled from
            - "Requires": For any software libraries that are required for this module

        Args:
            identifier: Identifier
            identifier_type: Identifier type
            relation_type: Relation between this identifier and the object you are describing
        """
        self.datacite.add_related_identifier(identifier, identifier_type, relation_type)

    def add_file(self, file, name=None):
        """Add a file to the list of files to be distributed with this object

        Args:
            file (string): Path to the file
            name (string): Optional. Name of the file, if it is a file that serves a specific
                purpose (e.g., "pickle" if this is a pickle file of a scikit-learn model)
        """

        if name is None or name == "other":
            if "other" not in self.dlhub.files:
                self.dlhub.files["other"] = []
            self.dlhub.files["other"].append(file)
        else:
            self.dlhub.files[name] = file
        return self

    def add_directory(self, directory: str, include: Union[str, Iterable[str]] = (),
                      exclude: Union[str, Iterable[str]] = (), recursive: bool = False):
        """Add all the files in a directory

        Args:
            include (string or [string]): Only add files that match any of these patterns
            exclude (string or [string]): Exclude all files that match any of these patterns
            directory (string): Path to a directory
            recursive (bool): Whether to add all files in a directory
        """

        # Make sure include and exclude are lists
        if isinstance(include, string_types):
            include = [include]
        if isinstance(exclude, string_types):
            exclude = [exclude]

        # Get potential files
        if recursive:
            if sys.version_info < (3, 5):
                raise RuntimeError('You must use Python >= 3.5 to do recursive listing')
            hits = glob('{}/**/*'.format(directory), recursive=True)
        else:
            hits = glob('{}/*'.format(directory))

        # Get only the files that match the filters
        files = [x for x in hits if os.path.isfile(x)]
        if len(include) > 0:  # Run inclusive filters
            files = [f for f in files
                     if any(fnmatch.fnmatch(os.path.basename(f), i) for i in include)]
        if len(exclude) > 0:
            files = [f for f in files
                     if not any(fnmatch.fnmatch(os.path.basename(f), e) for e in exclude)]

        return self.add_files(files)

    def add_files(self, files: Iterable[str]):
        """Add files that should be distributed with this artifact.

        Args:
            files: Paths of files that should be published
        """
        # Add the files
        for file in files:
            self.add_file(file)
        return self

    def to_dict(self, simplify_paths: bool = False):
        """Render the object to a JSON-ready dictionary

        Args:
            simplify_paths (bool): Whether to simplify the paths of each file
        Returns:
            (dict) A description of the dataset in a form suitable for download
        """

        # Check for required fields
        if len(self.datacite.titles) == 0:
            raise ValueError('Title must be specified. Use `set_title`')
        if self.name is None:
            raise ValueError('Name must be specified. Use `set_name`')

        # Make a copy of the output
        out = self.dict(exclude_none=True).copy()

        # Prepare the files
        if simplify_paths:
            # Get the common path
            common_path = self._get_common_path()

            files = {}
            for k, v in out["dlhub"]["files"].items():
                if k == "other":
                    files[k] = [os.path.relpath(f, common_path) for f in v]
                else:
                    files[k] = os.path.relpath(v, common_path)

            # Copy over the current files list
            out["dlhub"]["files"] = files
        return out

    def list_files(self):
        """Provide a list of files associated with this artifact.

        Returns:
            ([string]) list of file paths"""
        # Gather a list of all the files
        output = []
        for k, v in self.dlhub.files.items():
            if isinstance(v, string_types):
                output.append(v)
            else:  # It is a list
                output.extend(v)
        return output

    def get_zip_file(self, path):
        """Write all the listed files to a ZIP object

        Takes all of the files returned by `list_files`. First determines the largest common
        path of all files, and preserves directory structure by using this common path as the
        root directory. For example, if the files are "/home/a.pkl" and "/home/a/b.dat", the common
        directory is "/home" and the files will be stored in the Zip as "a.pkl" and "a/b.dat"

        Args:
            path (string): Path for the ZIP File
        Returns:
            (string): Base path for the ZIP file (useful for adjusting the paths of the files
                included in the metadata model)
        """

        # Open the zip file in "exclusively create" (x) mode
        with ZipFile(path, 'x') as newzip:
            if len(self.list_files()) == 0:
                return "."

            # Get the common path of all files
            root_path = self._get_common_path()

            # Add each file to the directory
            for file in self.list_files():
                newzip.write(file, arcname=os.path.relpath(file, root_path))

            return root_path

    def _get_common_path(self):
        """Determine the common path of all files

        Returns:
            (string) Common path
        """
        # Get the files
        files = self.list_files()

        # Shortcut: if no files
        if len(files) == 0:
            return '.'

        # Get the directories for all the files
        directories = [f if os.path.isdir(f) else os.path.dirname(f) for f in files]

        # Get the largest common path
        return os.path.commonpath(os.path.abspath(d) for d in directories)
