import getpass
import json
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime
from importlib.metadata import version
from pathlib import Path
from typing import Any, assert_never, overload
from uuid import UUID

from tierkreis.controller.data.core import PortID
from tierkreis.controller.data.graph import NodeDef, NodeDefModel
from tierkreis.controller.data.location import Loc, OutputLoc, WorkerCallArgs
from tierkreis.controller.storage.data import (
    ExecutorDebugData,
    NodeDebugData,
    WorkflowMetaData,
)
from tierkreis.controller.storage.exceptions import EntryNotFound
from tierkreis.exceptions import TierkreisError

logger = logging.getLogger(__name__)


@dataclass
class StorageEntryMetadata:
    """Collection of commonly found metadata.

    Storage implementations should decide which are applicable."""

    st_mtime: float | None = None


class ControllerStorage(ABC):
    tkr_dir: Path
    workflow_id: UUID
    name: str | None

    @abstractmethod
    def delete(self, path: Path) -> None:
        """Delete the storage entry at the specified path.

        Also delete any related data of the form \"{path}/**/*\"."""

    @abstractmethod
    def exists(self, path: Path) -> bool:
        """Is there an entry in the storage at the specified path?"""

    @abstractmethod
    def link(self, src: Path, dst: Path) -> None:
        """The storage entry at `dst` should have the same value as the entry at `src`."""

    @abstractmethod
    def list_subpaths(self, path: Path) -> list[Path]:
        """List all the paths starting with the specified path.

        This is used when the number of entries can only be determined at runtime.
        For example in a map node."""

    @abstractmethod
    def mkdir(self, path: Path) -> None:
        """Create an empty directory (and parents) at this path.

        Probably only required for file-based storage."""

    @abstractmethod
    def read(self, path: Path) -> bytes:
        """Read the storage entry at the specified path."""

    @abstractmethod
    def stat(self, path: Path) -> StorageEntryMetadata:
        """Get applicable stats for storage entry."""

    @abstractmethod
    def touch(self, path: Path) -> None:
        """Create empty storage entry at the specified path."""

    @abstractmethod
    def write(self, path: Path, value: bytes) -> None:
        """Write the given bytes to the storage entry at the specified path."""

    @property
    def workflow_dir(self) -> Path:
        return self.tkr_dir / str(self.workflow_id)

    @property
    def logs_path(self) -> Path:
        return self.workflow_dir / "logs"

    @property
    def debug_path(self) -> Path:
        return self.workflow_dir / "debug"

    def _exec_data_path(self, node_location: Loc) -> Path:
        return self.debug_path / "executors" / str(node_location)

    def _node_debug_path(self) -> Path:
        return self.debug_path / "nodes"

    def _nodedef_path(self, node_location: Loc) -> Path:
        return self.workflow_dir / str(node_location) / "nodedef"

    def _worker_call_args_path(self, node_location: Loc) -> Path:
        return self.workflow_dir / str(node_location) / "definition"

    def _metadata_path(self, node_location: Loc) -> Path:
        return self.workflow_dir / str(node_location) / "_metadata"

    def _outputs_dir(self, node_location: Loc) -> Path:
        return self.workflow_dir / str(node_location) / "outputs"

    def _output_path(self, node_location: Loc, port_name: PortID) -> Path:
        return self._outputs_dir(node_location) / port_name

    def _done_path(self, node_location: Loc) -> Path:
        return self.workflow_dir / str(node_location) / "_done"

    def _error_path(self, node_location: Loc) -> Path:
        return self.workflow_dir / str(node_location) / "_error"

    def _worker_logs_path(self, node_location: Loc) -> Path:
        return self.workflow_dir / str(node_location) / "logs"

    def clean_graph_files(self) -> None:
        self.delete(self.workflow_dir)

    def write_node_def(self, node_location: Loc, node: NodeDef):
        bs = NodeDefModel(root=node).model_dump_json().encode()
        self.write(self._nodedef_path(node_location), bs)

    def read_node_def(self, node_location: Loc) -> NodeDef:
        bs = self.read(self._nodedef_path(node_location))
        return NodeDefModel(**json.loads(bs)).root

    def write_worker_call_args(
        self,
        node_location: Loc,
        function_name: str,
        inputs: dict[PortID, OutputLoc],
        output_list: list[PortID],
    ) -> Path:
        call_args_path = self._worker_call_args_path(node_location)
        node_definition = WorkerCallArgs(
            function_name=function_name,
            inputs={
                k: self._output_path(loc, port).relative_to(self.tkr_dir)
                for k, (loc, port) in inputs.items()
            },
            outputs={
                k: self._output_path(node_location, k).relative_to(self.tkr_dir)
                for k in output_list
            },
            output_dir=self._outputs_dir(node_location).relative_to(self.tkr_dir),
            done_path=self._done_path(node_location).relative_to(self.tkr_dir),
            error_path=self._worker_logs_path(node_location).relative_to(self.tkr_dir),
            logs_path=self.logs_path.relative_to(self.tkr_dir),
        )
        self.write(call_args_path, node_definition.model_dump_json().encode())
        self.mkdir(self._outputs_dir(node_location))

        if (parent := node_location.parent()) is not None:
            self.touch(self._metadata_path(parent))

        return call_args_path.relative_to(self.tkr_dir)

    def read_worker_call_args(self, node_location: Loc) -> WorkerCallArgs:
        node_definition_path = self._worker_call_args_path(node_location)
        return WorkerCallArgs(**json.loads(self.read(node_definition_path)))

    def link_outputs(
        self,
        new_location: Loc,
        new_port: PortID,
        old_location: Loc,
        old_port: PortID,
    ) -> None:
        new_dir = self._output_path(new_location, new_port)
        try:
            self.link(self._output_path(old_location, old_port), new_dir)
        except EntryNotFound:
            logger.info(
                f"Could not find {old_location}. "
                "Tasks using this location will try to use a default value if specified."
            )
        except OSError as e:
            raise TierkreisError(
                "Workflow already exists. Try running with a different ID or do_cleanup."
            ) from e

    def write_output(
        self, node_location: Loc, output_name: PortID, value: bytes
    ) -> Path:
        output_path = self._output_path(node_location, output_name)
        self.write(output_path, bytes(value))
        return output_path

    def read_output(self, node_location: Loc, output_name: PortID) -> bytes:
        return self.read(self._output_path(node_location, output_name))

    def read_errors(self, node_location: Loc = Loc()) -> str:
        if self.exists(self._worker_logs_path(node_location)):
            return self.read(self._worker_logs_path(node_location)).decode()
        if self.exists(self._error_path(node_location)):
            return self.read(self._error_path(node_location)).decode()
        return ""

    def write_node_errors(self, node_location: Loc, error_logs: str) -> None:
        self.write(self._worker_logs_path(node_location), error_logs.encode())

    def read_output_ports(self, node_location: Loc) -> list[PortID]:
        dir_list = self.list_subpaths(self._outputs_dir(node_location))
        dir_list.sort()
        return [x.name for x in dir_list]

    def is_node_started(self, node_location: Loc) -> bool:
        return self.exists(Path(self._nodedef_path(node_location)))

    def is_node_finished(self, node_location: Loc) -> bool:
        return self.exists(self._done_path(node_location))

    def latest_loop_iteration(self, loc: Loc) -> Loc:
        i = 0
        while self.is_node_started(loc.L(i + 1)):
            i += 1
        return loc.L(i)

    def node_has_error(self, node_location: Loc) -> bool:
        return self.exists(self._error_path(node_location))

    def mark_node_finished(self, node_location: Loc) -> None:
        self.touch(self._done_path(node_location))

        if (parent := node_location.parent()) is not None:
            self.touch(self._metadata_path(parent))

    def write_metadata(self, node_location: Loc) -> None:
        if node_location == Loc(""):
            self.write_workflow_metadata()
            return
        j = json.dumps({"name": self.name, "start_time": datetime.now().isoformat()})
        self.write(self._metadata_path(node_location), j.encode())

    def read_metadata(self, node_location: Loc) -> dict[str, Any]:
        return json.loads(self.read(self._metadata_path(node_location)))

    def write_workflow_metadata(
        self,
    ) -> None:
        wf_data = WorkflowMetaData(
            workflow_id=str(self.workflow_id),
            name=self.name,
            user_id=getpass.getuser(),
            tierkreis_version=version("tierkreis"),
            start_time=datetime.now().isoformat(),
            execution_count=1,
        )
        self.write(self._metadata_path(Loc("")), wf_data.model_dump_json().encode())

    def write_workflow_completion_time(self) -> None:
        try:
            data = self.read_metadata(Loc(""))
        except json.JSONDecodeError as e:
            # Only in memory should trigger this
            logger.error(
                "Invalid json found. Dumping completion time anyway.\n Error: %s", e.msg
            )
            data = {}
        data["completion_time"] = datetime.now().isoformat()
        self.write(self._metadata_path(Loc()), json.dumps(data).encode())

    def read_started_time(self, node_location: Loc) -> str | None:
        node_def = Path(self._nodedef_path(node_location))
        if not self.exists(node_def):
            return None
        since_epoch = self.stat(node_def).st_mtime
        if since_epoch is None:
            return None
        return datetime.fromtimestamp(since_epoch).isoformat()

    def read_finished_time(self, node_location: Loc) -> str | None:
        done = Path(self._done_path(node_location))
        if not self.exists(done):
            return None
        since_epoch = self.stat(done).st_mtime
        if since_epoch is None:
            return None
        return datetime.fromtimestamp(since_epoch).isoformat()

    def read_loop_trace(self, node_location: Loc, output_name: PortID) -> list[bytes]:
        definition = self.read_node_def(node_location)
        if definition.type != "loop":
            raise TierkreisError("Can only read traces from loop nodes.")
        result = []

        i = 0
        while self.is_node_started(node_location.L(i)):
            result.append(self.read_output(node_location.L(i), output_name))
            i += 1
        return result

    def loc_from_node_name(self, node_name: str) -> Loc | None:
        debug_data = NodeDebugData(**self.read_debug_data(node_name))
        if debug_data.loop_loc is not None:
            return Loc(debug_data.loop_loc)

    def write_debug_data(self, name: str, loc: Loc) -> None:
        data = {name: NodeDebugData(loop_loc=loc).model_dump()}
        if not self.exists(self._node_debug_path()):
            self.write(self._node_debug_path(), json.dumps(data).encode())
            return

        existing_data = json.loads(self.read(self._node_debug_path()))
        if not isinstance(existing_data, dict):
            msg = f"Expecting executor data to be dict, found {type(existing_data)} instead."
            raise TierkreisError(msg)
        existing_data.update(data)
        self.write(
            self._node_debug_path(),
            json.dumps(existing_data).encode(),
        )

    def read_debug_data(self, name: str) -> dict[str, Any]:
        existing_data = json.loads(self.read(self._node_debug_path()))
        if not isinstance(existing_data, dict):
            msg = f"Expecting executor data to be dict, found {type(existing_data)} instead."
            raise TierkreisError(msg)
        return existing_data[name]

    def write_executor_data(self, loc: Loc, data: ExecutorDebugData) -> None:
        self.write(self._exec_data_path(loc), data.model_dump_json().encode())

    @overload
    def read_executor_data(self) -> dict[Loc, ExecutorDebugData]: ...

    @overload
    def read_executor_data(self, loc: Loc) -> ExecutorDebugData: ...

    def read_executor_data(
        self, loc: Loc | None = None
    ) -> dict[Loc, ExecutorDebugData] | ExecutorDebugData:
        if loc is None:
            result = {}
            for path in self.list_subpaths(self._exec_data_path(Loc()).parent):
                data = json.loads(self.read(path))
                result[Loc(path.parts[-1])] = ExecutorDebugData(**data)
            return result
        else:
            data = json.loads(self.read(self._exec_data_path(loc)))
            return ExecutorDebugData(**data)

    def dependents(self, loc: Loc) -> set[Loc]:
        """Nodes that are fully invalidated if the node at the given loc is invalidated.

        This does not include the direct parent Loc, which is only partially invalidated.
        """
        descs: set[Loc] = set()
        step, parent = loc.pop_last()
        match step:
            case "-":
                pass
            case ("N", _):
                nodedef = self.read_node_def(loc)
                if nodedef.type == "output":
                    descs.update(self.dependents(parent))
                for output_set in nodedef.outputs.values():
                    for output in output_set:
                        descs.add(parent.N(output))
                        descs.update(self.dependents(parent.N(output)))
            case ("M", _):
                descs.update(self.dependents(parent))
            case ("L", idx):
                latest_idx = self.latest_loop_iteration(parent).peek_index()
                [descs.add(parent.L(i)) for i in range(idx + 1, latest_idx + 1)]
                descs.update(self.dependents(parent))
            case _:
                assert_never(step)

        return descs

    def restart_task(self, loc: Loc) -> list[Loc]:
        """Restart the task/function node at the given loc.

        Fully dependent nodes will be removed from the storage.
        The parent locs will be partially invalidated.

        Returns the invalidated nodes."""

        nodedef = self.read_node_def(loc)
        if nodedef.type != "function":
            raise TierkreisError("Can only restart task/function nodes.")

        # Remove fully invalidated nodes.
        deps = self.dependents(loc)
        [self.delete(self.workflow_dir / a) for a in deps]

        # Mark partially invalidated nodes as not finished and remove their outputs.
        partials = loc.partial_locs()
        [self.delete(self._done_path(x)) for x in partials]
        [self.delete(self.workflow_dir / a / "outputs") for a in partials]

        # Mark given Loc as not started, so that the controller picks it up on the next tick.
        self.delete(self._nodedef_path(loc))

        return list(deps)
