import json
import logging
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, Iterable, Iterator, Optional, Set, TypeVar, Union
from urllib.parse import ParseResult

import dill
import petname
from sqlitedict import SqliteDict

from tango.common import PathOrStr
from tango.common.exceptions import StepStateError
from tango.common.file_lock import FileLock
from tango.common.logging import file_handler
from tango.common.util import exception_to_string, utc_now_datetime
from tango.step import Step
from tango.step_cache import StepCache
from tango.step_caches import LocalStepCache
from tango.step_info import StepInfo, StepState
from tango.workspace import Run, Workspace

logger = logging.getLogger(__name__)

T = TypeVar("T")


@Workspace.register("local")
class LocalWorkspace(Workspace):
    """
    This is a :class:`.Workspace` that keeps all its data in a local directory. This works great for single-machine
    jobs, or for multiple machines in a cluster if they can all access the same NFS drive.

    :param dir: The directory to store all the data in

    The directory will have three subdirectories, ``cache/`` for the step cache, ``runs/`` for the runs,
    and ``latest/`` for the results of the latest run. For the format of the ``cache/`` directory,
    refer to :class:`.LocalStepCache`. The ``runs/`` directory will contain one subdirectory for each
    registered run. Each one of those contains a symlink from the name of the step to the results directory
    in the step cache. Note that :class:`.LocalWorkspace` creates these symlinks even for steps that have not
    finished yet. You can tell the difference because either the symlink points to a directory that doesn't exist,
    or it points to a directory in the step cache that doesn't contain results.

    .. tip::

        Registered as a :class:`~tango.workspace.Workspace` under the name "local".

        You can also instantiate this workspace from a URL with the scheme ``local://``.
        For example, ``Workspace.from_url("local:///tmp/workspace")`` gives you a :class:`LocalWorkspace`
        in the directory ``/tmp/workspace``.

    """

    def __init__(self, dir: PathOrStr):
        super().__init__()
        self.dir = Path(dir)
        self.dir.mkdir(parents=True, exist_ok=True)
        self.cache = LocalStepCache(self.dir / "cache")
        self.locks: Dict[Step, FileLock] = {}
        self.runs_dir = self.dir / "runs"
        self.runs_dir.mkdir(parents=True, exist_ok=True)
        self.step_info_file = self.dir / "stepinfo.sqlite"
        self.latest_dir = self.dir / "latest"

        # Check the version of the local workspace
        try:
            with open(self.dir / "settings.json", "r") as settings_file:
                settings = json.load(settings_file)
        except FileNotFoundError:
            settings = {"version": 1}

        # Upgrade to version 2
        if settings["version"] == 1:
            with SqliteDict(self.step_info_file) as d:
                for stepinfo_file in self.cache.dir.glob("*/stepinfo.dill"):
                    with stepinfo_file.open("rb") as f:
                        stepinfo = dill.load(f)

                    # The `StepInfo` class changed from one version to the next. The deserialized version
                    # ends up being a `StepInfo` instance that is missing the `cacheable` member. This
                    # hack adds it in.
                    kwargs = stepinfo.__dict__
                    kwargs[
                        "cacheable"
                    ] = True  # Only cacheable steps were saved in v1. That's what v2 fixes.
                    d[stepinfo.unique_id] = StepInfo(**kwargs)
                d.commit()
            for stepinfo_file in self.cache.dir.glob("*/stepinfo.dill"):
                stepinfo_file.unlink()

            settings["version"] = 2
            with open(self.dir / "settings.json", "w") as settings_file:
                json.dump(settings, settings_file)

    def __getstate__(self):
        """
        We override `__getstate__()` to customize how instances of this class are pickled
        since we don't want to persist certain attributes.
        """
        out = super().__getstate__()
        out["locks"] = {}
        return out

    @property
    def url(self) -> str:
        return "local://" + str(self.dir)

    @classmethod
    def from_parsed_url(cls, parsed_url: ParseResult) -> "Workspace":
        workspace_dir: Path
        if parsed_url.netloc:
            workspace_dir = Path(parsed_url.netloc)
            if parsed_url.path:
                workspace_dir = workspace_dir / parsed_url.path.lstrip("/")
        elif parsed_url.path:
            workspace_dir = Path(parsed_url.path)
        else:
            workspace_dir = Path(".")
        return cls(workspace_dir.resolve())

    def step_dir(self, step_or_unique_id: Union[Step, str]) -> Path:
        return self.cache.step_dir(step_or_unique_id)

    @property
    def step_cache(self) -> StepCache:
        return self.cache

    def work_dir(self, step: Step) -> Path:
        result = self.step_dir(step) / "work"
        result.mkdir(parents=True, exist_ok=True)
        return result

    @classmethod
    def guess_step_dir_state(cls, dir: Path) -> Set[StepState]:
        """
        Returns the possible states of a given step dir, to the best of our knowledge.

        :param dir: the step dir to example
        :return: a set of possible states for the step
        """

        # If the directory doesn't exist, the step is incomplete or uncacheable.
        if not dir.exists():
            return {StepState.INCOMPLETE, StepState.UNCACHEABLE}

        # If the lock file exists and is locked, the step is running.
        lock_file = dir / "lock"
        if lock_file.exists():
            lock = FileLock(lock_file)
            try:
                lock.acquire(0.1)
                lock.release()
            except TimeoutError:
                return {StepState.RUNNING}

        # If the directory is empty except for the work dir and the lock file, the step is running, incomplete,
        # or failed. But it can't be running because then the lockfile would be locked, so it can only be
        # incomplete or failed.
        for dir_entry in dir.iterdir():
            if dir_entry.name == "work" and dir_entry.is_dir():
                continue
            if dir_entry.name == "lock" and dir_entry.is_file():
                continue
            break
        else:
            return {StepState.INCOMPLETE, StepState.FAILED}

        return set(StepState)

    @staticmethod
    def _fix_step_info(step_info: StepInfo) -> None:
        """
        Tragically we need to run a fix-up step over StepInfo objects that are freshly read from
        the database. This is for backwards compatibility.

        This function operates on the `step_info` object in place.
        """
        if isinstance(step_info.error, BaseException):
            step_info.error = exception_to_string(step_info.error)

    def step_info(self, step_or_unique_id: Union[Step, str]) -> StepInfo:
        with SqliteDict(self.step_info_file) as d:

            def find_or_add_step_info(step_or_unique_id: Union[Step, str]) -> StepInfo:
                if isinstance(step_or_unique_id, Step):
                    unique_id = step_or_unique_id.unique_id
                else:
                    unique_id = step_or_unique_id

                try:
                    step_info = d[unique_id]
                except KeyError:
                    if not isinstance(step_or_unique_id, Step):
                        raise
                    step = step_or_unique_id

                    for dep in step.dependencies:
                        find_or_add_step_info(dep)

                    step_info = StepInfo.new_from_step(step)
                    d[unique_id] = step_info
                    del step

                # Perform some sanity checks. Sqlite and the file system can get out of sync
                # when a process dies suddenly.
                step_dir = self.step_dir(unique_id)
                step_state_guesses = self.guess_step_dir_state(step_dir) or step_info.state
                if step_info.state not in step_state_guesses:
                    if step_info.state == StepState.RUNNING:
                        # We think the step is running, but it can't possibly be running, so we go ahead and
                        # assume the step is incomplete.
                        step_info.start_time = None
                        step_info.end_time = None
                        d[unique_id] = step_info
                    else:
                        possible_states = ", ".join(s.value for s in step_state_guesses)
                        raise IOError(
                            f"The step '{unique_id}' is marked as being {step_info.state.value}, but we "
                            f"determined it can only be one of {{{possible_states}}}. If you are positive "
                            f"this is a screw-up, delete the directory at '{step_dir}' and try again."
                        )

                return step_info

            result = find_or_add_step_info(step_or_unique_id)
            d.commit()
            self._fix_step_info(result)
            return result

    def _step_lock_file(self, step_or_unique_id: Union[Step, str]) -> Path:
        step_dir = self.step_dir(step_or_unique_id)
        step_dir.mkdir(parents=True, exist_ok=True)
        return step_dir / "lock"

    def step_starting(self, step: Step) -> None:
        # We don't do anything with uncacheable steps.
        if not step.cache_results:
            return

        # Gather the existing step info first. Step info automatically fixes itself if steps are
        # marked as "running" but are not locked. This happens, for example, when a process
        # gets killed. To make sure this works, we have to get the step info before we start
        # messing with locks.
        step_info = self.step_info(step)
        if step_info.state not in {StepState.INCOMPLETE, StepState.FAILED}:
            raise StepStateError(
                step,
                step_info.state,
                context="If you are certain the step is not running somewhere else, delete the lock "
                f"file at {self._step_lock_file(step)}.",
            )

        if step_info.state == StepState.FAILED:
            # Refresh environment metadata since it might be out-of-date now.
            step_info.refresh()

        lock = FileLock(self._step_lock_file(step), read_only_ok=True)
        lock.acquire_with_updates(desc=f"acquiring lock for '{step.name}'")
        self.locks[step] = lock

        try:
            step_info.start_time = utc_now_datetime()
            step_info.end_time = None
            step_info.error = None
            step_info.result_location = None
            with SqliteDict(self.step_info_file) as d:
                d[step.unique_id] = step_info
                d.commit()
        except:  # noqa: E722
            lock.release()
            del self.locks[step]
            raise

    def step_finished(self, step: Step, result: T) -> T:
        # We don't do anything with uncacheable steps.
        if not step.cache_results:
            return result

        lock = self.locks[step]

        step_info = self.step_info(step)
        if step_info.state != StepState.RUNNING:
            raise StepStateError(step, step_info.state)

        self.step_cache[step] = result
        if hasattr(result, "__next__"):
            assert isinstance(result, Iterator)
            # Caching the iterator will consume it, so we write it to the cache and then read from the cache
            # for the return value.
            result = self.step_cache[step]

        # Mark the step as finished
        step_info.end_time = utc_now_datetime()
        step_info.result_location = str(self.step_dir(step).absolute())
        with SqliteDict(self.step_info_file) as d:
            d[step.unique_id] = step_info
            d.commit()

        lock.release()
        del self.locks[step]

        return result

    def step_failed(self, step: Step, e: BaseException) -> None:
        # We don't do anything with uncacheable steps.
        if not step.cache_results:
            return

        lock = self.locks[step]

        try:
            step_info = self.step_info(step)
            if step_info.state != StepState.RUNNING:
                raise StepStateError(step, step_info.state)
            step_info.end_time = utc_now_datetime()
            step_info.error = exception_to_string(e)
            with SqliteDict(self.step_info_file) as d:
                d[step.unique_id] = step_info
                d.commit()
        finally:
            lock.release()
            del self.locks[step]

    def register_run(self, targets: Iterable[Step], name: Optional[str] = None) -> Run:
        # sanity check targets
        targets = list(targets)
        if name is None:
            while name is None or (self.runs_dir / name).exists():
                name = petname.generate()
        run_dir = self.runs_dir / name

        # clean any existing run directory
        if run_dir.exists():
            for filename in run_dir.iterdir():
                filename.unlink()
        else:
            run_dir.mkdir(parents=True, exist_ok=True)

        # write step info for all steps
        all_steps = set(targets)
        for step in targets:
            all_steps |= step.recursive_dependencies

        self._save_registered_run(name, all_steps)

        # write targets
        for target in targets:
            if target.cache_results:
                target_path = self.step_dir(target)
                (run_dir / target.name).symlink_to(os.path.relpath(target_path, run_dir))

        self.latest_dir.unlink(missing_ok=True)
        self.latest_dir.symlink_to(run_dir)

        return self.registered_run(name)

    def registered_runs(self) -> Dict[str, Run]:
        return {
            str(run_dir.name): self.registered_run(run_dir.name)
            for run_dir in self.runs_dir.iterdir()
            if run_dir.is_dir()
        }

    def registered_run(self, name: str) -> Run:
        run_dir = self.runs_dir / name
        if not run_dir.is_dir():
            raise KeyError(name)
        steps_for_run = self._load_registered_run(name)
        return Run(name, steps_for_run, datetime.fromtimestamp(run_dir.stat().st_ctime))

    def _run_step_info_file(self, name: str) -> Path:
        return self.runs_dir / name / "stepinfo.json"

    def _save_registered_run(self, name: str, all_steps: Iterable[Step]) -> None:
        step_unique_ids = {}
        with SqliteDict(self.step_info_file) as d:
            for step in all_steps:
                try:
                    step_info = d[step.unique_id]
                    step_info.name = step.name
                    d[step.unique_id] = step_info
                except KeyError:
                    d[step.unique_id] = StepInfo.new_from_step(step)
                step_unique_ids[step.name] = step.unique_id

            d.commit()

            run_step_info_file = self._run_step_info_file(name)
            with open(run_step_info_file, "w") as file_ref:
                json.dump(step_unique_ids, file_ref)

    def _load_registered_run(self, name: str) -> Dict[str, StepInfo]:
        run_step_info_file = self._run_step_info_file(name)
        try:
            with open(run_step_info_file, "r") as file_ref:
                step_ids = json.load(file_ref)
        except FileNotFoundError:
            # for backwards compatibility
            run_dir = self.runs_dir / name
            step_ids = {}
            for step_symlink in run_dir.iterdir():
                if not step_symlink.is_symlink():
                    continue
                step_name = str(step_symlink.name)
                unique_id = str(step_symlink.resolve().name)
                step_ids[step_name] = unique_id

        with SqliteDict(self.step_info_file, flag="r") as d:
            steps_for_run = {}
            for step_name, unique_id in step_ids.items():
                step_info = d[unique_id]
                assert isinstance(step_info, StepInfo)
                self._fix_step_info(step_info)
                steps_for_run[step_name] = step_info
            return steps_for_run

    def run_dir(self, name: str) -> Path:
        """
        Returns the directory where a given run is stored.

        :param name: The name of the run
        :return: The directory where the results of the run are stored

        If the run does not exist, this returns the directory where it will be stored if you call
        :meth:`register_run()` with that name.
        """
        return self.runs_dir / name

    def capture_logs_for_run(self, name: str):
        return file_handler(self.run_dir(name) / "out.log")
