"""Base ProcessManager class for stateful cross-domain orchestration.

Process managers coordinate long-running workflows across multiple aggregates.
They maintain their own event-sourced state (keyed by correlation_id) and
react to events from multiple domains.

Router Pattern: ProcessManager follows the MULTI-DOMAIN OO pattern.
- Multiple input domains: use @handles(EventType, input_domain="domain")
- Output domains: use @output_domain method decorator
- Uses @prepares for destination declaration
- Uses @rejected decorator for compensation handling
- Stateful: maintains event-sourced state via correlation_id

Two-phase protocol support:
    1. Prepare: Declare additional destinations needed (via @prepares(E))
    2. Handle: Produce commands and process events (via @handles(E, input_domain="x"))

Design Philosophy:
    PMs coordinate, aggregates decide. PMs should NOT rebuild destination
    state to make decisions. They receive only destination sequences for
    command stamping. Use sync mode for immediate feedback on command results.

Example usage:
    from angzarr_client.process_manager import (
        ProcessManager, applies, handles, output_domain, prepares, rejected
    )
    from angzarr_client.destinations import Destinations

    @dataclass
    class OrderWorkflowState:
        order_id: str = ""
        inventory_reserved: bool = False
        payment_received: bool = False

    class OrderWorkflowPM(ProcessManager[OrderWorkflowState]):
        name = "order-workflow"

        def _create_empty_state(self) -> OrderWorkflowState:
            return OrderWorkflowState()

        @applies(OrderCreated)
        def apply_order_created(self, state: OrderWorkflowState, event: OrderCreated):
            state.order_id = event.order_id

        @prepares(OrderCreated)
        def prepare_order(self, event: OrderCreated) -> list[Cover]:
            return [Cover(domain="inventory", root=...)]

        @output_domain("inventory")
        @handles(OrderCreated, input_domain="order")
        def on_order_created(
            self, event: OrderCreated, destinations: Destinations
        ) -> ReserveInventory:
            cmd = ReserveInventory(order_id=event.order_id, quantity=event.quantity)
            # Use destinations.stamp_command() to set correct sequence
            return cmd

        @output_domain("payment")
        @handles(InventoryReserved, input_domain="inventory")
        def on_inventory_reserved(self, event: InventoryReserved) -> ProcessPayment:
            return ProcessPayment(...)

"""

from __future__ import annotations

import inspect
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Generic, TypeVar

from google.protobuf.any_pb2 import Any

from .aggregate import applies
from .compensation import RejectionHandlerResponse
from .destinations import Destinations
from .proto.angzarr import process_manager_pb2 as pm_pb2
from .proto.angzarr import types_pb2 as types
from .router import _pack_any, domain, handles, output_domain, prepares, rejected

if TYPE_CHECKING:
    from .state_builder import StateRouter

# Re-export decorators
__all__ = [
    "ProcessManager",
    "applies",
    "domain",
    "handles",
    "output_domain",
    "prepares",
    "rejected",
]

StateT = TypeVar("StateT")


class ProcessManager(Generic[StateT], ABC):
    """Base class for stateful process managers.

    Router Pattern: Follows the MULTI-DOMAIN OO pattern.

    ProcessManager-specific notes:
    - Multi-domain: uses @handles(EventType, input_domain="domain") syntax
    - Output domains: uses @output_domain method decorator
    - Stateful: maintains event-sourced state keyed by correlation_id
    - Produces commands: can emit commands to any domain

    Process managers are event-sourced aggregates that:
    - React to events from multiple domains
    - Maintain state keyed by correlation_id
    - Produce commands to coordinate workflows
    - Handle rejection/compensation via @rejected handlers
    - Support two-phase protocol with @prepares(E) for destination declaration

    Subclasses must:
    - Set `name` class attribute
    - Implement `_create_empty_state() -> StateT`
    - Decorate event appliers with `@applies(EventType)` OR override `_apply_event()`
    - Decorate event handlers with `@handles(EventType, input_domain="...")`
    - Use `@output_domain("target")` to specify where commands go
    - Optionally decorate prepare handlers with `@prepares(EventType)`
    - Optionally decorate rejection handlers with `@rejected(domain, command)`

    Usage:
        class OrderWorkflowPM(ProcessManager[OrderWorkflowState]):
            name = "order-workflow"

            def _create_empty_state(self) -> OrderWorkflowState:
                return OrderWorkflowState()

            @applies(OrderCreated)
            def apply_order_created(self, state: OrderWorkflowState, event: OrderCreated):
                state.order_id = event.order_id

            @prepares(OrderCreated)
            def prepare_order(self, event: OrderCreated) -> list[Cover]:
                return [Cover(domain="inventory", root=...)]

            @output_domain("inventory")
            @handles(OrderCreated, input_domain="order")
            def on_order_created(
                self, event: OrderCreated, destinations: list[EventBook]
            ) -> ReserveInventory:
                ...

            @rejected(domain="inventory", command="ReserveInventory")
            def handle_reserve_rejected(self, revoke_cmd) -> WorkflowFailed:
                return WorkflowFailed(reason=revoke_cmd.rejection_reason)

    """

    name: str
    _dispatch_table: dict[str, tuple[str, type, str, str]] = (
        {}
    )  # suffix -> (method, type, input_domain, output_domain)
    _prepare_table: dict[str, tuple[str, type]] = {}  # suffix -> (method, type)
    _input_domains: dict[str, list[str]] = {}  # domain -> [event types]
    _rejection_table: dict[str, str] = {}  # "domain/command" -> method_name
    _applier_table: dict[str, tuple[str, type]] = (
        {}
    )  # suffix -> (method_name, event_type)

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)

        # Skip validation for abstract intermediate classes
        if inspect.isabstract(cls):
            return

        # Validate required class attributes
        if not getattr(cls, "name", None):
            raise TypeError(f"{cls.__name__} must define 'name' class attribute")

        cls._dispatch_table = {}
        cls._prepare_table = {}
        cls._input_domains = {}
        cls._rejection_table = {}
        cls._applier_table = {}
        cls._build_dispatch_table()
        cls._build_prepare_table()
        cls._build_rejection_table()
        cls._build_applier_table()

    @classmethod
    def _build_dispatch_table(cls):
        """Scan for @handles methods and build dispatch table."""
        for attr_name in dir(cls):
            attr = getattr(cls, attr_name, None)
            if callable(attr) and getattr(attr, "_is_handler", False):
                event_type = attr._event_type
                input_domain = attr._input_domain
                output_domain = attr._output_domain
                suffix = event_type.__name__

                if suffix in cls._dispatch_table:
                    raise TypeError(f"{cls.__name__}: duplicate handler for {suffix}")

                cls._dispatch_table[suffix] = (
                    attr_name,
                    event_type,
                    input_domain,
                    output_domain,
                )

                # Track input domains for descriptor
                if input_domain:
                    if input_domain not in cls._input_domains:
                        cls._input_domains[input_domain] = []
                    cls._input_domains[input_domain].append(suffix)

    @classmethod
    def _build_prepare_table(cls):
        """Scan for @prepares methods and build prepare table."""
        for attr_name in dir(cls):
            attr = getattr(cls, attr_name, None)
            if callable(attr) and getattr(attr, "_is_prepare_handler", False):
                event_type = attr._event_type
                suffix = event_type.__name__
                if suffix in cls._prepare_table:
                    raise TypeError(
                        f"{cls.__name__}: duplicate prepare handler for {suffix}"
                    )
                cls._prepare_table[suffix] = (attr_name, event_type)

    @classmethod
    def _build_rejection_table(cls):
        """Scan for @rejected methods and build rejection dispatch table."""
        for attr_name in dir(cls):
            attr = getattr(cls, attr_name, None)
            if callable(attr) and getattr(attr, "_is_rejection_handler", False):
                domain = attr._rejection_domain
                command = attr._rejection_command
                key = f"{domain}/{command}"
                if key in cls._rejection_table:
                    raise TypeError(
                        f"{cls.__name__}: duplicate rejection handler for {key}"
                    )
                cls._rejection_table[key] = attr_name

    @classmethod
    def _build_applier_table(cls):
        """Scan for @applies methods and build applier dispatch table.

        Returns:
            Dict mapping event type suffix to (method_name, event_type).
        """
        for attr_name in dir(cls):
            attr = getattr(cls, attr_name, None)
            if callable(attr) and getattr(attr, "_is_applier", False):
                event_type = attr._event_type
                suffix = event_type.__name__
                if suffix in cls._applier_table:
                    raise TypeError(f"{cls.__name__}: duplicate applier for {suffix}")
                cls._applier_table[suffix] = (attr_name, event_type)

    def _create_destinations(
        self,
        destination_sequences: dict[str, int],
    ) -> Destinations:
        """Create Destinations context from sequence map.

        Args:
            destination_sequences: Dict mapping domain name to next sequence.

        Returns:
            Destinations object for command stamping.
        """
        return Destinations(destination_sequences)

    def __init__(self, process_state: types.EventBook = None):
        """Initialize process manager with optional prior state.

        Args:
            process_state: Existing event book for this correlation_id.
        """
        if process_state is None:
            process_state = types.EventBook()
        self._event_book = process_state
        self._state: StateT = None
        self._new_events: list[Any] = []
        self._facts: list[types.EventBook] = []

    def emit_fact(self, fact: types.EventBook) -> None:
        """Emit a fact to inject to another aggregate.

        Args:
            fact: EventBook containing the event to inject.
                  The Cover should specify target domain and root.
        """
        self._facts.append(fact)

    def _get_state(self) -> StateT:
        """Get current state, rebuilding from events if needed."""
        if self._state is None:
            self._state = self._rebuild()
        return self._state

    def _rebuild(self) -> StateT:
        """Rebuild state from event book."""
        state = self._create_empty_state()
        for page in self._event_book.pages:
            if page.event:
                self._apply_event(state, page.event)
        return state

    def prepare(self, event_any: Any) -> list[types.Cover]:
        """Prepare destinations for an event.

        Dispatches to @prepares decorated method if one exists.

        Args:
            event_any: Packed event as google.protobuf.Any

        Returns:
            List of Covers identifying destination aggregates.
        """
        type_url = event_any.type_url

        for suffix, (method_name, event_type) in self._prepare_table.items():
            # Match exact type name (preceded by . or /) to avoid false matches
            if type_url.endswith(f".{suffix}") or type_url.endswith(f"/{suffix}"):
                # Unpack event
                event = event_type()
                event_any.Unpack(event)

                # Call prepare handler
                result = getattr(self, method_name)(event)
                return result if result else []

        return []

    def dispatch(
        self,
        event_any: Any,
        root: bytes = None,
        correlation_id: str = "",
        destination_sequences: dict[str, int] = None,
    ) -> list[types.CommandBook]:
        """Dispatch event to matching handler.

        Args:
            event_any: Packed event as google.protobuf.Any
            root: Source aggregate root
            correlation_id: Correlation ID for the workflow
            destination_sequences: Map of domain to next sequence number

        Returns:
            List of CommandBooks to send.
        """
        type_url = event_any.type_url

        for suffix, (
            method_name,
            event_type,
            _,
            handler_domain,
        ) in self._dispatch_table.items():
            # Match exact type name (preceded by . or /) to avoid false matches
            if type_url.endswith(f".{suffix}") or type_url.endswith(f"/{suffix}"):
                # Unpack event
                event = event_type()
                event_any.Unpack(event)

                # Check which parameters handler accepts
                method = getattr(self, method_name)
                sig = inspect.signature(method)
                params = list(sig.parameters.keys())

                # Build kwargs based on what handler accepts
                kwargs = {}
                if "destinations" in params:
                    # Pass Destinations object for command stamping
                    kwargs["destinations"] = self._create_destinations(
                        destination_sequences or {}
                    )
                if "root" in params:
                    kwargs["root"] = root

                # Call handler with appropriate parameters
                result = method(event, **kwargs)

                # Pack result into CommandBooks
                return self._pack_commands(result, handler_domain, root, correlation_id)

        return []

    def _pack_commands(
        self,
        result,
        output_domain: str = None,
        root: bytes = None,
        correlation_id: str = "",
    ) -> list[types.CommandBook]:
        """Pack command(s) into CommandBooks."""
        if result is None:
            return []

        # Handle pre-packed CommandBooks (advanced usage)
        if isinstance(result, types.CommandBook):
            return [result]
        if (
            isinstance(result, list)
            and result
            and isinstance(result[0], types.CommandBook)
        ):
            return result

        commands = result if isinstance(result, tuple) else (result,)
        books = []

        for cmd in commands:
            cmd_any = _pack_any(cmd)
            # Use handler's output_domain if specified
            domain = output_domain or ""
            cover = types.Cover(
                domain=domain,
                correlation_id=correlation_id,
            )
            if root:
                cover.root.value = root

            book = types.CommandBook(
                cover=cover,
                pages=[types.CommandPage(command=cmd_any)],
            )
            books.append(book)

        return books

    def _apply_and_record(self, event) -> None:
        """Pack event, apply to cached state, record for persistence."""
        event_any = _pack_any(event)

        # Apply directly to cached state
        if self._state is not None:
            self._apply_event(self._state, event_any)

        self._new_events.append(event_any)

    def process_events(self) -> types.EventBook:
        """Return new process events for persistence."""
        pages = [types.EventPage(event=e) for e in self._new_events]
        return types.EventBook(pages=pages)

    @classmethod
    def replay_state(cls, process_state: types.EventBook) -> Any:
        """Replay events to rebuild typed state packed in Any.

        Called by framework before Prepare/Handle to convert EventBook to typed state.
        This matches the aggregate pattern where framework handles state reconstruction.

        Default implementation returns empty Any. Subclasses that need to pass
        typed state to handlers should override this method to pack their state
        into a protobuf Any.

        Args:
            process_state: PM events to replay.

        Returns:
            Typed PM state packed in google.protobuf.Any.
        """
        # Default: return empty Any
        # Subclasses can override to pack their state into a protobuf message
        return Any()

    @classmethod
    def prepare_destinations(
        cls,
        trigger: types.EventBook,
        process_state: types.EventBook,
    ) -> list[types.Cover]:
        """Phase 1: Declare additional destinations needed.

        Args:
            trigger: The triggering event book.
            process_state: Current process manager state.

        Returns:
            List of Covers identifying destination aggregates to fetch.
        """
        pm = cls(process_state)
        destinations: list[types.Cover] = []

        for page in trigger.pages:
            if page.HasField("event"):
                destinations.extend(pm.prepare(page.event))

        return destinations

    @classmethod
    def handle(
        cls,
        trigger: types.EventBook,
        process_state: types.EventBook,
        destination_sequences: dict[str, int] = None,
    ) -> pm_pb2.ProcessManagerHandleResponse:
        """Phase 2: Handle a trigger event with current process state.

        Args:
            trigger: The triggering event book.
            process_state: Current process manager state.
            destination_sequences: Map of domain to next sequence number for stamping.

        Returns:
            ProcessManagerHandleResponse with process_events, commands, and facts.
        """
        pm = cls(process_state)
        root = trigger.cover.root.value if trigger.HasField("cover") else None
        correlation_id = (
            trigger.cover.correlation_id if trigger.HasField("cover") else ""
        )

        commands = []
        for page in trigger.pages:
            if page.HasField("event"):
                commands.extend(
                    pm.dispatch(page.event, root, correlation_id, destination_sequences)
                )

        return pm_pb2.ProcessManagerHandleResponse(
            process_events=pm.process_events(),
            commands=commands,
            facts=pm._facts,
        )

    def handle_revocation(
        self,
        notification: types.Notification,
    ) -> RejectionHandlerResponse:
        """Handle a rejection notification for commands this PM issued.

        Called when a command produced by this PM is rejected by the target aggregate.
        Dispatches to @rejected decorated methods based on target domain
        and rejected command type.

        If no matching @rejected handler is found, returns empty response
        (framework handles default behavior).

        Args:
            notification: Notification containing RejectionNotification payload.

        Returns:
            RejectionHandlerResponse with optional events (compensation)
            and/or notification (upstream propagation).

        Usage:
            @rejected(domain="inventory", command="ReserveInventory")
            def handle_reserve_rejected(self, notification) -> RejectionHandlerResponse:
                ctx = CompensationContext.from_notification(notification)
                self._apply_and_record(WorkflowFailed(reason=ctx.rejection_reason))
                return RejectionHandlerResponse(events=self.process_events())
        """
        # Unpack rejection details from notification payload
        rejection = types.RejectionNotification()
        if notification.HasField("payload"):
            notification.payload.Unpack(rejection)

        # Extract domain and command type from rejected_command
        domain = ""
        command_suffix = ""

        if rejection.HasField("rejected_command") and rejection.rejected_command.pages:
            rejected_cmd = rejection.rejected_command
            if rejected_cmd.HasField("cover"):
                domain = rejected_cmd.cover.domain
            if rejected_cmd.pages[0].HasField("command"):
                type_url = rejected_cmd.pages[0].command.type_url
                # Extract suffix (e.g., "ReserveInventory" from "type.googleapis.com/.../ReserveInventory")
                command_suffix = (
                    type_url.rsplit("/", 1)[-1] if "/" in type_url else type_url
                )

        # Dispatch to @rejected handler if found (use suffix matching like regular dispatch)
        for key, method_name in self._rejection_table.items():
            expected_domain, expected_command = key.split("/", 1)
            if domain == expected_domain and command_suffix.endswith(expected_command):
                # Ensure state is built before calling handler
                _ = self._get_state()
                # Call the handler
                result = getattr(self, method_name)(notification)

                # Handler may return RejectionHandlerResponse directly
                if isinstance(result, RejectionHandlerResponse):
                    return result

                # Handler may return None or a single event
                process_events = self.process_events()
                return RejectionHandlerResponse(
                    events=process_events if process_events.pages else None,
                )

        # Default: no handler, return empty response
        return RejectionHandlerResponse()

    @property
    def state(self) -> StateT:
        """Get current state (convenience property for _get_state)."""
        return self._get_state()

    @abstractmethod
    def _create_empty_state(self) -> StateT:
        """Create an empty state instance. Must be implemented by subclasses."""
        ...

    def _apply_event(self, state: StateT, event_any: Any) -> None:
        """Apply a single event to state.

        If @applies decorated methods exist, dispatches to them automatically.
        Otherwise, subclasses must override this method.

        Args:
            state: Current state to mutate.
            event_any: Packed event as google.protobuf.Any.
        """
        if not self._applier_table:
            raise NotImplementedError(
                f"{self.__class__.__name__} must either define @applies methods "
                "or override _apply_event()"
            )

        type_url = event_any.type_url
        for suffix, (method_name, event_type) in self._applier_table.items():
            # Match exact type name (preceded by . or /) to avoid false matches
            if type_url.endswith(f".{suffix}") or type_url.endswith(f"/{suffix}"):
                event = event_type()
                event_any.Unpack(event)
                getattr(self, method_name)(state, event)
                return

        # Unknown event type - silently ignore (forward compatibility)
