import shutil
from pathlib import Path
from typing import ClassVar, Final, Literal, override

from cstar.base.log import LoggingMixin


class JobFileSystemManager(LoggingMixin):
    """JobFileSystem establishes a common directory structure convention across
    all jobs.
    """

    _INPUT_NAME: ClassVar[Literal["input"]] = "input"
    _WORK_NAME: ClassVar[Literal["work"]] = "work"
    _TASKS_NAME: ClassVar[Literal["tasks"]] = "tasks"
    _LOGS_NAME: ClassVar[Literal["logs"]] = "logs"
    _OUTPUT_NAME: ClassVar[Literal["output"]] = "output"

    root: Final[Path]
    """The root directory of the job. It is provided via user configuration."""

    def __init__(self, root_directory: Path) -> None:
        """Initialize the job file system.

        Parameters
        ----------
        root_directory : Path
            The root directory that contains all job byproducts.
        """
        self.root = root_directory

    def _dir_set(self) -> set[Path]:
        """Return the complete set of directories for the job.

        Returns
        -------
        set[Path]
        """
        return {
            self.input_dir,
            self.work_dir,
            self.tasks_dir,
            self.logs_dir,
            self.output_dir,
        }

    @property
    def input_dir(self) -> Path:
        """The directory for job inputs, such as datasets, code repositories, and
        runtime configuration.
        """
        return self.root / self._INPUT_NAME

    @property
    def work_dir(self) -> Path:
        """The directory for job work items, such as shell scripts generated by the
        system.
        """
        return self.root / self._WORK_NAME

    @property
    def tasks_dir(self) -> Path:
        """The directory for subtasks of a job, such as automatic system tasks
        or those generated via transformations of the inputs.
        """
        return self.root / self._TASKS_NAME

    @property
    def logs_dir(self) -> Path:
        """The directory for log files created during the job."""
        return self.root / self._LOGS_NAME

    @property
    def output_dir(self) -> Path:
        """The directory for writing any outputs from the job."""
        return self.root / self._OUTPUT_NAME

    def prepare(self) -> None:
        """Construct the directory tree for the job."""
        for task_dir in self._dir_set():
            task_dir.mkdir(parents=True, exist_ok=True)
            self.log.debug(f"Created {task_dir.name} directory `{task_dir}`")

    def clear(self) -> None:
        """Ensure the job's working directory is empty."""
        shutil.rmtree(self.root)
        self.root.mkdir(parents=True)
        self.log.debug(f"Created empty working directory for job `{self.root.name}`")

    def __getstate__(self) -> dict[str, str]:
        """Return the state of the object."""
        return {"root": self.root.as_posix()}

    def __setstate__(self, state: dict[str, str]) -> None:
        """Restore the object from a pickle."""
        self.__dict__.update({"root": Path(state["root"])})


class RomsFileSystemManager(JobFileSystemManager):
    _COMPILE_TIME_NAME: ClassVar[Literal["compile_time_code"]] = "compile_time_code"
    _RUNTIME_NAME: ClassVar[Literal["runtime_code"]] = "runtime_code"
    _INPUT_DATASETS_NAME: ClassVar[Literal["input_datasets"]] = "input_datasets"
    _CODEBASES_NAME: ClassVar[Literal["codebases"]] = "codebases"
    _JOINED_OUTPUT_NAME: ClassVar[Literal["joined_output"]] = "joined_output"

    def __init__(self, root_directory: Path) -> None:
        super().__init__(root_directory)

    @override
    def _dir_set(self) -> set[Path]:
        return (
            super()
            ._dir_set()
            .union(
                {
                    self.compile_time_code_dir,
                    self.runtime_code_dir,
                    self.input_datasets_dir,
                    self._codebases_dir,
                    self.joined_output_dir,
                }
            )
        )

    @property
    def compile_time_code_dir(self) -> Path:
        """The directory for compile-time code."""
        return self.root / self._COMPILE_TIME_NAME

    @property
    def runtime_code_dir(self) -> Path:
        """The directory for runtime code."""
        return self.root / self._RUNTIME_NAME

    @property
    def input_datasets_dir(self) -> Path:
        """The directory for input datasets."""
        return self.root / self._INPUT_DATASETS_NAME

    @property
    def _codebases_dir(self) -> Path:
        """The directory for codebases."""
        return self.root / self._CODEBASES_NAME

    @property
    def joined_output_dir(self) -> Path:
        """The directory for de-partitioned outputs."""
        return self.root / self._JOINED_OUTPUT_NAME

    def codebase_subdir(self, key: str) -> Path:
        """Return a codebase subdirectory path.

        Returns
        -------
        str
        """
        return self._codebases_dir / key
