# -----------------------------------------------------------
# Astra - WhatsApp Userbot Framework
# Licensed under the MIT License.
# -----------------------------------------------------------

"""
The EventDispatcher is the central brain of Astra's event system.
It routes incoming data from the bridge to appropriate handlers and waiters.
"""

import logging
import asyncio
import uuid
import time
from typing import Any, Callable, Dict, List, Optional, Union, Awaitable
from dataclasses import dataclass, field

from .filters import Criterion
from .context import EventContext

from ..protocol.serializers import DataTransformer

logger = logging.getLogger("Astra.Events")

@dataclass
class EventSubscription:
    """Represents a registered event or command handler."""
    id: str
    event_name: str
    callback: Callable[..., Awaitable[None]]
    criteria: Optional[Criterion] = None
    priority: int = 0
    is_command: bool = False
    command_name: Optional[str] = None

class EventDispatcher:
    """
    Orchestrates the flow of events across the application.
    
    Manages priority-based execution, command parsing, and event interception 
    via middlewares.
    """

    def __init__(self, client):
        self._client = client
        self._subscriptions: Dict[str, List[EventSubscription]] = {}
        self._waiters: Dict[str, List[Dict[str, Any]]] = {}
        self._middlewares: List[Callable[[str, Any], Awaitable[bool]]] = []
        self._on_event_received: Optional[Callable[[], None]] = None
        self._event_count: int = 0

    def add_middleware(self, middleware: Callable[[str, Any], Awaitable[bool]]):
        """Injects a middleware to pre-process or block events."""
        self._middlewares.append(middleware)

    async def wait_for(self, event_name: str, criteria: Optional[Criterion] = None, timeout: Optional[float] = None) -> Any:
        """
        Suspends execution until a specific event arrives.
        """
        future = asyncio.get_event_loop().create_future()
        waiter = {"future": future, "criteria": criteria}
        
        self._waiters.setdefault(event_name, []).append(waiter)
        
        try:
            return await asyncio.wait_for(future, timeout=timeout)
        finally:
            # Ensure the waiter is removed regardless of success or timeout
            if event_name in self._waiters:
                self._waiters[event_name] = [w for w in self._waiters[event_name] if w["future"] is not future]

    def subscribe(
        self, 
        event_name: str, 
        callback: Callable[..., Awaitable[None]], 
        criteria: Optional[Criterion] = None, 
        priority: int = 0,
        is_command: bool = False,
        command_name: Optional[str] = None
    ) -> str:
        """
        Registers a new event handler.
        """
        sub_id = str(uuid.uuid4())
        sub = EventSubscription(
            id=sub_id,
            event_name=event_name,
            callback=callback,
            criteria=criteria,
            priority=priority,
            is_command=is_command,
            command_name=command_name
        )
        
        self._subscriptions.setdefault(event_name, []).append(sub)
        # Higher priority runs first
        self._subscriptions[event_name].sort(key=lambda s: s.priority, reverse=True)
        
        logger.debug(f"Subscribed: {event_name} (id={sub_id[:8]}, priority={priority})")
        return sub_id

    def set_event_callback(self, callback: Callable[[], None]):
        """Registers a callback invoked on every event dispatch (used by SyncEngine)."""
        self._on_event_received = callback

    async def dispatch(self, name: str, payload: Any):
        """
        Dispatches an event through the engine.
        """
        correlation_id = str(uuid.uuid4())[:8]
        self._event_count += 1
        logger.debug(f"Dispatcher received event '{name}' (ID: {correlation_id})")

        # Notify SyncEngine that an event just arrived (resets stall timer)
        if self._on_event_received:
            try:
                self._on_event_received()
            except Exception:
                pass
        
        # 1. Transform raw payload to Astra models
        is_msg = name in {"message", "msg"}
        if is_msg: name = "message"
        is_reaction = name == "reaction"
        
        if (is_msg or is_reaction) and isinstance(payload, dict):
            payload = DataTransformer.to_message(payload, self._client)

        # 2. Contextualize (for Message/Reaction events)
        prefix, cmd, args = (None, None, [])
        if is_msg:
            text = (getattr(payload, 'text', None) or getattr(payload, 'body', None) or "").strip()
            logger.debug(f"[{correlation_id}] Captured text: '{text}'")
            if text and text[0] in "/.!":
                prefix = text[0]
                parts = text[1:].split()
                if parts:
                    cmd = parts[0].lower()
                    # Sanity check: command names are rarely longer than 32 chars
                    if len(cmd) > 32:
                        logger.warning(f"[{correlation_id}] Ignoring suspicious long command name: {cmd[:50]}...")
                        cmd = None
                        parts = []
                    else:
                        args = parts[1:]
            
            logger.debug(f"[{correlation_id}] Parsed Command: {cmd}, Prefix: {prefix}")
            
            # Upgrade payload to EventContext
            payload = EventContext(self._client, payload, cmd, args, prefix)
        elif is_reaction:
            # Wrap reaction in context as well (for consistency)
            payload = EventContext(self._client, payload)

        # 3. Middleware Execution
        for middleware in self._middlewares:
            try:
                if not await middleware(name, payload):
                    logger.debug(f"[{correlation_id}] Event '{name}' intercepted by middleware.")
                    return
            except Exception as e:
                logger.error(f"[{correlation_id}] Middleware error: {e}")

        # 4. Notify Listeners (via Client EventEmitter)
        if hasattr(self._client, "events"):
            logger.debug(f"[{correlation_id}] Emitting '{name}' to event emitter")
            self._client.events.emit(name, payload)

        # 5. Handle Waiters
        await self._process_waiters(name, payload, cmd, args, prefix, correlation_id)

        # 6. Run Registered Subscriptions
        if name in self._subscriptions:
            for sub in self._subscriptions[name]:
                # Command Check
                if sub.is_command and (not is_msg or cmd != sub.command_name):
                    continue
                
                # Criteria Check
                if sub.criteria:
                    if not await sub.criteria.passes(payload):
                        continue
                
                # Run the handler in a new task
                asyncio.create_task(self._execute_safe(sub, payload, cmd, args, prefix, correlation_id))

    # --- Private Execution Logic ---

    async def _process_waiters(self, name, payload, cmd, args, prefix, correlation_id):
        if name not in self._waiters: return
        
        for waiter in self._waiters[name][:]:
            if waiter["future"].done(): continue
            
            if waiter["criteria"]:
                ctx = EventContext(self._client, payload, cmd, args, prefix) if name == "message" else payload
                if not await waiter["criteria"].passes(ctx): continue
            
            waiter["future"].set_result(payload)

    async def _execute_safe(self, sub, payload, cmd, args, prefix, cid):
        start = time.time()
        try:
            if sub.is_command:
                ctx = EventContext(self._client, payload, cmd, args, prefix)
                await sub.callback(ctx)
            else:
                await sub.callback(payload)
            
            ms = (time.time() - start) * 1000
            logger.debug(f"[{cid}] '{sub.callback.__name__}' done in {ms:.2f}ms")
        except Exception as e:
            logger.error(f"[{cid}] Handler '{sub.callback.__name__}' failed: {e}", exc_info=True)
