from abc import ABC, abstractmethod
from pydantic import BaseModel, ConfigDict, Field
from typing import AsyncIterator, Protocol, runtime_checkable, Any, Self
from types import TracebackType
from asyncio import Lock


class State(BaseModel):
    """
    Holds variables for the nodes of serializable types. Arbitrary types are not allowed.

    The state is copied for each parallel node and merged after the node has finished.
    Therefore, the state is not shared between nodes and operations are safe.

    Parallel changes of the same variable will be detected as a conflict and raise an error.

    Implements pydantic's BaseModel.
    """
    model_config = ConfigDict(arbitrary_types_allowed=False) # for deep copy


class Shared(BaseModel):
    """
    Holds shared variables for the nodes of any type.

    The shared state is shared between all parallel nodes and operations are not safe without using the Lock.
    The lock can be accessed via `shared.lock`.

    Implements pydantic's BaseModel with arbitrary types allowed.
    """
    lock: Lock = Field(default_factory=Lock)
    model_config = ConfigDict(arbitrary_types_allowed=True)


@runtime_checkable
class StateProtocol(Protocol):
    """
    Protocol for the state. Used for type hinting.
    """
    
    def model_copy(self, *, deep: bool = False) -> Self: ...
    def model_dump(self) -> dict[str, Any]: ...
    @classmethod
    def model_validate(cls, obj: Any) -> Self: ...

@runtime_checkable
class SharedProtocol(Protocol):
    """
    Protocol for the shared state. Used for type hinting.
    """

    def model_copy(self, *, deep: bool = False) -> Self: ...
    def model_dump(self) -> dict[str, Any]: ...
    @classmethod
    def model_validate(cls, obj: Any) -> Self: ...

    lock: Lock


class StateAttribute(BaseModel):
    """
    Base class for state attributes. State attributes are used to define variables in the state that can be accessed by the nodes.
    Implements pydantic's BaseModel.
    """
    pass


class SharedAttribute(BaseModel):
    """
    Base class for shared state attributes. Shared state attributes are used to define variables in the shared state that can be accessed by the nodes.
    Implements pydantic's BaseModel.
    """
    pass




class Stream[T: object](ABC, AsyncIterator[T]):
    """
    Standardized wrapper interface for streams of data.

    Set the type of the data that the stream will yield with the generic type parameter `T`.

    Implements an async iterator and async context manager.
    """

    @abstractmethod
    async def aclose(self) -> None:
        pass

    @abstractmethod
    async def __anext__(self) -> T:
        pass

    async def __aenter__(self) -> "Stream[T]":
        return self

    async def __aexit__(
            self, exc_type: type[BaseException] | None, 
            exc: BaseException | None, 
            tb: TracebackType | None
        ) -> None: # Not handling exceptions here -> returns None

        await self.aclose()