import sys
from os import getcwd, getenv
from os.path import abspath
from pathlib import Path
from typing import Optional

from kymata.io.file import PathType
from kymata.system.reflection import kymata_installed_as_dependency


_DATA_PATH_ENVIRONMENT_VAR_NAME = "KYMATA_DATA_ROOT"
DATA_DIR_NAME = "kymata-core-data"


def _default_location_when_source() -> Path:
    """The default location for the data dir (parent) when kymata is being run from source."""
    return Path(
        Path(__file__)  # data_root.py
        .parent         # datasets
        .parent         # kymata
        .parent         # kymata-core
    )  # kymata/../data_dir (next to kymata dir)


def _default_location_when_dependency() -> Path:
    """The default location for the data dir (parent) when kymata is being run as an installed package."""

    # Generic "if __name__ == "__main__""
    main_module = sys.modules.get("__main__")
    if main_module and hasattr(main_module, "__file__"):
        return Path(abspath(main_module.__file__))

    # Otherwise return the location of the called script (may not work if the called script was an invoker)
    elif sys.argv[0]:
        return Path(abspath(sys.argv[0]))

    # Otherwise fall back to the existing default
    else:
        return _default_location_when_source()


# Places downloaded datasets could go, in order of preference
_preferred_default_data_locations = [
    _default_location_when_dependency() if kymata_installed_as_dependency() else _default_location_when_source(),
    Path(getcwd()),  # <cwd>/data_dir
    Path(Path.home(), "Documents"),  # ~/Documents/data_dir
    Path(Path.home()),  # ~/data_dir
]


def data_root_path(data_root: Optional[PathType] = None) -> Path:
    """
    Retrieve or create the path to the data root directory.

    This function checks for the data root directory in the following order:
    1. If provided as an argument.
    2. If specified in the environment variable defined by `$KYMATA_DATA_ROOT`.
    3. If not specified, it checks preferred default locations
       and attempts to create the directory if it does not exist.

    Parameters:
        data_root (Optional[PathType]): An optional path to the data root.
                                         If None, the function will check the
                                         environment variable and default locations.

    Returns:
        Path: Path object representing the data root directory.

    Raises:
        FileNotFoundError: If the specified data root does not exist or cannot be created.
        NotADirectoryError: If the specified data root is not a directory.
    """

    # Check if the data root has been specified

    # Might be in an environmental variable
    if data_root is None:
        data_root: PathType | None = getenv(
            _DATA_PATH_ENVIRONMENT_VAR_NAME, default=None
        )

    # Might have been supplied as an argument
    if data_root is not None:
        if isinstance(data_root, str):
            data_root = Path(data_root)
        # Data root specified
        if not data_root.exists():
            raise FileNotFoundError(
                f"data_root {str(data_root)} specified but does not exist"
            )
        if not data_root.is_dir():
            raise NotADirectoryError(
                f"Please specify a directory ({str(data_root)} is not a directory)"
            )

        return data_root

    else:
        # Data root not specified

        # Check if the data root already exists
        for loc in _preferred_default_data_locations:
            if (here := Path(loc, DATA_DIR_NAME)).exists():
                data_root = here
                break

        # If not, attempt to create it
        if data_root is None:
            here: Path | None = None
            for loc in _preferred_default_data_locations:
                here = Path(loc, DATA_DIR_NAME)
                try:
                    here.mkdir()
                    break
                # If it fails for sensible reasons, no sweat, we'll fall through to the next option
                except (FileNotFoundError, OSError):
                    # Parent didn't exist, not writeable, etc
                    pass
            # Did we make it?
            if here is not None and here.exists():
                data_root = here
            else:
                raise FileNotFoundError("Failed to create data root directory")

        # Data root location has been derived, rather than prespecified, so feed that back to the user to avoid a
        # different location somehow being derived next time
        print(f"Data root set at {str(data_root)}.")
        print(
            f"Consider setting this as environmental variable {_DATA_PATH_ENVIRONMENT_VAR_NAME} to ensure it's reused"
            f" next time."
        )
        print(f'Hint: $> {_DATA_PATH_ENVIRONMENT_VAR_NAME}="{str(data_root)}"')
        return data_root
