from __future__ import annotations

import warnings
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Self, TypeAlias

from typing_extensions import override

from waku.di import Provider, WithParents, many, object_, scoped
from waku.eventsourcing._generics import resolve_generic_args
from waku.eventsourcing.contracts.aggregate import IDecider
from waku.eventsourcing.contracts.event import IMetadataEnricher
from waku.eventsourcing.decider.repository import DeciderRepository
from waku.eventsourcing.exceptions import DuplicateAggregateNameError, RegistryFrozenError, UpcasterChainError
from waku.eventsourcing.projection.interfaces import ICatchUpProjection, ICheckpointStore, IProjection
from waku.eventsourcing.serialization.interfaces import IEventSerializer, ISnapshotStateSerializer
from waku.eventsourcing.serialization.registry import EventTypeRegistry
from waku.eventsourcing.snapshot.interfaces import ISnapshotStore, ISnapshotStrategy
from waku.eventsourcing.store.in_memory import InMemoryEventStore
from waku.eventsourcing.store.interfaces import IEventStore
from waku.eventsourcing.upcasting.chain import UpcasterChain
from waku.extensions import OnModuleConfigure, OnModuleRegistration
from waku.modules import DynamicModule, ModuleMetadataRegistry, module

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator, Mapping, Sequence

    from waku.cqrs.contracts.notification import INotification
    from waku.eventsourcing.repository import EventSourcedRepository
    from waku.eventsourcing.upcasting.interfaces import IEventUpcaster
    from waku.modules import ModuleMetadata, ModuleType


__all__ = [
    'EventSourcingConfig',
    'EventSourcingExtension',
    'EventSourcingModule',
    'EventType',
    'EventTypeSpec',
]


@dataclass(frozen=True, slots=True)
class EventType:
    event_type: type[INotification]
    name: str | None = field(default=None, kw_only=True)
    aliases: Sequence[str] = field(default=(), kw_only=True)
    version: int = field(default=1, kw_only=True)
    upcasters: Sequence[IEventUpcaster] = field(default=(), kw_only=True)


EventTypeSpec: TypeAlias = 'type[INotification] | EventType'


@dataclass(frozen=True, slots=True, kw_only=True)
class AggregateBinding:
    repository: type[EventSourcedRepository[Any]]
    event_types: Sequence[EventTypeSpec] = ()
    projections: Sequence[type[IProjection]] = ()
    snapshot_strategy: ISnapshotStrategy | None = None


@dataclass(frozen=True, slots=True, kw_only=True)
class DeciderBinding:
    repository: type[DeciderRepository[Any, Any, Any]]
    decider: type[IDecider[Any, Any, Any]]
    event_types: Sequence[EventTypeSpec] = ()
    projections: Sequence[type[IProjection]] = ()
    snapshot_strategy: ISnapshotStrategy | None = None


@dataclass(frozen=True, slots=True, kw_only=True)
class EventSourcingConfig:
    store: type[IEventStore] | Callable[..., IEventStore] | None = None
    event_serializer: type[IEventSerializer] | Callable[..., IEventSerializer] | None = None
    snapshot_store: type[ISnapshotStore] | Callable[..., ISnapshotStore] | None = None
    snapshot_state_serializer: type[ISnapshotStateSerializer] | Callable[..., ISnapshotStateSerializer] | None = None
    checkpoint_store: type[ICheckpointStore] | Callable[..., ICheckpointStore] | None = None
    enrichers: Sequence[type[IMetadataEnricher]] = ()


@dataclass(slots=True)
class EventSourcingRegistry:
    projection_types: list[type[IProjection]] = field(default_factory=list)
    catch_up_projection_types: list[type[ICatchUpProjection]] = field(default_factory=list)
    event_type_bindings: list[EventTypeSpec] = field(default_factory=list)
    _frozen: bool = field(default=False, init=False, repr=False)

    def merge(self, other: EventSourcingRegistry) -> None:
        self._check_not_frozen()
        self.projection_types.extend(other.projection_types)
        self.catch_up_projection_types.extend(other.catch_up_projection_types)
        self.event_type_bindings.extend(other.event_type_bindings)

    def freeze(self) -> None:
        self._frozen = True

    def handler_providers(self) -> Iterator[Provider]:
        if self.projection_types:
            yield many(IProjection, *self.projection_types, collect=False)
        if self.catch_up_projection_types:
            yield many(ICatchUpProjection, *self.catch_up_projection_types, collect=False)

    @staticmethod
    def collector_providers() -> Iterator[Provider]:
        yield many(IProjection, collect=True)
        yield many(ICatchUpProjection, collect=True)

    def _check_not_frozen(self) -> None:
        if self._frozen:
            raise RegistryFrozenError


@module()
class EventSourcingModule:
    @classmethod
    def register(cls, config: EventSourcingConfig | None = None, /) -> DynamicModule:
        config_ = config or EventSourcingConfig()
        providers: list[Provider] = []

        if config_.store is not None:
            providers.append(scoped(IEventStore, config_.store))
        else:
            providers.append(scoped(WithParents[IEventStore], InMemoryEventStore))  # ty:ignore[not-subscriptable]

        if config_.event_serializer is not None:
            providers.append(scoped(IEventSerializer, config_.event_serializer))

        if config_.snapshot_store is not None:
            providers.append(scoped(ISnapshotStore, config_.snapshot_store))

        if config_.snapshot_state_serializer is not None:
            providers.append(scoped(ISnapshotStateSerializer, config_.snapshot_state_serializer))

        if config_.checkpoint_store is not None:
            providers.append(scoped(ICheckpointStore, config_.checkpoint_store))

        providers.append(many(IMetadataEnricher, *config_.enrichers))

        return DynamicModule(
            parent_module=cls,
            providers=providers,
            extensions=[
                EventSourcingRegistryAggregator(has_serializer=config_.event_serializer is not None),
            ],
            is_global=True,
        )


@dataclass
class EventSourcingExtension(OnModuleConfigure):
    _bindings: list[AggregateBinding] = field(default_factory=list, init=False)
    _decider_bindings: list[DeciderBinding] = field(default_factory=list, init=False)
    _registry: EventSourcingRegistry = field(default_factory=EventSourcingRegistry, init=False)

    def bind_aggregate(
        self,
        repository: type[EventSourcedRepository[Any]],
        event_types: Sequence[EventTypeSpec] = (),
        projections: Sequence[type[IProjection]] = (),
        snapshot_strategy: ISnapshotStrategy | None = None,
    ) -> Self:
        self._bindings.append(
            AggregateBinding(
                repository=repository,
                event_types=event_types,
                projections=projections,
                snapshot_strategy=snapshot_strategy,
            )
        )
        self._registry.projection_types.extend(projections)
        self._registry.event_type_bindings.extend(event_types)
        return self

    def bind_decider(
        self,
        repository: type[DeciderRepository[Any, Any, Any]],
        decider: type[IDecider[Any, Any, Any]],
        event_types: Sequence[EventTypeSpec] = (),
        projections: Sequence[type[IProjection]] = (),
        snapshot_strategy: ISnapshotStrategy | None = None,
    ) -> Self:
        self._decider_bindings.append(
            DeciderBinding(
                repository=repository,
                decider=decider,
                event_types=event_types,
                projections=projections,
                snapshot_strategy=snapshot_strategy,
            )
        )
        self._registry.projection_types.extend(projections)
        self._registry.event_type_bindings.extend(event_types)
        return self

    def bind_catch_up_projections(self, projections: Sequence[type[ICatchUpProjection]]) -> Self:
        self._registry.catch_up_projection_types.extend(projections)
        return self

    @property
    def registry(self) -> EventSourcingRegistry:
        return self._registry

    def aggregate_names(self) -> Iterator[tuple[str, type]]:
        for binding in self._bindings:
            yield binding.repository.aggregate_name, binding.repository
        for binding in self._decider_bindings:
            yield binding.repository.aggregate_name, binding.repository

    def on_module_configure(self, metadata: ModuleMetadata) -> None:
        for binding in self._bindings:
            repo_type = binding.repository
            metadata.providers.append(scoped(WithParents[repo_type], repo_type))  # type: ignore[misc,valid-type]
            if binding.snapshot_strategy is not None:
                metadata.providers.append(object_(binding.snapshot_strategy, provided_type=ISnapshotStrategy))

        for binding in self._decider_bindings:
            repo_type = binding.repository
            metadata.providers.append(scoped(WithParents[repo_type], repo_type))  # type: ignore[misc,valid-type]
            decider_iface = self._resolve_decider_interface(repo_type)
            metadata.providers.append(scoped(decider_iface, binding.decider))
            if binding.snapshot_strategy is not None:
                metadata.providers.append(object_(binding.snapshot_strategy, provided_type=ISnapshotStrategy))

    @staticmethod
    def _resolve_decider_interface(repo_type: type[Any]) -> type[Any]:
        args = resolve_generic_args(repo_type, DeciderRepository)
        if args and len(args) == 3:  # noqa: PLR2004
            return IDecider[args[0], args[1], args[2]]  # type: ignore[valid-type]
        return IDecider


class EventSourcingRegistryAggregator(OnModuleRegistration):
    def __init__(self, *, has_serializer: bool = False) -> None:
        self._has_serializer = has_serializer

    @override
    def on_module_registration(
        self,
        registry: ModuleMetadataRegistry,
        owning_module: ModuleType,
        context: Mapping[Any, Any] | None,
    ) -> None:
        aggregated = EventSourcingRegistry()
        all_aggregate_names: dict[str, list[type]] = {}

        for module_type, ext in registry.find_extensions(EventSourcingExtension):
            aggregated.merge(ext.registry)
            for provider in ext.registry.handler_providers():
                registry.add_provider(module_type, provider)
            for name, repo_type in ext.aggregate_names():
                all_aggregate_names.setdefault(name, []).append(repo_type)

        for name, repos in all_aggregate_names.items():
            if len(repos) > 1:
                raise DuplicateAggregateNameError(name, repos)

        for provider in aggregated.collector_providers():
            registry.add_provider(owning_module, provider)

        aggregated.freeze()
        registry.add_provider(owning_module, object_(aggregated))

        event_type_registry, upcaster_chain = self._build_type_registry(aggregated)
        registry.add_provider(owning_module, object_(event_type_registry))
        registry.add_provider(owning_module, object_(upcaster_chain))

        if self._has_serializer and len(event_type_registry) == 0:
            warnings.warn(
                'A serializer is configured but no event types were registered via '
                'bind_aggregate(event_types=[...]). Deserialization will fail at runtime.',
                UserWarning,
                stacklevel=1,
            )

    @classmethod
    def _build_type_registry(cls, aggregated: EventSourcingRegistry) -> tuple[EventTypeRegistry, UpcasterChain]:
        registry = EventTypeRegistry()
        upcasters: dict[str, Sequence[IEventUpcaster]] = {}

        for spec in cls._deduplicate(aggregated.event_type_bindings):
            item = spec if isinstance(spec, EventType) else EventType(spec)

            registry.register(item.event_type, name=item.name, version=item.version)
            for alias in item.aliases:
                registry.add_alias(item.event_type, alias)

            if item.upcasters:
                type_name = item.name or item.event_type.__name__
                cls._validate_upcaster_versions(item.upcasters, type_name, item.version)
                if type_name in upcasters and upcasters[type_name] is not item.upcasters:
                    msg = f'Conflicting upcaster definitions for event type {type_name!r}'
                    raise UpcasterChainError(msg)
                upcasters[type_name] = item.upcasters

        registry.freeze()
        return registry, UpcasterChain({k: list(v) for k, v in upcasters.items()})

    @staticmethod
    def _deduplicate(bindings: Sequence[EventTypeSpec]) -> Iterator[EventTypeSpec]:
        seen: set[int] = set()
        for item in bindings:
            item_id = id(item)
            if item_id not in seen:
                seen.add(item_id)
                yield item

    @staticmethod
    def _validate_upcaster_versions(upcasters: Sequence[IEventUpcaster], type_name: str, version: int) -> None:
        for u in upcasters:
            if u.from_version >= version:
                msg = (
                    f'Upcaster from_version {u.from_version} for event type {type_name!r} '
                    f'must be < event version {version}'
                )
                raise UpcasterChainError(msg)
