from __future__ import annotations
from typing import Generator, Iterator

from contextlib import contextmanager

from tgzr.context import Context, ContextData

#
#   Contexts Definition
#


@Context.context_type
class ProjectCtx(ContextData):
    project_name: str
    project_id: int
    config_name: str

    def task_depts(self) -> list[str]:
        return ["core", "service", "app"]

    def task_names(self) -> list[str]:
        if self.config_name == "DEV":
            return ["impl", "build", "release"]
        else:
            return ["spec", "test", "report"]

    def user_names(self) -> list[str]:
        if self.config_name == "DEV":
            return ["Bob-Dev", "Alice-Dev"]
        else:
            return ["Bob", "Alice"]

    @contextmanager
    def config(self, config_name: str) -> Generator[ProjectCtx, None, None]:
        with self.get(
            ProjectCtx,
            config_name=config_name,
        ) as p:
            yield p

    @contextmanager
    def user(self, user_name: str) -> Generator[UserCtx, None, None]:
        with self.get(UserCtx, user_name=user_name) as user:
            yield user


@Context.context_type
class UserCtx(ProjectCtx):
    user_name: str

    def iter_tasks(self) -> Iterator[TaskCtx]:
        for task_dept in self.task_depts():
            for name in self.task_names():
                with self.task(task_name=name, task_dept=task_dept) as task:
                    yield task

    @contextmanager
    def task(self, task_name: str, task_dept: str) -> Generator[TaskCtx, None, None]:
        with self.get(TaskCtx, task_name=task_name, task_dept=task_dept) as task:
            yield task


@Context.context_type
class TaskCtx(UserCtx):
    task_name: str
    task_dept: str


def demo_project_context_usage():

    # We can define a function which will work with the current
    # contexts, without the need to pass them as arguments:
    def process_config(config_name: str):
        """
        Prints all the tasks in the current project when
        the given config_name is applied.
        """
        # Get the current project:
        project = ProjectCtx.current()
        # Get the same project with a different config:
        with project.config(config_name=config_name) as project_config:
            # This is now the current Project:
            assert project_config is ProjectCtx.current()
            # The previously current project is not the same because it had a different config:
            assert project is not ProjectCtx.current()
            # But the previously current project has the same id:
            assert project.project_id == ProjectCtx.current().project_id
            # NB: the previous line was 90% written by autocompletion <3

            # There is no current User yet:
            assert UserCtx.current(False) is None

            # Let's iterate over the project's users:
            for user_name in project_config.user_names():
                with project_config.user(user_name=user_name):
                    # There is a current User now:
                    assert UserCtx.current() is not None

                    # Let's iterate over the user's tasks:
                    for task in UserCtx.current().iter_tasks():
                        # The current Task is the one we're using:
                        assert task == TaskCtx.current()
                        # The current User is the current Task,
                        # but since Task is a sub-context of User, you get all the completion for a User
                        assert task == UserCtx.current()
                        # The current Project is the current Task,
                        # but since Task inherit Project you get all the completion for a Project
                        assert task == ProjectCtx.current() and isinstance(
                            task, UserCtx
                        )
                        # The current Context is the current Task
                        # and you get completion for a Task
                        assert task == Context.current() and isinstance(task, TaskCtx)
                        print(
                            "Task:",
                            # All these attributes are avaible in completion:
                            task.project_name,
                            task.project_id,
                            task.config_name,
                            task.user_name,
                            task.task_dept,
                            task.task_name,
                        )

    # We create a new context using the Context.get() classmethod:
    with Context.get(
        ProjectCtx, project_name="MyProject", project_id=1234, config_name=None
    ) as project:
        # The ProjectCtx context became "the current ProjectCtx", and since we
        # gave it `project_name="MyProject"`, we made this Project the current one:
        assert project == ProjectCtx.current()

        # We can call procedure acting on this project without passing the project,
        # the procedure will use the current one:
        process_config(config_name="DEV")
        process_config(config_name="PROD")


if __name__ == "__main__":
    demo_project_context_usage()
