"""Extension points for fair task queue.

This module includes:
- Event hooks system for task lifecycle events
- Middleware system for cross-cutting concerns
"""

import logging
from datetime import UTC, datetime
from typing import Any, Protocol

from haolib.background.fair.abstract import FairTask

logger = logging.getLogger(__name__)

# ============================================================================
# Event Hooks
# ============================================================================


class TaskEnqueuedEvent:
    """Event fired when a task is enqueued."""

    def __init__(
        self,
        task: FairTask,
        group_id: str,
        timestamp: datetime,
    ) -> None:
        """Initialize event.

        Args:
            task: Enqueued task
            group_id: Group ID
            timestamp: Event timestamp

        """
        self.task = task
        self.group_id = group_id
        self.timestamp = timestamp


class TaskDequeuedEvent:
    """Event fired when a task is dequeued."""

    def __init__(
        self,
        task: FairTask,
        group_id: str,
        timestamp: datetime,
    ) -> None:
        """Initialize event.

        Args:
            task: Dequeued task
            group_id: Group ID
            timestamp: Event timestamp

        """
        self.task = task
        self.group_id = group_id
        self.timestamp = timestamp


class TaskProcessingEvent:
    """Event fired when a task starts processing."""

    def __init__(
        self,
        task: FairTask,
        group_id: str,
        timestamp: datetime,
    ) -> None:
        """Initialize event.

        Args:
            task: Task being processed
            group_id: Group ID
            timestamp: Event timestamp

        """
        self.task = task
        self.group_id = group_id
        self.timestamp = timestamp


class TaskCompletedEvent:
    """Event fired when a task completes successfully."""

    def __init__(
        self,
        task: FairTask,
        group_id: str,
        processing_time: float,
        timestamp: datetime,
    ) -> None:
        """Initialize event.

        Args:
            task: Completed task
            group_id: Group ID
            processing_time: Time taken to process (seconds)
            timestamp: Event timestamp

        """
        self.task = task
        self.group_id = group_id
        self.processing_time = processing_time
        self.timestamp = timestamp


class TaskFailedEvent:
    """Event fired when a task fails."""

    def __init__(
        self,
        task: FairTask,
        group_id: str,
        error: Exception,
        retry_count: int,
        timestamp: datetime,
    ) -> None:
        """Initialize event.

        Args:
            task: Failed task
            group_id: Group ID
            error: Exception that caused failure
            retry_count: Current retry count
            timestamp: Event timestamp

        """
        self.task = task
        self.group_id = group_id
        self.error = error
        self.retry_count = retry_count
        self.timestamp = timestamp


class TaskHook(Protocol):
    """Protocol for task event hooks."""

    async def on_task_enqueued(self, event: TaskEnqueuedEvent) -> None:
        """Called when a task is enqueued.

        Args:
            event: Enqueued event

        """
        ...

    async def on_task_dequeued(self, event: TaskDequeuedEvent) -> None:
        """Called when a task is dequeued.

        Args:
            event: Dequeued event

        """
        ...

    async def on_task_processing(self, event: TaskProcessingEvent) -> None:
        """Called when a task starts processing.

        Args:
            event: Processing event

        """
        ...

    async def on_task_completed(self, event: TaskCompletedEvent) -> None:
        """Called when a task completes successfully.

        Args:
            event: Completed event

        """
        ...

    async def on_task_failed(self, event: TaskFailedEvent) -> None:
        """Called when a task fails.

        Args:
            event: Failed event

        """
        ...


class HookManager:
    """Manages event hooks for fair task queue."""

    def __init__(self) -> None:
        """Initialize hook manager."""
        self._hooks: list[TaskHook] = []

    def register(self, hook: TaskHook) -> None:
        """Register a hook.

        Args:
            hook: Hook to register

        """
        self._hooks.append(hook)

    def unregister(self, hook: TaskHook) -> None:
        """Unregister a hook.

        Args:
            hook: Hook to unregister

        """
        if hook in self._hooks:
            self._hooks.remove(hook)

    async def fire_task_enqueued(self, task: FairTask, group_id: str) -> None:
        """Fire task enqueued event.

        Args:
            task: Enqueued task
            group_id: Group ID

        """
        event = TaskEnqueuedEvent(task=task, group_id=group_id, timestamp=datetime.now(UTC))
        for hook in self._hooks:
            try:
                await hook.on_task_enqueued(event)
            except Exception as e:
                # Log hook errors but don't fail the operation
                logger.warning(
                    f"Hook {hook.__class__.__name__} failed in on_task_enqueued: {e}",
                    exc_info=True,
                    extra={"task_id": task.task_id, "group_id": group_id},
                )

    async def fire_task_dequeued(self, task: FairTask, group_id: str) -> None:
        """Fire task dequeued event.

        Args:
            task: Dequeued task
            group_id: Group ID

        """
        event = TaskDequeuedEvent(task=task, group_id=group_id, timestamp=datetime.now(UTC))
        for hook in self._hooks:
            try:
                await hook.on_task_dequeued(event)
            except Exception as e:
                logger.warning(
                    f"Hook {hook.__class__.__name__} failed in on_task_dequeued: {e}",
                    exc_info=True,
                    extra={"task_id": task.task_id, "group_id": group_id},
                )

    async def fire_task_processing(self, task: FairTask, group_id: str) -> None:
        """Fire task processing event.

        Args:
            task: Task being processed
            group_id: Group ID

        """
        event = TaskProcessingEvent(task=task, group_id=group_id, timestamp=datetime.now(UTC))
        for hook in self._hooks:
            try:
                await hook.on_task_processing(event)
            except Exception as e:
                logger.warning(
                    f"Hook {hook.__class__.__name__} failed in on_task_processing: {e}",
                    exc_info=True,
                    extra={"task_id": task.task_id, "group_id": group_id},
                )

    async def fire_task_completed(self, task: FairTask, group_id: str, processing_time: float) -> None:
        """Fire task completed event.

        Args:
            task: Completed task
            group_id: Group ID
            processing_time: Time taken to process (seconds)

        """
        event = TaskCompletedEvent(
            task=task,
            group_id=group_id,
            processing_time=processing_time,
            timestamp=datetime.now(UTC),
        )
        for hook in self._hooks:
            try:
                await hook.on_task_completed(event)
            except Exception as e:
                logger.warning(
                    f"Hook {hook.__class__.__name__} failed in on_task_completed: {e}",
                    exc_info=True,
                    extra={"task_id": task.task_id, "group_id": group_id},
                )

    async def fire_task_failed(self, task: FairTask, group_id: str, error: Exception, retry_count: int) -> None:
        """Fire task failed event.

        Args:
            task: Failed task
            group_id: Group ID
            error: Exception that caused failure
            retry_count: Current retry count

        """
        event = TaskFailedEvent(
            task=task,
            group_id=group_id,
            error=error,
            retry_count=retry_count,
            timestamp=datetime.now(UTC),
        )
        for hook in self._hooks:
            try:
                await hook.on_task_failed(event)
            except Exception as e:
                logger.warning(
                    f"Hook {hook.__class__.__name__} failed in on_task_failed: {e}",
                    exc_info=True,
                    extra={"task_id": task.task_id, "group_id": group_id},
                )


class LoggingHook:
    """Hook that logs task events with structured logging.

    Provides structured logging with consistent fields for all events.
    Uses the 'extra' parameter for structured logging compatibility.
    """

    def __init__(self, logger: Any | None = None, log_level: int = logging.INFO) -> None:
        """Initialize logging hook.

        Args:
            logger: Logger instance (defaults to standard logging)
            log_level: Log level for non-error events (default: INFO)

        """
        self._logger = logger or logging.getLogger(__name__)
        self._log_level = log_level

    def _log_with_context(
        self,
        message: str,
        level: int,
        task_id: str,
        group_id: str,
        **extra_fields: Any,
    ) -> None:
        """Log with structured context.

        Args:
            message: Log message
            level: Log level
            task_id: Task ID
            group_id: Group ID
            **extra_fields: Additional structured fields

        """
        structured_extra = {
            "task_id": task_id,
            "group_id": group_id,
            "component": "fair_task_queue",
            **extra_fields,
        }
        self._logger.log(level, message, extra=structured_extra)

    async def on_task_enqueued(self, event: TaskEnqueuedEvent) -> None:
        """Log task enqueued event."""
        self._log_with_context(
            "Task enqueued",
            self._log_level,
            event.task.task_id,
            event.group_id,
            priority=event.task.priority,
            scheduled_at=event.task.scheduled_at.isoformat() if event.task.scheduled_at else None,
            expires_at=event.task.expires_at.isoformat() if event.task.expires_at else None,
            idempotency_key=event.task.idempotency_key,
            max_retries=event.task.max_retries,
        )

    async def on_task_dequeued(self, event: TaskDequeuedEvent) -> None:
        """Log task dequeued event."""
        self._log_with_context(
            "Task dequeued",
            logging.DEBUG,
            event.task.task_id,
            event.group_id,
            priority=event.task.priority,
            retry_count=event.task.retry_count,
        )

    async def on_task_processing(self, event: TaskProcessingEvent) -> None:
        """Log task processing event."""
        self._log_with_context(
            "Task processing started",
            self._log_level,
            event.task.task_id,
            event.group_id,
            priority=event.task.priority,
            retry_count=event.task.retry_count,
        )

    async def on_task_completed(self, event: TaskCompletedEvent) -> None:
        """Log task completed event."""
        self._log_with_context(
            "Task completed",
            self._log_level,
            event.task.task_id,
            event.group_id,
            processing_time_seconds=event.processing_time,
            retry_count=event.task.retry_count,
        )

    async def on_task_failed(self, event: TaskFailedEvent) -> None:
        """Log task failed event."""
        self._log_with_context(
            "Task failed",
            logging.ERROR,
            event.task.task_id,
            event.group_id,
            retry_count=event.retry_count,
            max_retries=event.task.max_retries,
            error_type=type(event.error).__name__,
            error_message=str(event.error),
        )
        # Also log exception with traceback
        self._logger.exception(
            f"Task {event.task.task_id} failed after {event.retry_count} retries",
            extra={
                "task_id": event.task.task_id,
                "group_id": event.group_id,
                "component": "fair_task_queue",
                "retry_count": event.retry_count,
                "max_retries": event.task.max_retries,
                "error_type": type(event.error).__name__,
            },
            exc_info=event.error,
        )


class MetricsHook:
    """Hook that records metrics for task queue operations.

    Uses OpenTelemetry metrics if available, otherwise no-op.
    This hook tracks:
    - Tasks enqueued/dequeued
    - Task processing time
    - Task failures
    - Queue lengths
    """

    def __init__(self) -> None:
        """Initialize metrics hook."""
        self._meter = None
        self._enqueued_counter = None
        self._dequeued_counter = None
        self._completed_counter = None
        self._failed_counter = None
        self._processing_time_histogram = None
        self._queue_length_gauge = None

        # Try to get OpenTelemetry meter
        try:
            from opentelemetry import metrics  # noqa: PLC0415

            meter_provider = metrics.get_meter_provider()
            if meter_provider:
                self._meter = meter_provider.get_meter("haolib.background.fair")
                self._enqueued_counter = self._meter.create_counter(
                    "fair_queue_tasks_enqueued",
                    description="Number of tasks enqueued",
                    unit="1",
                )
                self._dequeued_counter = self._meter.create_counter(
                    "fair_queue_tasks_dequeued",
                    description="Number of tasks dequeued",
                    unit="1",
                )
                self._completed_counter = self._meter.create_counter(
                    "fair_queue_tasks_completed",
                    description="Number of tasks completed",
                    unit="1",
                )
                self._failed_counter = self._meter.create_counter(
                    "fair_queue_tasks_failed",
                    description="Number of tasks failed",
                    unit="1",
                )
                self._processing_time_histogram = self._meter.create_histogram(
                    "fair_queue_task_processing_time",
                    description="Time taken to process tasks",
                    unit="s",
                )
                self._queue_length_gauge = self._meter.create_up_down_counter(
                    "fair_queue_length",
                    description="Current queue length",
                    unit="1",
                )
        except ImportError:
            # OpenTelemetry not available, metrics will be no-op
            logger.debug("OpenTelemetry not available, metrics will be no-op")
        except Exception as e:
            # Meter provider not configured, metrics will be no-op
            logger.debug(f"Meter provider not configured, metrics will be no-op: {e}")

    async def on_task_enqueued(self, event: TaskEnqueuedEvent) -> None:
        """Record task enqueued metric."""
        if self._enqueued_counter:
            self._enqueued_counter.add(
                1,
                {
                    "group_id": event.group_id,
                    "priority": str(event.task.priority),
                },
            )
        if self._queue_length_gauge:
            self._queue_length_gauge.add(1, {"group_id": event.group_id})

    async def on_task_dequeued(self, event: TaskDequeuedEvent) -> None:
        """Record task dequeued metric."""
        if self._dequeued_counter:
            self._dequeued_counter.add(
                1,
                {
                    "group_id": event.group_id,
                    "priority": str(event.task.priority),
                },
            )
        if self._queue_length_gauge:
            self._queue_length_gauge.add(-1, {"group_id": event.group_id})

    async def on_task_processing(self, event: TaskProcessingEvent) -> None:
        """Record task processing metric."""
        # Processing start is tracked via dequeued event
        # No metrics needed here

    async def on_task_completed(self, event: TaskCompletedEvent) -> None:
        """Record task completed metric."""
        if self._completed_counter:
            self._completed_counter.add(
                1,
                {
                    "group_id": event.group_id,
                    "priority": str(event.task.priority),
                },
            )
        if self._processing_time_histogram:
            self._processing_time_histogram.record(
                event.processing_time,
                {
                    "group_id": event.group_id,
                    "priority": str(event.task.priority),
                },
            )

    async def on_task_failed(self, event: TaskFailedEvent) -> None:
        """Record task failed metric."""
        if self._failed_counter:
            error_type = type(event.error).__name__
            self._failed_counter.add(
                1,
                {
                    "group_id": event.group_id,
                    "priority": str(event.task.priority),
                    "error_type": error_type,
                    "retry_count": str(event.retry_count),
                },
            )


# ============================================================================
# Middleware
# ============================================================================


class TaskMiddleware(Protocol):
    """Protocol for task middleware.

    Middleware allows adding cross-cutting concerns like logging, metrics,
    validation, transformation, etc. without modifying core queue logic.

    Example:
        ```python
        class LoggingMiddleware:
            async def before_enqueue(self, task: FairTask) -> FairTask:
                logger.info(f"Enqueuing task {task.task_id}")
                return task

            async def after_enqueue(self, task: FairTask) -> None:
                logger.info(f"Task {task.task_id} enqueued")

            async def before_process(self, task: FairTask) -> FairTask:
                logger.info(f"Processing task {task.task_id}")
                return task

            async def after_process(self, task: FairTask, result: Any) -> None:
                logger.info(f"Task {task.task_id} processed with result: {result}")
        ```

    """

    async def before_enqueue(self, task: FairTask) -> FairTask:
        """Called before a task is enqueued.

        Can modify the task before it's enqueued. Must return a FairTask.

        Args:
            task: Task to be enqueued

        Returns:
            Modified task (or same task if no modifications)

        """
        ...

    async def after_enqueue(self, task: FairTask) -> None:
        """Called after a task is successfully enqueued.

        Args:
            task: Enqueued task

        """
        ...

    async def before_process(self, task: FairTask) -> FairTask:
        """Called before a task is processed.

        Can modify the task before it's processed. Must return a FairTask.

        Args:
            task: Task to be processed

        Returns:
            Modified task (or same task if no modifications)

        """
        ...

    async def after_process(self, task: FairTask, result: Any) -> None:
        """Called after a task is successfully processed.

        Args:
            task: Processed task
            result: Task result

        """
        ...

    async def on_error(self, task: FairTask, error: Exception) -> None:
        """Called when a task encounters an error.

        Args:
            task: Task that encountered an error
            error: Exception that occurred

        """
        ...


class MiddlewareManager:
    """Manages middleware chain for fair task queue."""

    def __init__(self) -> None:
        """Initialize middleware manager."""
        self._middleware: list[TaskMiddleware] = []

    def register(self, middleware: TaskMiddleware) -> None:
        """Register a middleware.

        Args:
            middleware: Middleware to register

        """
        self._middleware.append(middleware)

    def unregister(self, middleware: TaskMiddleware) -> None:
        """Unregister a middleware.

        Args:
            middleware: Middleware to unregister

        """
        if middleware in self._middleware:
            self._middleware.remove(middleware)

    async def apply_before_enqueue(self, task: FairTask) -> FairTask:
        """Apply all middleware before enqueue.

        Args:
            task: Task to process

        Returns:
            Modified task after all middleware

        """
        current_task = task
        for middleware in self._middleware:
            current_task = await middleware.before_enqueue(current_task)
        return current_task

    async def apply_after_enqueue(self, task: FairTask) -> None:
        """Apply all middleware after enqueue.

        Args:
            task: Enqueued task

        """
        for middleware in self._middleware:
            await middleware.after_enqueue(task)

    async def apply_before_process(self, task: FairTask) -> FairTask:
        """Apply all middleware before process.

        Args:
            task: Task to process

        Returns:
            Modified task after all middleware

        """
        current_task = task
        for middleware in self._middleware:
            current_task = await middleware.before_process(current_task)
        return current_task

    async def apply_after_process(self, task: FairTask, result: Any) -> None:
        """Apply all middleware after process.

        Args:
            task: Processed task
            result: Task result

        """
        for middleware in self._middleware:
            await middleware.after_process(task, result)

    async def apply_on_error(self, task: FairTask, error: Exception) -> None:
        """Apply all middleware on error.

        Args:
            task: Task that encountered an error
            error: Exception that occurred

        """
        for middleware in self._middleware:
            await middleware.on_error(task, error)

