"""Collector scheduling and management."""

from __future__ import annotations

import asyncio
import contextlib
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

import structlog

from curadise_agent.collectors.base import CollectorResult, DomainCollector

if TYPE_CHECKING:
    from curadise_agent.collectors.registry import CollectorRegistry

log = structlog.get_logger(__name__)


@dataclass
class CollectionManagerStats:
    """Statistics for the collection manager."""

    total_runs: int = 0
    successful_runs: int = 0
    failed_runs: int = 0
    total_metrics: int = 0
    last_run_time: datetime | None = None


@dataclass
class ScheduledCollector:
    """Wrapper for a scheduled collector."""

    collector: DomainCollector
    task: asyncio.Task[None] | None = None
    last_run: datetime | None = None
    next_run: datetime | None = None


class CollectionManager:
    """
    Manages scheduling and execution of collectors.

    Coordinates multiple collectors, handles their lifecycle,
    and collects their metrics.
    """

    def __init__(
        self,
        registry: CollectorRegistry,
        on_metrics: Any | None = None,
    ) -> None:
        """
        Initialize collection manager.

        Args:
            registry: Collector registry
            on_metrics: Callback for collected metrics (async function)
        """
        self._registry = registry
        self._on_metrics = on_metrics
        self._scheduled: dict[str, ScheduledCollector] = {}
        self._running = False
        self._stats = CollectionManagerStats()
        self._shutdown_event = asyncio.Event()

    @property
    def is_running(self) -> bool:
        """Check if manager is running."""
        return self._running

    @property
    def stats(self) -> CollectionManagerStats:
        """Get manager statistics."""
        return self._stats

    async def start(self) -> None:
        """Start the collection manager."""
        if self._running:
            return

        self._running = True
        self._shutdown_event.clear()

        # Initialize all collectors
        for collector in self._registry.get_all_collectors():
            await collector.initialize()
            self._scheduled[collector.name] = ScheduledCollector(collector=collector)

        # Start collection loops
        for name, scheduled in self._scheduled.items():
            if scheduled.collector.enabled:
                scheduled.task = asyncio.create_task(
                    self._collection_loop(scheduled),
                    name=f"collector-{name}",
                )

        log.info("collection_manager_started", collector_count=len(self._scheduled))

    async def stop(self) -> None:
        """Stop the collection manager."""
        if not self._running:
            return

        self._running = False
        self._shutdown_event.set()

        # Cancel all tasks
        for scheduled in self._scheduled.values():
            if scheduled.task and not scheduled.task.done():
                scheduled.task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await scheduled.task

        # Shutdown all collectors
        for collector in self._registry.get_all_collectors():
            await collector.shutdown()

        self._scheduled.clear()
        log.info("collection_manager_stopped")

    async def _collection_loop(self, scheduled: ScheduledCollector) -> None:
        """Run collection loop for a single collector."""
        collector = scheduled.collector

        while self._running and collector.enabled:
            try:
                # Calculate next run time
                scheduled.next_run = datetime.now(UTC)

                # Run collection
                result = await asyncio.wait_for(
                    collector.run(),
                    timeout=collector.timeout,
                )

                scheduled.last_run = datetime.now(UTC)
                self._stats.total_runs += 1

                if result.success:
                    self._stats.successful_runs += 1
                    self._stats.total_metrics += len(result.metrics)
                    self._stats.last_run_time = scheduled.last_run

                    # Send metrics to callback
                    if self._on_metrics and result.metrics:
                        await self._on_metrics(result.metrics)
                else:
                    self._stats.failed_runs += 1

                # Wait for next interval
                await asyncio.wait_for(
                    self._shutdown_event.wait(),
                    timeout=collector.interval,
                )
                break  # Shutdown requested

            except TimeoutError:
                # Normal timeout for interval wait - continue loop
                continue
            except asyncio.CancelledError:
                break
            except Exception as e:
                log.error(
                    "collection_loop_error",
                    collector=collector.name,
                    error=str(e),
                )
                self._stats.failed_runs += 1
                # Wait before retry
                try:
                    await asyncio.wait_for(
                        self._shutdown_event.wait(),
                        timeout=min(collector.interval, 30.0),
                    )
                    break
                except TimeoutError:
                    continue

    async def run_once(self, collector_name: str | None = None) -> list[CollectorResult]:
        """
        Run collection once for specified or all collectors.

        Args:
            collector_name: Specific collector to run, or None for all

        Returns:
            List of collection results
        """
        results: list[CollectorResult] = []

        if collector_name:
            collector = self._registry.get_collector(collector_name)
            if collector:
                result = await collector.run()
                results.append(result)
                if result.success and self._on_metrics and result.metrics:
                    await self._on_metrics(result.metrics)
        else:
            for collector in self._registry.get_enabled_collectors():
                try:
                    result = await asyncio.wait_for(
                        collector.run(),
                        timeout=collector.timeout,
                    )
                    results.append(result)
                    if result.success and self._on_metrics and result.metrics:
                        await self._on_metrics(result.metrics)
                except TimeoutError:
                    results.append(
                        CollectorResult(
                            metrics=[],
                            duration_ms=collector.timeout * 1000,
                            success=False,
                            error="Collection timed out",
                        )
                    )

        return results

    def get_status(self) -> dict[str, Any]:
        """Get manager status."""
        return {
            "running": self._running,
            "stats": {
                "total_runs": self._stats.total_runs,
                "successful_runs": self._stats.successful_runs,
                "failed_runs": self._stats.failed_runs,
                "total_metrics": self._stats.total_metrics,
                "last_run": (
                    self._stats.last_run_time.isoformat() if self._stats.last_run_time else None
                ),
            },
            "collectors": {
                name: {
                    "enabled": sched.collector.enabled,
                    "last_run": sched.last_run.isoformat() if sched.last_run else None,
                    "next_run": sched.next_run.isoformat() if sched.next_run else None,
                }
                for name, sched in self._scheduled.items()
            },
        }
