from __future__ import annotations
from typing import Generator, Type, TypeVar, ClassVar, Any, Iterable
from types import GenericAlias

from pathlib import Path
from contextlib import contextmanager
from dataclasses import asdict, is_dataclass

from tgzr.context import Context, ContextData

#
#   Contexts Definition
#

# We will use contexts to provide the traits of an object.
# To do so, we manage a list of known Traits in a registry.
# Each trait is a ContextData so there's always a current one
# and you can ask to elevate it to a sub-trait by creating
# a context of that sub-trait. If it fails, the object doesn't
# have this trait(s)

# Traits registry
# Traits are actually called "DataType" because "Trait" is way
# too pedantic for me ;)

DataType = TypeVar("DataType", bound="Data")


class DataTypeRegistry:
    """
    The DataTypeRegistry knows all the DataType you want.
    It is able to generate one of its DataType from the
    current context (when compatible).
    """

    def __init__(self):
        self._data_types: dict[str, Type[Data]] = {}

    def register(self, cls: Type[DataType]) -> Type[DataType]:
        if not issubclass(cls, Data):
            raise TypeError("Can only register subclasses of Data.")
        self._data_types[cls.__name__] = cls
        return cls

    def get(self, type_name: str) -> Type[Data] | None:
        return self._data_types.get(type_name)


registry = DataTypeRegistry()


@Context.context_type
class Data(ContextData):
    """
    The Data is the root context for our DataTypes (traits).
    It has a registry as class variable, and automatically fills
    it with all the classes  decorated as `@context_type()`

    It is able to dump a DataType to json and to rebuild one from
    json (making it the current one).

    """

    registry: ClassVar[DataTypeRegistry] = DataTypeRegistry()

    @classmethod
    def context_type(cls, class_to_decorate: Type[DataType]) -> Type[DataType]:
        decorated_class = super(Data, cls).context_type(class_to_decorate)
        cls.registry.register(decorated_class)
        return decorated_class

    def to_dict(self) -> dict[str, Any]:
        data = asdict(self)  # type: ignore - we know it's a dataclass
        data["_type"] = self.__class__.__name__
        return data

    @classmethod
    def as_dataclass(cls, DataClass: Type[DataType], data: dict[str, Any]) -> DataType:
        values = {}
        for f in fields(DataClass):  # type: ignore - it's a DataType we know it's a dataclass type
            if isinstance(f.type, GenericAlias) and f.type.__origin__ == list:
                values[f.name] = [
                    cls.as_dataclass(f.type.__args__[0], d2) for d2 in data[f.name]
                ]
            elif not is_dataclass(data[f.name]):
                values[f.name] = data[f.name]
            else:
                values[f.name] = cls.as_dataclass(
                    f.type,  # type: ignore - it's a dataclass because of the earlier "elif no is_dataclass():"
                    data[f.name],
                )

        return DataClass(**values)

    @classmethod
    @contextmanager
    def from_dict(
        cls: Type[DataType], data: dict[str, Any]
    ) -> Generator[DataType, None, None]:
        data = data.copy()
        type_name = data.pop("_type", None)

        if not type_name:
            raise ValueError("Data missing '_type' field for reconstruction.")

        DataType = cls.registry.get(type_name)
        if not DataType:
            raise ValueError(f"Class '{type_name}' not found in Context registry.")

        with cls.get(cls, **data) as d:
            yield d
        # return cls.as_dataclass(DataType, data)


@Data.context_type
class FSData(Data):
    """
    A filesystem DataType, i.e: a Data with a path.

    Not directly usable, you need to use Folder or File
    instead, or make you own subclass which implements
    the `filenames()->list[path]` method.
    """

    path: str

    def filenames(self) -> list[Path]:
        raise NotImplementedError()


@Data.context_type
class Folder(FSData):
    """
    A filesystem Data, i.e: a Data with a path.
    The `filenames()` method returns all the contained files
    or folders.
    """

    def filenames(self) -> list[Path]:
        return list(Path(self.path).iterdir())


@Data.context_type
class File(FSData):
    """
    A File, i.e: a FSData with a name and an extension
    containing a single path.
    """

    name: str
    extension: str

    def filenames(self) -> list[Path]:
        return [(Path(self.path) / self.name).with_suffix(self.extension)]


@Data.context_type
class Sequence(File):
    """
    An ordered Sequence of named files.
    It contains all the files at `{path}/{name}.[{first}->{last}]{.extension}`
    and is represented as `{path}/{name}.{padding}{.extension}`.

    It can iterate over its frame numbers with `iter_fram_numbers()`.
    It can iterate over its `Frame` contexts with `with iter_frames() as frame:`

    You can also get a specific contained Frame:
        with my_sequence.frame(frame_number) as frame:
            assert isinstance(frame, Frame)

    """

    padding: str
    first: int
    last: int

    def iter_frame_numbers(self) -> Iterable[int]:
        for frame_number in range(self.first, self.last + 1):
            yield frame_number

    @contextmanager
    def frame(self, frame_number: int) -> Generator[Frame, None, None]:
        with self.get(Frame, frame_number=frame_number) as frame:
            yield frame

    def iter_frames(self) -> Iterable[Frame]:
        for i in self.iter_frame_numbers():
            with self.frame(i) as frame:
                yield frame

    def filenames(self) -> list[Path]:
        return sum([frame.filenames() for frame in self.iter_frames()], [])


@Data.context_type
class Frame(Sequence):
    """
    A Frame contains one of the files contained in a Sequence.
    The file is identifier by the `frame_number` field.
    """

    frame_number: str

    def filenames(self) -> list[Path]:
        padding_width = self.padding.count("#") * 4 + self.padding.count("@")
        suffix = f".{str(self.frame_number).zfill(padding_width)}{self.extension}"
        return [(Path(self.path) / self.name).with_suffix(suffix)]


def demo_data_traits_context():

    # We create a Sequence context by providing all its fields:
    with Data.get(
        Sequence,
        path="/path/to",
        name="my_shot",
        extension=".exr",
        first=51,
        last=150,
        padding="@",
    ) as seq:
        # It's a Sequence with exactrly 100 frames:
        assert len(list(seq.iter_frame_numbers())) == 100
        assert len(seq.filenames()) == 100

        # It's the current Sequence:
        assert Sequence.current() == seq

        # We can iterate over the sequence, making
        # each frame the current one as we get it.
        # This give you the opportunity to implement
        # methods which can process a frame without
        # providing the frame as argument. It essentially
        # is a dependency injection, which will make your
        # applications easily configurable and extensible
        # while keeping coupling to a minimum (the Frame
        # API but nothing about how to instanciate it).
        for frame in seq.iter_frames():
            # This Frame is the current Frame
            assert Frame.current() is frame

            # We have completion on frame attributes:
            frame.filenames
            frame.iter_frame_numbers
            frame.iter_frames
            frame.first
            frame.last
            frame.path
            frame.name
            frame.padding
            frame.frame_number

        # we can store the state of the context:
        seq_dict = seq.to_dict()
        # We will now use `seq_dict` to "restore" the
        # frame's context.

    # NB: the current (WIP) implementation needs this,
    # but we're working on removing that need:
    del seq_dict["_type"]

    # This is the sequence context dumped into a dict
    # that you could store as json or similar:
    print("->", seq_dict)

    # We can "restore" the original context with
    # `Data.get_as()` or with `Data.get()`
    # Read on to discover the difference!

    # Restoring with `Data.get_as(**dump)`

    # A Sequence is a FSData so
    # we can create a FSData context from the Sequence's dump:
    with Data.get_as(FSData, **seq_dict) as fsdata:
        # It becomes the current context:
        assert Data.current() is fsdata
        assert FSData.current() is fsdata
        # But the current context is just a FSData
        # and cannot be accessed as File
        try:
            File.current()
        except TypeError:
            pass
        else:
            raise AssertionError()
        # and cannot be accessed as Sequence
        try:
            Sequence.current()
        except TypeError:
            pass
        else:
            raise AssertionError()
        # and cannot be accessed as Frame
        try:
            Frame.current()
        except TypeError:
            pass
        else:
            raise AssertionError()

        # It has the exact requested type:
        assert isinstance(fsdata, FSData)
        # With completion for that type:
        fsdata.path

    # A Sequence is a File so
    # we can create a File context from the Sequence's dump:
    with Data.get_as(File, **seq_dict) as file:
        # It becomes the current context:
        assert Data.current() is file
        assert FSData.current() is file
        assert File.current() is file
        # But the current context is just a File
        # and cannot be accessed as Sequence
        try:
            Sequence.current()
        except TypeError:
            pass
        else:
            raise AssertionError()
        # and cannot be accessed as Frame
        try:
            Frame.current()
        except TypeError:
            pass
        else:
            raise AssertionError()

        # It has the exact requested type:
        assert isinstance(file, File)
        # With completion for that type:
        file.path
        file.filenames

    # If you know it's a sequence
    # you can recreate a Sequence context:
    with Data.get_as(Sequence, **seq_dict) as file:
        # It becomes the current context:
        assert Data.current() is file
        assert FSData.current() is file
        assert File.current() is file
        assert Sequence.current() is file
        # And cannot be accessed as Frame
        # since a Sequence is not a Frame
        try:
            Frame.current()
        except TypeError:
            pass
        else:
            raise AssertionError()

    # Restoring with `Data.get(**dump)`

    # Since it's a Sequence dump, we can simply
    # recreate the context by passing the dump
    # as fields for a Sequence context:
    with Data.get(Sequence, **seq_dict) as sequence:
        # It become the current context:
        assert Sequence.current() is sequence
        assert File.current() is sequence
        assert FSData.current() is sequence
        assert Data.current() is sequence
        # It cannot be accessed as Frame
        # since a sequence is not a Frame:
        try:
            Frame.current()
        except TypeError:
            # TypeError: The current context is a <calls `[...].Sequence'>, not a <class '[...].Frame'>
            pass
        else:
            raise AssertionError()

        # It has the exact requested type:
        assert isinstance(sequence, Sequence)
        # With completion on that type:
        sequence.path
        sequence.name
        sequence.filenames
        sequence.padding
        sequence.extension
        sequence.iter_frame_numbers
        sequence.iter_frames

    # Data.get() will raise TypeError if the requested
    # type is lower than the serialized type
    # (or more precisely, if the requested type can't
    # handle the serialized type attribute):
    try:
        with Data.get(File, **seq_dict) as file:
            pass
    except TypeError:
        # Since the File context does not have a 'padding' field:
        # TypeError: __init__() got an unexpected keyword argument 'padding'
        pass
    else:
        raise AssertionError()

    # Data.get() will raise TypeError if the requested
    # type is high than the serialized type
    # (or more precisely, if the serialized type can't
    # provide all the attribute of the requested type):
    try:
        with Data.get(Frame, **seq_dict) as frame:
            pass
    except TypeError:
        # Since the Sequence dumped context did not have a 'frame_number' field:
        # TypeError: __init__() missing 1 required positional argument: 'frame_number'
        pass
    else:
        raise AssertionError()

    # TODO: implement something like Data.get_highest(dump) returning
    # an instannce of the highest class compatible with dump.
    # -> This could be usefull in GUIs but we wouldn't get code completion on
    # it so it's a low priority.


if __name__ == "__main__":
    demo_data_traits_context()
