from typing import ClassVar
from typing_extensions import Self
from contextvars import ContextVar


class PageState:
    """
    Returns the current PageState instance for this subclass, creating one if it does not yet exist.
    This ensures each page request has an isolated state container via a ContextVar.

    Args:
        cls (type): The subclass of PageState requesting its context-bound instance.

    Returns:
        Self: The existing or newly created PageState instance associated with the current context.

    Example:
    .. code-block:: python
        class MyState(PageState):
            def __init__(self):
                self.a1 = ui.state("foo")

        # Must call prepare() at the beginning of the page function
        MyState.prepare()

        # Then you can get the state instance anywhere
        my_state = MyState.get()
        value = my_state.a1
    """

    _ctx: ClassVar[ContextVar["PageState"]]
    _prepared: ClassVar[bool] = False

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        cls._ctx = ContextVar(f"{cls.__name__}_ctx")
        cls._prepared = False

    @classmethod
    def prepare(cls) -> Self:
        """
        Prepare the page state instance. Must be called at the beginning of the page function,
        outside of any component scope.

        Returns:
            Self: The newly created PageState instance.

        Raises:
            RuntimeError: If prepare() is called multiple times.
        """
        if cls._prepared:
            raise RuntimeError(
                f"{cls.__name__}.prepare() has already been called. "
                "Do not call prepare() multiple times in the same page."
            )

        inst = cls()
        cls._ctx.set(inst)
        cls._prepared = True
        return inst  # type: ignore

    @classmethod
    def get(cls) -> Self:
        """
        Get the page state instance for the current context.

        Returns:
            Self: The page state instance.

        Raises:
            RuntimeError: If get() is called before prepare().
        """
        if not cls._prepared:
            raise RuntimeError(
                f"Must call {cls.__name__}.prepare() at the beginning of the page function "
                f"before calling {cls.__name__}.get()"
            )

        inst = cls._ctx.get(None)

        # This line should never execute since _prepared guarantees instance exists
        # But kept for robustness
        if inst is None:
            raise RuntimeError(
                f"Internal error: {cls.__name__}._prepared is True but no instance found in context"
            )

        return inst  # type: ignore
