"""Metric batching for efficient transmission."""

from __future__ import annotations

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

import orjson
import structlog

from curadise_agent.transport.compression import compress

if TYPE_CHECKING:
    from curadise_agent.models.domain.metric import DomainMetric

log = structlog.get_logger(__name__)


@dataclass
class BatchConfig:
    """Configuration for batching behavior."""

    max_size: int = 100
    max_bytes: int = 1_048_576  # 1MB
    max_age_seconds: float = 5.0
    compression_enabled: bool = True
    compression_level: int = 3


@dataclass
class Batch:
    """A batch of metrics ready for transmission."""

    id: str
    metrics: list[dict[str, Any]]
    created_at: datetime
    compressed: bool = False
    _raw_data: bytes | None = field(default=None, repr=False)

    @property
    def size(self) -> int:
        """Number of metrics in batch."""
        return len(self.metrics)

    @property
    def data(self) -> bytes:
        """Get serialized batch data."""
        if self._raw_data is not None:
            return self._raw_data
        return orjson.dumps(
            {"batch_id": self.id, "metrics": self.metrics, "timestamp": self.created_at.isoformat()}
        )

    @property
    def age_seconds(self) -> float:
        """Age of batch in seconds."""
        return (datetime.now(UTC) - self.created_at).total_seconds()


class BatchBuilder:
    """
    Builds batches of metrics for efficient transmission.

    Accumulates metrics until batch limits are reached, then yields
    completed batches for sending.
    """

    def __init__(self, config: BatchConfig | None = None) -> None:
        """
        Initialize batch builder.

        Args:
            config: Batching configuration
        """
        self._config = config or BatchConfig()
        self._current_batch: list[dict[str, Any]] = []
        self._current_bytes: int = 0
        self._batch_created_at: datetime | None = None
        self._batch_counter: int = 0
        self._lock = asyncio.Lock()

    @property
    def config(self) -> BatchConfig:
        """Get batch configuration."""
        return self._config

    @property
    def pending_count(self) -> int:
        """Number of metrics waiting to be batched."""
        return len(self._current_batch)

    async def add(self, metric: DomainMetric) -> Batch | None:
        """
        Add a metric to the current batch.

        Args:
            metric: Metric to add

        Returns:
            Completed batch if limits reached, None otherwise
        """
        metric_dict = metric.to_dict()
        metric_bytes = len(orjson.dumps(metric_dict))

        async with self._lock:
            # Check if this metric would exceed byte limit
            if self._current_bytes + metric_bytes > self._config.max_bytes and self._current_batch:
                batch = await self._finalize_batch()
                self._start_new_batch(metric_dict, metric_bytes)
                return batch

            # Add to current batch
            if not self._current_batch:
                self._batch_created_at = datetime.now(UTC)

            self._current_batch.append(metric_dict)
            self._current_bytes += metric_bytes

            # Check if size limit reached
            if len(self._current_batch) >= self._config.max_size:
                return await self._finalize_batch()

            return None

    async def add_many(self, metrics: list[DomainMetric]) -> list[Batch]:
        """
        Add multiple metrics, returning any completed batches.

        Args:
            metrics: Metrics to add

        Returns:
            List of completed batches
        """
        batches: list[Batch] = []
        for metric in metrics:
            batch = await self.add(metric)
            if batch:
                batches.append(batch)
        return batches

    async def flush(self) -> Batch | None:
        """
        Force flush current batch regardless of limits.

        Returns:
            Current batch if non-empty, None otherwise
        """
        async with self._lock:
            if self._current_batch:
                return await self._finalize_batch()
            return None

    async def check_age(self) -> Batch | None:
        """
        Check if current batch has exceeded max age.

        Returns:
            Current batch if aged out, None otherwise
        """
        async with self._lock:
            if self._batch_created_at is None or not self._current_batch:
                return None

            age = (datetime.now(UTC) - self._batch_created_at).total_seconds()
            if age >= self._config.max_age_seconds:
                return await self._finalize_batch()

            return None

    async def _finalize_batch(self) -> Batch:
        """Finalize and return the current batch."""
        self._batch_counter += 1
        batch_id = f"batch-{self._batch_counter:08d}"

        batch = Batch(
            id=batch_id,
            metrics=self._current_batch.copy(),
            created_at=self._batch_created_at or datetime.now(UTC),
        )

        # Compress if enabled
        if self._config.compression_enabled:
            raw_data = batch.data
            compressed_data = compress(raw_data, self._config.compression_level)
            batch._raw_data = compressed_data
            batch.compressed = True

            ratio = len(compressed_data) / len(raw_data) if raw_data else 1.0
            log.debug(
                "batch_compressed",
                batch_id=batch_id,
                original_bytes=len(raw_data),
                compressed_bytes=len(compressed_data),
                ratio=f"{ratio:.2%}",
            )

        # Reset state
        self._current_batch = []
        self._current_bytes = 0
        self._batch_created_at = None

        log.debug("batch_finalized", batch_id=batch_id, metric_count=batch.size)
        return batch

    def _start_new_batch(self, metric_dict: dict[str, Any], metric_bytes: int) -> None:
        """Start a new batch with the given metric."""
        self._current_batch = [metric_dict]
        self._current_bytes = metric_bytes
        self._batch_created_at = datetime.now(UTC)
