from collections.abc import Generator
from pathlib import Path
from typing import (
    Union,
)

import luigi
from luigi import Task

from exasol_integration_test_docker_environment.lib.base.stoppable_base_task import (
    StoppableBaseTask,
)
from exasol_integration_test_docker_environment.lib.base.task_dependency import (
    DependencyState,
    DependencyType,
    TaskDependency,
    TaskDescription,
)


class DependencyLoggerBaseTask(StoppableBaseTask):

    def _get_dependencies_path_for_job(self) -> Path:
        return Path(super()._get_output_path_for_job(), "dependencies")

    def _get_dependencies_path(self) -> Path:
        path = Path(self._get_dependencies_path_for_job(), self.task_id)
        path.mkdir(parents=True, exist_ok=True)
        return path

    def _get_dependencies_requires_path(self) -> Path:
        return Path(self._get_dependencies_path(), "requires")

    def _get_dependencies_dynamic_path(self) -> Path:
        return Path(self._get_dependencies_path(), "dynamic")

    def handle_requires_value(
        self, tasks: Union[Task, list, tuple, dict]
    ) -> Union[Task, list, tuple, dict]:
        dependency_path = self._get_dependencies_requires_path()
        if not dependency_path.exists():
            with dependency_path.open("w") as dependencies_file:
                self.write_dependency(
                    dependencies_file=dependencies_file,
                    dependency_type=DependencyType.requires,
                    dependency_state=DependencyState.requested,
                    index=0,
                    value=tasks,
                )
                return tasks
        else:
            return tasks

    def handle_requires_generator(self, tasks):
        dependency_path = self._get_dependencies_requires_path()
        if not dependency_path.exists():
            with dependency_path.open("w") as dependencies_file:
                result = list(
                    self.write_dependencies_for_generator(
                        dependencies_file=dependencies_file,
                        task_generator=tasks,
                        dependency_type=DependencyType.requires,
                    )
                )
                return result
        else:
            return tasks

    def write_dependencies_for_generator(
        self, dependencies_file, task_generator, dependency_type: DependencyType
    ):
        index = 0
        try:
            element = next(task_generator)
            while True:
                self.write_dependency(
                    dependencies_file=dependencies_file,
                    dependency_type=dependency_type,
                    dependency_state=DependencyState.requested,
                    index=index,
                    value=element,
                )
                result = yield element
                element = task_generator.send(result)
                index += 1
        except StopIteration:
            pass

    def write_dependency(
        self,
        dependencies_file,
        dependency_type: DependencyType,
        dependency_state: DependencyState,
        index: int,
        value,
    ):
        for task in self.flatten_tasks(value):
            dependency = TaskDependency(
                source=self.get_task_description(),
                target=task.get_task_description(),
                type=dependency_type,
                index=index,
                state=dependency_state,
            )
            dependencies_file.write(f"{dependency.to_json()}")
            dependencies_file.write("\n")

    def requires(self):
        tasks = super().requires()
        if tasks is not None:
            if isinstance(tasks, (Task, list, tuple, dict)):
                return self.handle_requires_value(tasks)
            else:
                return self.handle_requires_generator(tasks)
        else:
            return []

    def run(self):
        dependency_path = self._get_dependencies_dynamic_path()
        with dependency_path.open("a") as dependencies_file:
            task_generator = super().run()
            if task_generator is not None:
                yield from self.write_dependencies_for_generator(
                    dependencies_file=dependencies_file,
                    task_generator=task_generator,
                    dependency_type=DependencyType.dynamic,
                )

    def get_task_description(self) -> TaskDescription:
        return TaskDescription(id=self.task_id, representation=str(self))

    def flatten_tasks(self, generator: Generator) -> list["DependencyLoggerBaseTask"]:
        return [
            task
            for task in luigi.task.flatten(generator)
            if isinstance(task, DependencyLoggerBaseTask)
        ]
