from functools import update_wrapper
from typing import Any, Callable, Optional, Union

from dagster import check
from dagster.core.errors import DagsterInvalidDefinitionError

from ..graph import GraphDefinition
from ..partition import PartitionSetDefinition
from ..pipeline import PipelineDefinition
from ..repository import (
    VALID_REPOSITORY_DATA_DICT_KEYS,
    CachingRepositoryData,
    RepositoryData,
    RepositoryDefinition,
)
from ..schedule import ScheduleDefinition
from ..sensor import SensorDefinition


class _Repository:
    def __init__(self, name: Optional[str] = None, description: Optional[str] = None):
        self.name = check.opt_str_param(name, "name")
        self.description = check.opt_str_param(description, "description")

    def __call__(self, fn: Callable[[], Any]) -> RepositoryDefinition:
        check.callable_param(fn, "fn")

        if not self.name:
            self.name = fn.__name__

        repository_definitions = fn()

        if not (
            isinstance(repository_definitions, list)
            or isinstance(repository_definitions, dict)
            or isinstance(repository_definitions, RepositoryData)
        ):
            raise DagsterInvalidDefinitionError(
                "Bad return value of type {type_} from repository construction function: must "
                "return list, dict, or RepositoryData. See the @repository decorator docstring for "
                "details and examples".format(type_=type(repository_definitions)),
            )

        if isinstance(repository_definitions, list):
            bad_definitions = []
            for i, definition in enumerate(repository_definitions):
                if not (
                    isinstance(definition, PipelineDefinition)
                    or isinstance(definition, PartitionSetDefinition)
                    or isinstance(definition, ScheduleDefinition)
                    or isinstance(definition, SensorDefinition)
                    or isinstance(definition, GraphDefinition)
                ):
                    bad_definitions.append((i, type(definition)))
            if bad_definitions:
                bad_definitions_str = ", ".join(
                    [
                        "value of type {type_} at index {i}".format(type_=type_, i=i)
                        for i, type_ in bad_definitions
                    ]
                )
                raise DagsterInvalidDefinitionError(
                    "Bad return value from repository construction function: all elements of list "
                    "must be of type JobDefinition, GraphDefinition, PipelineDefinition, "
                    "PartitionSetDefinition, ScheduleDefinition, or SensorDefinition. "
                    f"Got {bad_definitions_str}."
                )
            repository_data = CachingRepositoryData.from_list(repository_definitions)

        elif isinstance(repository_definitions, dict):
            if not set(repository_definitions.keys()).issubset(VALID_REPOSITORY_DATA_DICT_KEYS):
                raise DagsterInvalidDefinitionError(
                    "Bad return value from repository construction function: dict must not contain "
                    "keys other than {{'pipelines', 'partition_sets', 'schedules', 'jobs'}}: found "
                    "{bad_keys}".format(
                        bad_keys=", ".join(
                            [
                                "'{key}'".format(key=key)
                                for key in repository_definitions.keys()
                                if key not in VALID_REPOSITORY_DATA_DICT_KEYS
                            ]
                        )
                    )
                )
            repository_data = CachingRepositoryData.from_dict(repository_definitions)
        elif isinstance(repository_definitions, RepositoryData):
            repository_data = repository_definitions

        repository_def = RepositoryDefinition(
            name=self.name, description=self.description, repository_data=repository_data
        )

        update_wrapper(repository_def, fn)
        return repository_def


def repository(
    name: Union[Optional[str], Callable[..., Any]] = None, description: Optional[str] = None
) -> Union[_Repository, RepositoryDefinition]:
    """Create a repository from the decorated function.

    The decorated function should take no arguments and its return value should one of:

    1. ``List[Union[JobDefinition, PipelineDefinition, PartitionSetDefinition, ScheduleDefinition, SensorDefinition]]``.
        Use this form when you have no need to lazy load pipelines or other definitions. This is the
        typical use case.

    2. A dict of the form:

    .. code-block:: python

        {
            'jobs': Dict[str, Callable[[], JobDefinition]],
            'pipelines': Dict[str, Callable[[], PipelineDefinition]],
            'partition_sets': Dict[str, Callable[[], PartitionSetDefinition]],
            'schedules': Dict[str, Callable[[], ScheduleDefinition]]
            'sensors': Dict[str, Callable[[], SensorDefinition]]
        }

    This form is intended to allow definitions to be created lazily when accessed by name,
    which can be helpful for performance when there are many definitions in a repository, or
    when constructing the definitions is costly.

    3. :py:class:`RepositoryData`. Return this object if you need fine-grained
        control over the construction and indexing of definitions within the repository, e.g., to
        create definitions dynamically from .yaml files in a directory.

    Args:
        name (Optional[str]): The name of the repository. Defaults to the name of the decorated
            function.
        description (Optional[str]): A string description of the repository.

    Example:

    .. code-block:: python

        ######################################################################
        # A simple repository using the first form of the decorated function
        ######################################################################

        @op(config_schema={n: Field(Int)})
        def return_n(context):
            return context.op_config['n']

        @job
        def simple_job():
            return_n()

        @job
        def some_job():
            ...

        @sensor(job=some_job)
        def some_sensor():
            if foo():
                yield RunRequest(
                    run_key= ...,
                    run_config={
                        'ops': {'return_n': {'config': {'n': bar()}}}
                    }
                )

        @job
        def my_job():
            ...

        my_schedule = ScheduleDefinition(cron_schedule="0 0 * * *", job=my_job)

        @repository
        def simple_repository():
            return [simple_job, some_sensor, my_schedule]


        ######################################################################
        # A lazy-loaded repository
        ######################################################################

        def make_expensive_job():
            @job
            def expensive_job():
                for i in range(10000):
                    return_n.alias(f'return_n_{i}')()

            return expensive_job

        def make_expensive_schedule():
            @job
            def other_expensive_job():
                for i in range(11000):
                    return_n.alias(f'my_return_n_{i}')()

            return ScheduleDefinition(cron_schedule="0 0 * * *", job=other_expensive_job)

        @repository
        def lazy_loaded_repository():
            return {
                'jobs': {'expensive_job': make_expensive_job},
                'schedules': {'expensive_schedule: make_expensive_schedule}
            }


        ######################################################################
        # A complex repository that lazily constructs jobs from a directory
        # of files in a bespoke YAML format
        ######################################################################

        class ComplexRepositoryData(RepositoryData):
            def __init__(self, yaml_directory):
                self._yaml_directory = yaml_directory

            def get_all_jobs(self):
                return [
                    self._construct_job_def_from_yaml_file(
                      self._yaml_file_for_job_name(file_name)
                    )
                    for file_name in os.listdir(self._yaml_directory)
                ]

            ...

        @repository
        def complex_repository():
            return ComplexRepositoryData('some_directory')

    """
    if callable(name):
        check.invariant(description is None)

        return _Repository()(name)

    return _Repository(name=name, description=description)
