"""Observability support for fair task queue.

This module includes:
- Prometheus metrics export
- Structured logging hooks
- Distributed tracing hooks
"""

import asyncio
import concurrent.futures
import json
import logging
import traceback
import warnings
from datetime import UTC, datetime
from typing import Any

from haolib.background.fair.abstract import QueueMetrics
from haolib.background.fair.extensions import (
    TaskCompletedEvent,
    TaskDequeuedEvent,
    TaskEnqueuedEvent,
    TaskFailedEvent,
    TaskProcessingEvent,
)
from haolib.background.fair.queue import FairTaskQueue

# Try to import OpenTelemetry for tracing
try:
    from opentelemetry import trace
except ImportError:
    trace = None  # type: ignore[assignment, misc]

logger = logging.getLogger(__name__)

# ============================================================================
# Prometheus Metrics Export
# ============================================================================


def export_prometheus_metrics(queue: FairTaskQueue, include_help: bool = True) -> str:  # noqa: PLR0915
    """Export queue metrics in Prometheus format.

    This function generates Prometheus-compatible metrics that can be exposed
    via an HTTP endpoint for Prometheus scraping.

    Example:
        ```python
        from fastapi import FastAPI
        from haolib.background.fair import export_prometheus_metrics

        app = FastAPI()

        @app.get("/metrics")
        async def metrics():
            return export_prometheus_metrics(queue)
        ```

    Args:
        queue: FairTaskQueue instance
        include_help: If True, include HELP comments in output

    Returns:
        Prometheus-formatted metrics string

    """
    # Get metrics synchronously (assuming we're in an async context)
    # This is a best-effort approach - prefer using export_prometheus_metrics_async() in async contexts
    try:
        asyncio.get_running_loop()
        # If we're in a running loop, we can't use asyncio.run()
        # In this case, we need to use the async version or create a new thread

        async def get_metrics() -> QueueMetrics:
            return await queue.get_queue_metrics()

        with concurrent.futures.ThreadPoolExecutor() as executor:
            future = executor.submit(asyncio.run, get_metrics())
            metrics_dict = future.result(timeout=30.0)  # 30 second timeout
    except RuntimeError:
        # No event loop, create a new one
        async def get_metrics() -> QueueMetrics:
            return await queue.get_queue_metrics()

        metrics_dict = asyncio.run(get_metrics())
    except Exception:
        # Fallback: try to get metrics in current context
        # This may fail if we're in an async context without proper handling
        warnings.warn(
            "export_prometheus_metrics() called in async context. Use export_prometheus_metrics_async() instead.",
            UserWarning,
            stacklevel=2,
        )

        # Try one more time with a new event loop
        async def get_metrics() -> QueueMetrics:
            return await queue.get_queue_metrics()

        metrics_dict = asyncio.run(get_metrics())

    lines: list[str] = []

    if include_help:
        lines.append("# HELP fair_queue_total_groups Number of active groups in the queue")
        lines.append("# TYPE fair_queue_total_groups gauge")
        lines.append("# HELP fair_queue_total_queue_length Total number of pending tasks")
        lines.append("# TYPE fair_queue_total_queue_length gauge")
        lines.append("# HELP fair_queue_total_processing Number of tasks currently processing")
        lines.append("# TYPE fair_queue_total_processing gauge")
        lines.append("# HELP fair_queue_total_completed_sample Sample count of completed tasks")
        lines.append("# TYPE fair_queue_total_completed_sample gauge")
        lines.append("# HELP fair_queue_total_failed_sample Sample count of failed tasks")
        lines.append("# TYPE fair_queue_total_failed_sample gauge")
        lines.append("# HELP fair_queue_is_paused Whether the queue is paused (1) or active (0)")
        lines.append("# TYPE fair_queue_is_paused gauge")
        lines.append("# HELP fair_queue_scheduler_active_tasks Number of active tasks in scheduler")
        lines.append("# TYPE fair_queue_scheduler_active_tasks gauge")
        lines.append("# HELP fair_queue_scheduler_max_concurrent_tasks Maximum concurrent tasks in scheduler")
        lines.append("# TYPE fair_queue_scheduler_max_concurrent_tasks gauge")
        lines.append("# HELP fair_queue_scheduler_healthy Whether scheduler is healthy (1) or not (0)")
        lines.append("# TYPE fair_queue_scheduler_healthy gauge")
        lines.append("# HELP fair_queue_group_queue_length Queue length for a specific group")
        lines.append("# TYPE fair_queue_group_queue_length gauge")
        lines.append("# HELP fair_queue_group_weight Weight for a specific group")
        lines.append("# TYPE fair_queue_group_weight gauge")

    # Global metrics
    total_groups = metrics_dict.total_groups
    total_queue_length = metrics_dict.total_queue_length
    total_processing = metrics_dict.total_processing
    total_completed_sample = metrics_dict.total_completed_sample
    total_failed_sample = metrics_dict.total_failed_sample
    is_paused = 1 if metrics_dict.is_paused else 0

    lines.append(f"fair_queue_total_groups {total_groups}")
    lines.append(f"fair_queue_total_queue_length {total_queue_length}")
    lines.append(f"fair_queue_total_processing {total_processing}")
    lines.append(f"fair_queue_total_completed_sample {total_completed_sample}")
    lines.append(f"fair_queue_total_failed_sample {total_failed_sample}")
    lines.append(f"fair_queue_is_paused {is_paused}")

    # Scheduler metrics
    scheduler = metrics_dict.scheduler
    scheduler_active_tasks = scheduler.active_tasks
    scheduler_max_concurrent_tasks = scheduler.max_concurrent_tasks
    scheduler_healthy = 1 if scheduler.healthy else 0

    lines.append(f"fair_queue_scheduler_active_tasks {scheduler_active_tasks}")
    lines.append(f"fair_queue_scheduler_max_concurrent_tasks {scheduler_max_concurrent_tasks}")
    lines.append(f"fair_queue_scheduler_healthy {scheduler_healthy}")

    # Per-group metrics
    groups = metrics_dict.groups.metrics
    for group_id, group_data in groups.items():
        # Escape group_id for Prometheus label (replace special characters)
        safe_group_id = group_id.replace('"', '\\"').replace("\n", "\\n").replace("\\", "\\\\")
        queue_length = group_data.queue_length
        weight = group_data.weight

        lines.append(f'fair_queue_group_queue_length{{group_id="{safe_group_id}"}} {queue_length}')
        lines.append(f'fair_queue_group_weight{{group_id="{safe_group_id}"}} {weight}')

    return "\n".join(lines) + "\n"


async def export_prometheus_metrics_async(queue: FairTaskQueue, include_help: bool = True) -> str:  # noqa: PLR0915
    """Export queue metrics in Prometheus format (async version).

    Args:
        queue: FairTaskQueue instance
        include_help: If True, include HELP comments in output

    Returns:
        Prometheus-formatted metrics string

    """
    metrics_dict: QueueMetrics = await queue.get_queue_metrics()

    lines: list[str] = []

    if include_help:
        lines.append("# HELP fair_queue_total_groups Number of active groups in the queue")
        lines.append("# TYPE fair_queue_total_groups gauge")
        lines.append("# HELP fair_queue_total_queue_length Total number of pending tasks")
        lines.append("# TYPE fair_queue_total_queue_length gauge")
        lines.append("# HELP fair_queue_total_processing Number of tasks currently processing")
        lines.append("# TYPE fair_queue_total_processing gauge")
        lines.append("# HELP fair_queue_total_completed_sample Sample count of completed tasks")
        lines.append("# TYPE fair_queue_total_completed_sample gauge")
        lines.append("# HELP fair_queue_total_failed_sample Sample count of failed tasks")
        lines.append("# TYPE fair_queue_total_failed_sample gauge")
        lines.append("# HELP fair_queue_is_paused Whether the queue is paused (1) or active (0)")
        lines.append("# TYPE fair_queue_is_paused gauge")
        lines.append("# HELP fair_queue_scheduler_active_tasks Number of active tasks in scheduler")
        lines.append("# TYPE fair_queue_scheduler_active_tasks gauge")
        lines.append("# HELP fair_queue_scheduler_max_concurrent_tasks Maximum concurrent tasks in scheduler")
        lines.append("# TYPE fair_queue_scheduler_max_concurrent_tasks gauge")
        lines.append("# HELP fair_queue_scheduler_healthy Whether scheduler is healthy (1) or not (0)")
        lines.append("# TYPE fair_queue_scheduler_healthy gauge")
        lines.append("# HELP fair_queue_group_queue_length Queue length for a specific group")
        lines.append("# TYPE fair_queue_group_queue_length gauge")
        lines.append("# HELP fair_queue_group_weight Weight for a specific group")
        lines.append("# TYPE fair_queue_group_weight gauge")

    # Global metrics
    total_groups = metrics_dict.total_groups
    total_queue_length = metrics_dict.total_queue_length
    total_processing = metrics_dict.total_processing
    total_completed_sample = metrics_dict.total_completed_sample
    total_failed_sample = metrics_dict.total_failed_sample
    is_paused = 1 if metrics_dict.is_paused else 0

    lines.append(f"fair_queue_total_groups {total_groups}")
    lines.append(f"fair_queue_total_queue_length {total_queue_length}")
    lines.append(f"fair_queue_total_processing {total_processing}")
    lines.append(f"fair_queue_total_completed_sample {total_completed_sample}")
    lines.append(f"fair_queue_total_failed_sample {total_failed_sample}")
    lines.append(f"fair_queue_is_paused {is_paused}")

    # Scheduler metrics
    scheduler = metrics_dict.scheduler
    scheduler_active_tasks = scheduler.active_tasks
    scheduler_max_concurrent_tasks = scheduler.max_concurrent_tasks
    scheduler_healthy = 1 if scheduler.healthy else 0

    lines.append(f"fair_queue_scheduler_active_tasks {scheduler_active_tasks}")
    lines.append(f"fair_queue_scheduler_max_concurrent_tasks {scheduler_max_concurrent_tasks}")
    lines.append(f"fair_queue_scheduler_healthy {scheduler_healthy}")

    # Per-group metrics
    groups = metrics_dict.groups.metrics
    for group_id, group_data in groups.items():
        # Escape group_id for Prometheus label (replace special characters)
        safe_group_id = group_id.replace('"', '\\"').replace("\n", "\\n").replace("\\", "\\\\")
        queue_length = group_data.queue_length
        weight = group_data.weight

        lines.append(f'fair_queue_group_queue_length{{group_id="{safe_group_id}"}} {queue_length}')
        lines.append(f'fair_queue_group_weight{{group_id="{safe_group_id}"}} {weight}')

    return "\n".join(lines) + "\n"


# ============================================================================
# Structured Logging
# ============================================================================


class StructuredLoggingHook:
    """Hook that logs task events in structured JSON format.

    This hook provides structured logging that can be easily parsed by log aggregation
    systems like ELK, Splunk, or CloudWatch Logs Insights.

    Example:
        ```python
        from haolib.background.fair.observability import StructuredLoggingHook

        hook = StructuredLoggingHook(
            logger=logger,
            include_task_data=False,  # Set to True to include full task data
        )
        queue = (
            FairTaskQueueBuilder()
            .with_redis_storage(redis)
            .with_taskiq_adapter(broker)
            .with_hook(hook)
            .build()
        )
        ```

    """

    def __init__(
        self,
        logger: logging.Logger | None = None,
        include_task_data: bool = False,
        include_error_traceback: bool = True,
    ) -> None:
        """Initialize structured logging hook.

        Args:
            logger: Logger instance (defaults to standard logging)
            include_task_data: If True, include full task data in logs (may be large)
            include_error_traceback: If True, include error traceback in failed task logs

        """
        self._logger = logger or logging.getLogger(__name__)
        self._include_task_data = include_task_data
        self._include_error_traceback = include_error_traceback

    def _format_log_entry(
        self,
        event_type: str,
        task_id: str,
        group_id: str,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Format a log entry as structured data.

        Args:
            event_type: Type of event (e.g., "task_enqueued")
            task_id: Task ID
            group_id: Group ID
            **kwargs: Additional fields to include

        Returns:
            Dictionary with structured log data

        """
        entry: dict[str, Any] = {
            "event_type": event_type,
            "timestamp": datetime.now(UTC).isoformat(),
            "task_id": task_id,
            "group_id": group_id,
        }
        entry.update(kwargs)
        return entry

    async def on_task_enqueued(self, event: TaskEnqueuedEvent) -> None:
        """Log task enqueued event in structured format."""
        log_data = self._format_log_entry(
            event_type="task_enqueued",
            task_id=event.task.task_id,
            group_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,
        )
        if self._include_task_data:
            log_data["task_data"] = str(event.task.task_data)
        self._logger.info(json.dumps(log_data))

    async def on_task_dequeued(self, event: TaskDequeuedEvent) -> None:
        """Log task dequeued event in structured format."""
        log_data = self._format_log_entry(
            event_type="task_dequeued",
            task_id=event.task.task_id,
            group_id=event.group_id,
            priority=event.task.priority,
            retry_count=event.task.retry_count,
        )
        self._logger.debug(json.dumps(log_data))

    async def on_task_processing(self, event: TaskProcessingEvent) -> None:
        """Log task processing event in structured format."""
        log_data = self._format_log_entry(
            event_type="task_processing",
            task_id=event.task.task_id,
            group_id=event.group_id,
            priority=event.task.priority,
            retry_count=event.task.retry_count,
        )
        self._logger.info(json.dumps(log_data))

    async def on_task_completed(self, event: TaskCompletedEvent) -> None:
        """Log task completed event in structured format."""
        log_data = self._format_log_entry(
            event_type="task_completed",
            task_id=event.task.task_id,
            group_id=event.group_id,
            priority=event.task.priority,
            processing_time_seconds=event.processing_time,
            retry_count=event.task.retry_count,
        )
        self._logger.info(json.dumps(log_data))

    async def on_task_failed(self, event: TaskFailedEvent) -> None:
        """Log task failed event in structured format."""
        log_data = self._format_log_entry(
            event_type="task_failed",
            task_id=event.task.task_id,
            group_id=event.group_id,
            priority=event.task.priority,
            retry_count=event.retry_count,
            max_retries=event.task.max_retries,
            error_type=type(event.error).__name__,
            error_message=str(event.error),
        )
        if self._include_error_traceback:
            log_data["error_traceback"] = traceback.format_exception(
                type(event.error),
                event.error,
                event.error.__traceback__,
            )
        self._logger.error(json.dumps(log_data))


# ============================================================================
# Distributed Tracing
# ============================================================================


class TracingHook:
    """Hook that creates distributed traces for task queue operations.

    Uses OpenTelemetry tracing if available, otherwise no-op.
    This hook creates spans for:
    - Task enqueue operations
    - Task processing lifecycle
    - Task completion/failure
    """

    def __init__(self) -> None:
        """Initialize tracing hook."""
        self._tracer = None

        # Try to get OpenTelemetry tracer
        try:
            from opentelemetry import trace  # noqa: PLC0415

            tracer_provider = trace.get_tracer_provider()
            if tracer_provider:
                self._tracer = trace.get_tracer("haolib.background.fair")
        except ImportError:
            # OpenTelemetry not available, tracing will be no-op
            pass
        except Exception:  # noqa: S110
            # Tracer provider not configured, tracing will be no-op
            pass

    async def on_task_enqueued(self, event: TaskEnqueuedEvent) -> None:
        """Create span for task enqueued event."""
        if not self._tracer:
            return

        with self._tracer.start_as_current_span(
            "fair_queue.task.enqueue",
            attributes={
                "task.id": event.task.task_id,
                "task.group_id": event.group_id,
                "task.priority": event.task.priority,
                "task.status": event.task.status.value,
            },
        ):
            pass

    async def on_task_dequeued(self, event: TaskDequeuedEvent) -> None:
        """Create span for task dequeued event."""
        if not self._tracer:
            return

        with self._tracer.start_as_current_span(
            "fair_queue.task.dequeue",
            attributes={
                "task.id": event.task.task_id,
                "task.group_id": event.group_id,
            },
        ):
            pass

    async def on_task_processing(self, event: TaskProcessingEvent) -> None:
        """Create span for task processing event."""
        if not self._tracer:
            return

        span = self._tracer.start_span(
            "fair_queue.task.process",
            attributes={
                "task.id": event.task.task_id,
                "task.group_id": event.group_id,
                "task.retry_count": event.task.retry_count,
            },
        )
        # Store span in task metadata for later completion
        event.task.metadata._processing_span = span  # noqa: SLF001
        span.__enter__()

    async def on_task_completed(self, event: TaskCompletedEvent) -> None:
        """Complete span for task completed event."""
        if not self._tracer:
            return

        # Try to get span from task metadata if available
        span = event.task.metadata._processing_span  # noqa: SLF001

        try:
            from opentelemetry import trace  # noqa: PLC0415
        except ImportError:
            return

        if span:
            span.set_attribute("task.processing_time", event.processing_time)
            span.set_status(trace.Status(trace.StatusCode.OK))
            span.__exit__(None, None, None)
        else:
            # Create new span if processing span not found
            with self._tracer.start_as_current_span(
                "fair_queue.task.complete",
                attributes={
                    "task.id": event.task.task_id,
                    "task.group_id": event.group_id,
                    "task.processing_time": event.processing_time,
                },
            ) as span:
                span.set_status(trace.Status(trace.StatusCode.OK))

    async def on_task_failed(self, event: TaskFailedEvent) -> None:
        """Complete span for task failed event with error."""
        if not self._tracer:
            return

        # Try to get span from task metadata if available
        span = event.task.metadata._processing_span  # noqa: SLF001

        try:
            from opentelemetry import trace  # noqa: PLC0415
        except ImportError:
            # OpenTelemetry not available, skip tracing
            return

        if span:
            span.set_attribute("task.retry_count", event.retry_count)
            span.set_attribute("error.type", type(event.error).__name__)
            span.set_attribute("error.message", str(event.error))
            span.record_exception(event.error)
            span.set_status(trace.Status(trace.StatusCode.ERROR, str(event.error)))
            span.__exit__(type(event.error), event.error, None)
        else:
            # Create new span if processing span not found

            with self._tracer.start_as_current_span(
                "fair_queue.task.fail",
                attributes={
                    "task.id": event.task.task_id,
                    "task.group_id": event.group_id,
                    "task.retry_count": event.retry_count,
                    "error.type": type(event.error).__name__,
                    "error.message": str(event.error),
                },
            ) as span:
                span.record_exception(event.error)
                span.set_status(trace.Status(trace.StatusCode.ERROR, str(event.error)))
