import logging
from collections.abc import Callable, Iterable
from typing import Any, overload

logger = logging.getLogger(__name__)

type Updater[T] = Callable[[T], T]
type Callback[T] = Callable[[T], None]
type Unsubscribe = Callable[[], None]
type Selector[T, R] = Callable[[T], R]
type EqualityFn[T] = Callable[[T, T], bool]


class State[T]:
    _value: T

    @overload
    def __init__(
        self,
        initial_value: T,
        *,
        equality_fn: EqualityFn[T] | None = None,
    ) -> None: ...

    @overload
    def __init__(
        self,
        *,
        factory_fn: Callable[[], T],
        dependencies: "list[State[Any]]",
        equality_fn: EqualityFn[T] | None = None,
    ) -> None: ...

    def __init__(
        self,
        initial_value: T | None = None,
        *,
        factory_fn: Callable[[], T] | None = None,
        dependencies: "Iterable[State[Any]] | None" = None,
        equality_fn: EqualityFn[T] | None = None,
    ) -> None:
        self._callbacks: dict[int, Callback[T]] = {}
        self._next_id = 0
        self._equals = equality_fn or (lambda a, b: a == b)

        if initial_value is not None:
            self._value = initial_value

        if factory_fn is not None and dependencies is not None:
            self._value = factory_fn()
            for dependency in dependencies:
                dependency.subscribe(lambda _: self.set(factory_fn()))

    def get(self) -> T:
        return self._value

    def set(self, new_value: T | Updater[T]) -> None:
        computed = new_value(self._value) if callable(new_value) else new_value

        if self._equals(self._value, computed):
            logger.debug("State unchanged, skipping update (%s)")
            return

        logger.debug("State updated, notifying %d callback(s) (%s)", len(self._callbacks))

        self._value = computed

        for callback in list(self._callbacks.values()):
            callback(self._value)

    def subscribe(self, callback: Callback[T]) -> Unsubscribe:
        callback_id = self._next_id
        self._next_id += 1
        self._callbacks[callback_id] = callback

        def unsubscribe() -> None:
            self._callbacks.pop(callback_id, None)

        return unsubscribe

    def derive[R](
        self, selector: Selector[T, R], equality_fn: EqualityFn[R] | None = None
    ) -> "State[R]":
        derived_state = State(selector(self._value), equality_fn=equality_fn)

        def sync(new_parent_value: T) -> None:
            derived_state.set(selector(new_parent_value))

        self.subscribe(sync)

        return derived_state
