"""Field with automatic serialization and deserialization."""
import dataclasses
import json
import logging
import pathlib
import typing

import pandas as pd
import yaml
import znflow
import znflow.utils
import zninit
import znjson
from znflow import handler

from zntrack import exceptions
from zntrack.fields.field import (
    DataIsLazyError,
    Field,
    FieldGroup,
    LazyField,
    PlotsMixin,
)
from zntrack.utils import module_handler, update_key_val

if typing.TYPE_CHECKING:
    from zntrack import Node
log = logging.getLogger(__name__)


class ConnectionConverter(znjson.ConverterBase):
    """Convert a znflow.Connection object to dict and back."""

    level = 100
    representation = "znflow.Connection"
    instance = znflow.Connection

    def encode(self, obj: znflow.Connection) -> dict:
        """Convert the znflow.Connection object to dict."""
        if obj.item is not None:
            raise NotImplementedError("znflow.Connection getitem is not supported yet.")
        return dataclasses.asdict(obj)

    def decode(self, value: str) -> znflow.Connection:
        """Create znflow.Connection object from dict."""
        return znflow.Connection(**value)


class CombinedConnectionsConverter(znjson.ConverterBase):
    """Convert a znflow.Connection object to dict and back."""

    level = 100
    representation = "znflow.CombinedConnections"
    instance = znflow.CombinedConnections

    def encode(self, obj: znflow.CombinedConnections) -> dict:
        """Convert the znflow.Connection object to dict."""
        if obj.item is not None:
            raise NotImplementedError(
                "znflow.CombinedConnections getitem is not supported yet."
            )
        return dataclasses.asdict(obj)

    def decode(self, value: str) -> znflow.CombinedConnections:
        """Create znflow.Connection object from dict."""
        connections = []
        for item in value["connections"]:
            if isinstance(item, dict):
                # @nodify functions aren't support as 'zn.deps'
                # Nodes directly aren't supported because they aren't lists
                connections.append(znflow.Connection(**item))
            else:
                # For the case that item is already a znflow.Connection
                connections.append(item)
        value["connections"] = connections
        return znflow.CombinedConnections(**value)


class SliceConverter(znjson.ConverterBase):
    """Convert a znflow.Connection object to dict and back."""

    level = 100
    representation = "slice"
    instance = slice

    def encode(self, obj: slice) -> dict:
        """Convert the znflow.Connection object to dict."""
        return {"start": obj.start, "stop": obj.stop, "step": obj.step}

    def decode(self, value: dict) -> znflow.Connection:
        """Create znflow.Connection object from dict."""
        return slice(value["start"], value["stop"], value["step"])


znjson.config.register(SliceConverter)


class Params(Field):
    """A parameter field.

    Parameters
    ----------
    dvc_option: str
        The DVC option to use. Default is "params".
    """

    dvc_option: str = "params"
    group = FieldGroup.PARAMETER

    def get_files(self, instance: "Node") -> list:
        """Get the list of files affected by this field.

        Returns
        -------
        list
            A list of file paths.
        """
        return ["params.yaml"]

    def save(self, instance: "Node"):
        """Save the field to disk.

        Parameters
        ----------
        instance : Node
            The node instance associated with this field.
        """
        file = self.get_files(instance)[0]

        try:
            params_dict = yaml.safe_load(pathlib.Path(file).read_text())
        except FileNotFoundError:
            params_dict = {instance.name: {}}

        if instance.name not in params_dict:
            params_dict[instance.name] = {}

        params_dict[instance.name][self.name] = getattr(instance, self.name)
        params_dict = json.loads(json.dumps(params_dict, cls=znjson.ZnEncoder))

        with open(file, "w") as f:
            yaml.safe_dump(params_dict, f, indent=4)

    def get_data(self, instance: "Node") -> any:
        """Get the value of the field from the file."""
        file = self.get_files(instance)[0]
        params_dict = yaml.safe_load(instance.state.fs.read_text(file))
        value = params_dict[instance.name][self.name]
        return json.loads(json.dumps(value), cls=znjson.ZnDecoder)

    def get_stage_add_argument(self, instance: "Node") -> typing.List[tuple]:
        """Get the DVC stage add argument for this field.

        Parameters
        ----------
        instance : Node
            The node instance associated with this field.

        Returns
        -------
        list
            A list of tuples containing the DVC option and the file path.
        """
        file = self.get_files(instance)[0]
        return [(f"--{self.dvc_option}", f"{file}:{instance.name}")]


class Output(LazyField):
    """A field that is saved to disk."""

    group = FieldGroup.RESULT

    def __init__(self, dvc_option: str, **kwargs):
        """Create a new Output field.

        Parameters
        ----------
        dvc_option : str
            The DVC option used to specify the output file.
        **kwargs
            Additional arguments to pass to the parent constructor.
        """
        self.dvc_option = dvc_option
        super().__init__(**kwargs)

    def get_files(self, instance) -> list:
        """Get the path of the file in the node directory.

        Parameters
        ----------
        instance : Node
            The node instance.

        Returns
        -------
        list
            A list containing the path of the file.
        """
        return [instance.nwd / f"{self.name}.json"]

    def save(self, instance: "Node"):
        """Save the field to disk.

        Parameters
        ----------
        instance : Node
            The node instance.
        """
        try:
            value = self.get_value_except_lazy(instance)
        except DataIsLazyError:
            return

        instance.nwd.mkdir(exist_ok=True, parents=True)
        file = self.get_files(instance)[0]
        file.write_text(json.dumps(value, cls=znjson.ZnEncoder, indent=4))

    def get_data(self, instance: "Node") -> any:
        """Get the value of the field from the file."""
        file = self.get_files(instance)[0]
        return json.loads(
            instance.state.fs.read_text(file.as_posix()),
            cls=znjson.ZnDecoder,
        )

    def get_stage_add_argument(self, instance) -> typing.List[tuple]:
        """Get the DVC command for this field.

        Parameters
        ----------
        instance : Node
            The node instance.

        Returns
        -------
        list
            A list containing the DVC command for this field.
        """
        file = self.get_files(instance)[0]
        return [(f"--{self.dvc_option}", file.as_posix())]


class Plots(PlotsMixin, LazyField):
    """A field that is saved to disk."""

    dvc_option: str = "plots"
    group = FieldGroup.RESULT

    def get_files(self, instance) -> list:
        """Get the path of the file in the node directory."""
        return [instance.nwd / f"{self.name}.csv"]

    def save(self, instance: "Node"):
        """Save the field to disk."""
        super().save(instance)
        try:
            value = self.get_value_except_lazy(instance)
        except DataIsLazyError:
            return
        if value is None:
            return

        instance.nwd.mkdir(exist_ok=True, parents=True)
        file = self.get_files(instance)[0]
        value.to_csv(file)

    def get_data(self, instance: "Node") -> any:
        """Get the value of the field from the file."""
        file = self.get_files(instance)[0]
        return pd.read_csv(instance.state.fs.open(file.as_posix()), index_col=0)

    def get_stage_add_argument(self, instance) -> typing.List[tuple]:
        """Get the dvc command for this field."""
        file = self.get_files(instance)[0]
        return [(f"--{self.dvc_option}", file.as_posix())]


_default = object()


def _get_all_connections_and_instances(value) -> list["Node"]:
    """Get Nodes from Connections and CombinedConnections."""
    connections = []
    stack = [value]
    while stack:
        node = stack.pop()
        if isinstance(node, znflow.CombinedConnections):
            stack.extend(node.connections)
        elif isinstance(node, znflow.Connection):
            instance = node.instance
            while isinstance(instance, znflow.Connection):
                instance = instance.instance
            connections.append(instance)
    return connections


class Dependency(LazyField):
    """A dependency field."""

    dvc_option = "deps"
    group = FieldGroup.PARAMETER

    def __init__(self, default=_default):
        """Create a new dependency field.

        A `zn.deps` does not support default values.
        To build a dependency graph, the values must be passed at runtime.
        """
        if default is _default:
            super().__init__()
        elif default is None:
            super().__init__(default=default)
        else:
            raise ValueError(
                "A dependency field does not support default dependencies. You can only"
                " use 'None' to declare this an optional dependency"
                f"and not {default}."
            )

    def get_files(self, instance) -> list:
        """Get the affected files of the respective Nodes."""
        files = []

        value = getattr(instance, self.name)
        # TODO use IterableHandler?

        if isinstance(value, dict):
            value = list(value.values())
        if not isinstance(value, (list, tuple)):
            value = [value]
        if isinstance(value, tuple):
            value = list(value)

        others = []
        for node in value:
            if isinstance(node, (znflow.CombinedConnections, znflow.Connection)):
                others.extend(_get_all_connections_and_instances(node))
            else:
                others.append(node)

        value = others

        for node in value:
            node: Node
            if node is None:
                continue
            if node._external_:
                from zntrack.utils import run_dvc_cmd

                # TODO save these files in a specific directory called `external`
                # TODO the `dvc import cmd` should not run here but rather be a stage?

                deps_file = pathlib.Path("external", f"{node.uuid}.json")
                deps_file.parent.mkdir(exist_ok=True, parents=True)

                # zntrack run node.name --external \
                # --remote node.state.remote --rev node.state.rev

                # when combining with zn.nodes this should be used
                # dvc stage add <name> --params params.yaml:<name>
                # --outs nodes/<name>/node-meta.json zntrack run <name> --external

                cmd = [
                    "import",
                    node.state.remote if node.state.remote is not None else ".",
                    (node.nwd / "node-meta.json").as_posix(),
                    "-o",
                    deps_file.as_posix(),
                ]
                if node.state.rev is not None:
                    cmd.extend(["--rev", node.state.rev])
                # TODO how can we test, that the loaded file truly is the correct one?
                if not deps_file.exists():
                    run_dvc_cmd(cmd)
                files.append(deps_file.as_posix())
                # dvc import node-meta.json + add as dependency file
                continue
            # if node.state.rev is not None or node.state.remote is not None:
            #     # TODO if the Node has a `rev` or `remote` attribute, we need to
            #     #  get the UUID file of the respective Node through node.state.fs.open
            #     # save that somewhere (can't use NWD, because we can now have multiple
            #     # nodes with the same name...)
            #     # and make the uuid a dependency of the node.
            #     continue
            files.append(node.nwd / "node-meta.json")
            for field in zninit.get_descriptors(Field, self=node):
                if field.dvc_option in ["params", "deps"]:
                    # We do not want to depend on parameter files or
                    # recursively on dependencies.
                    continue
                files.extend(field.get_files(node))
                log.debug(f"Found field {field} and extended files to {files}")
        return files

    def save(self, instance: "Node"):
        """Save the field to disk."""
        try:
            value = self.get_value_except_lazy(instance)
        except DataIsLazyError:
            return

        self._write_value_to_config(
            value,
            instance,
            encoder=znjson.ZnEncoder.from_converters(
                [ConnectionConverter, CombinedConnectionsConverter], add_default=True
            ),
        )

    def get_data(self, instance: "Node") -> any:
        """Get the value of the field from the file."""
        zntrack_dict = json.loads(
            instance.state.fs.read_text("zntrack.json"),
        )
        value = zntrack_dict[instance.name][self.name]

        value = update_key_val(value, instance=instance)

        value = json.loads(
            json.dumps(value),
            cls=znjson.ZnDecoder.from_converters(
                [ConnectionConverter, CombinedConnectionsConverter], add_default=True
            ),
        )

        # Up until here we have connection objects. Now we need
        # to resolve them to Nodes. The Nodes, as in 'connection.instance'
        #  are already loaded by the ZnDecoder.
        return handler.UpdateConnectors()(value)

    def get_stage_add_argument(self, instance) -> typing.List[tuple]:
        """Get the dvc command for this field."""
        return [
            (f"--{self.dvc_option}", pathlib.Path(file).as_posix())
            for file in self.get_files(instance)
        ]


class _SaveNodes(znflow.utils.IterableHandler):
    def default(self, value, **kwargs):
        name = kwargs["name"]
        if isinstance(value, znflow.Node):
            value.name = name
            value.save(results=False)
        return value


class NodeField(Dependency):
    """Add another Node as a field.

    The other Node will provide its parameters and methods to be used.
    From a graph standpoint, it will add these parameters and methods to the Node
    it is attached to.
    The Node will not execute its run method or save any results to disk.
    """

    def __set__(self, instance, value):
        """Disable the _graph_ in the value 'Node'."""
        if value is None:
            return super().__set__(instance, value)

        for entry in value if isinstance(value, (list, tuple)) else [value]:
            if hasattr(entry, "_graph_"):
                entry._graph_ = None
                if entry.uuid in instance._graph_:
                    raise exceptions.ZnNodesOnGraphError(
                        node=entry, field=self, instance=instance
                    )
            else:
                raise TypeError(f"The value must be a Node and not {entry}.")
        return super().__set__(instance, value)

    def get_node_names(self, instance) -> list:
        """Get the name of the other Node."""
        try:
            value = self.get_value_except_lazy(instance)
        except DataIsLazyError:
            return []

        if value is None:  # the zn.nodes(None) case
            return []
        if isinstance(value, (list, tuple)):
            return [f"{instance.name}_{self.name}_{idx}" for idx in range(len(value))]
        return [f"{instance.name}_{self.name}"]

    def save(self, instance: "Node"):
        """Save the Node parameters to disk."""
        try:
            value = self.get_value_except_lazy(instance)
        except DataIsLazyError:
            return

        if value is not None:
            if not isinstance(value, (list, tuple)):
                value = [value]

            for node, name in zip(value, self.get_node_names(instance)):
                _SaveNodes()(node, name=name)
        super().save(instance)

    def _get_nwd(self, instance: "Node", name: str) -> pathlib.Path:
        """Get the node working directory."""
        # get the name of the parent directory as string
        # e.g. we have nodes/AL_0/AL_0_ASEMD_checker_list_0
        # but want nodes/AL_0/ASEMD_checker_list_0
        if name.startswith(instance.nwd.parent.name):
            return instance.nwd.parent / name[len(instance.nwd.parent.name) + 1 :]
        else:
            return instance.nwd.parent / name

    def get_optional_dvc_cmd(self, instance: "Node") -> typing.List[list]:
        """Get the dvc command for this field."""
        nodes = getattr(instance, self.name)
        if nodes is None:
            return []

        names = self.get_node_names(instance)
        if not isinstance(nodes, (list, tuple)):
            nodes = [nodes]

        cmd = []

        for node, name in zip(nodes, names):
            if not isinstance(node, znflow.Node):
                raise TypeError(f"The value must be a Node and not {node}.")
            node.name = name
            module = module_handler(node.__class__)

            nwd = self._get_nwd(instance, name)
            node.__dict__["nwd"] = nwd

            _cmd = [
                "stage",
                "add",
                "--name",
                name,
                "--force",
                "--metrics-no-cache",
                (nwd / "node-meta.json").as_posix(),  # HOW DO I MOVE THIS TO GROUP ?
                "--params",
                f"zntrack.json:{instance.name}.{self.name}",
            ]
            field_cmds = []
            for attr in zninit.get_descriptors(Params, self=node):
                field_cmds += attr.get_stage_add_argument(node)

            for field_cmd in set(field_cmds):
                _cmd += list(field_cmd)

            _cmd += [
                f"zntrack run {module}.{node.__class__.__name__} --name"
                f" {name} --meta-only"
            ]

            cmd.append(_cmd)

        return cmd

    def get_files(self, instance: "Node") -> list:
        """Get the files affected by this field."""
        return [
            self._get_nwd(instance, name) / "node-meta.json"
            for name in self.get_node_names(instance)
        ]


def params(*args, **kwargs) -> Params:
    """Create a params field."""
    return Params(*args, **kwargs)


def deps(*args, **kwargs) -> Dependency:
    """Create a dependency field."""
    return Dependency(*args, **kwargs)


def outs() -> Output:
    """Create an output field."""
    return Output(dvc_option="outs", use_repr=False)


def metrics() -> Output:
    """Create a metrics output field."""
    return Output(dvc_option="metrics-no-cache")


def plots(*args, **kwargs) -> Plots:
    """Create a metrics output field."""
    return Plots(*args, **kwargs)


def nodes(*args, **kwargs) -> NodeField:
    """Create a node field."""
    return NodeField(*args, **kwargs)
