"""
Smart chunk batching for streaming to reduce HTTP requests.

Instead of sending one POST per chunk (50-70 requests), batch chunks
with configurable time/size windows (5-10 requests).

Batching Strategy:
- Time window: Flush after X ms (default: 100ms)
- Size window: Flush when batch reaches Y bytes (default: 100 bytes)
- Immediate flush: On tool events, errors, or completion

This provides:
- 90%+ reduction in HTTP requests
- Still feels real-time (100ms is imperceptible)
- Lower latency (fewer round trips)
- Better serverless performance (fewer cold starts)
- Lower costs (fewer invocations)
"""

import asyncio
import time
from typing import Dict, Any, Optional, Callable
from dataclasses import dataclass, field
import structlog

logger = structlog.get_logger()


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

    # Time-based batching: flush after this many milliseconds
    time_window_ms: int = 100

    # Size-based batching: flush when accumulated content reaches this size
    size_window_bytes: int = 100

    # Maximum batch size before forced flush (safety limit)
    max_batch_size_bytes: int = 1000

    # Enable/disable batching (for testing/debugging)
    enabled: bool = True

    @classmethod
    def from_env(cls) -> "BatchConfig":
        """
        Create configuration from environment variables.

        Environment variables:
            CHUNK_BATCHING_ENABLED: Enable/disable batching (default: true)
            CHUNK_BATCHING_TIME_WINDOW_MS: Time window in ms (default: 100)
            CHUNK_BATCHING_SIZE_WINDOW_BYTES: Size window in bytes (default: 100)
            CHUNK_BATCHING_MAX_SIZE_BYTES: Max batch size in bytes (default: 1000)

        Returns:
            BatchConfig instance with values from environment
        """
        import os

        return cls(
            enabled=os.getenv("CHUNK_BATCHING_ENABLED", "true").lower() == "true",
            time_window_ms=int(os.getenv("CHUNK_BATCHING_TIME_WINDOW_MS", "100")),
            size_window_bytes=int(os.getenv("CHUNK_BATCHING_SIZE_WINDOW_BYTES", "100")),
            max_batch_size_bytes=int(os.getenv("CHUNK_BATCHING_MAX_SIZE_BYTES", "1000")),
        )


@dataclass
class ContentBatch:
    """Accumulated content chunks waiting to be flushed."""

    chunks: list[str] = field(default_factory=list)
    total_size: int = 0
    first_chunk_time: Optional[float] = None

    def add(self, content: str) -> None:
        """Add content to the batch."""
        self.chunks.append(content)
        self.total_size += len(content.encode('utf-8'))

        if self.first_chunk_time is None:
            self.first_chunk_time = time.time()

    def get_combined_content(self) -> str:
        """Get all chunks combined into single string."""
        return ''.join(self.chunks)

    def clear(self) -> None:
        """Clear the batch after flushing."""
        self.chunks.clear()
        self.total_size = 0
        self.first_chunk_time = None

    def is_empty(self) -> bool:
        """Check if batch is empty."""
        return len(self.chunks) == 0

    def age_ms(self) -> float:
        """Get age of batch in milliseconds."""
        if self.first_chunk_time is None:
            return 0
        return (time.time() - self.first_chunk_time) * 1000


class ChunkBatcher:
    """
    Smart batching for streaming chunks to reduce HTTP requests.

    Usage:
        batcher = ChunkBatcher(
            publish_func=control_plane.publish_event,
            execution_id=execution_id,
            message_id=message_id,
            config=BatchConfig(time_window_ms=100, size_window_bytes=100)
        )

        # Add chunks as they arrive
        await batcher.add_chunk("Hello")
        await batcher.add_chunk(" world")

        # Flush remaining chunks when done
        await batcher.flush()
    """

    def __init__(
        self,
        publish_func: Callable,
        execution_id: str,
        message_id: str,
        config: Optional[BatchConfig] = None
    ):
        self.publish_func = publish_func
        self.execution_id = execution_id
        self.message_id = message_id
        self.config = config or BatchConfig()

        self.batch = ContentBatch()
        self._flush_task: Optional[asyncio.Task] = None
        self._stats = {
            "chunks_received": 0,
            "batches_sent": 0,
            "bytes_sent": 0,
            "flushes_by_time": 0,
            "flushes_by_size": 0,
            "flushes_manual": 0,
        }

    async def add_chunk(self, content: str) -> None:
        """
        Add a chunk to the batch.

        Automatically flushes if:
        - Batch size exceeds size_window_bytes
        - Batch age exceeds time_window_ms
        - Max batch size is reached (safety)
        """
        if not self.config.enabled:
            # Batching disabled - send immediately
            await self._publish_batch([content])
            return

        self._stats["chunks_received"] += 1
        self.batch.add(content)

        # Check if we should flush immediately due to size
        should_flush_size = self.batch.total_size >= self.config.size_window_bytes
        should_flush_max = self.batch.total_size >= self.config.max_batch_size_bytes

        if should_flush_max:
            # Safety: flush immediately if max size reached
            logger.debug(
                "Flushing batch (max size reached)",
                execution_id=self.execution_id[:8],
                batch_size=self.batch.total_size,
                chunk_count=len(self.batch.chunks),
            )
            await self.flush(reason="max_size")
        elif should_flush_size:
            # Size threshold reached - flush now
            await self.flush(reason="size")
        else:
            # Start/reset timer for time-based flush
            await self._schedule_time_flush()

    async def _schedule_time_flush(self) -> None:
        """Schedule a time-based flush if not already scheduled."""
        if self._flush_task is not None and not self._flush_task.done():
            # Timer already running
            return

        self._flush_task = asyncio.create_task(self._time_based_flush())

    async def _time_based_flush(self) -> None:
        """Wait for time window, then flush."""
        await asyncio.sleep(self.config.time_window_ms / 1000.0)

        if not self.batch.is_empty():
            await self.flush(reason="time")

    async def flush(self, reason: str = "manual") -> None:
        """
        Flush current batch immediately.

        Args:
            reason: Why flush was triggered (for stats/debugging)
        """
        if self.batch.is_empty():
            return

        # Cancel pending timer if any
        if self._flush_task is not None and not self._flush_task.done():
            self._flush_task.cancel()
            try:
                await self._flush_task
            except asyncio.CancelledError:
                pass

        # Publish the batch
        chunks = self.batch.chunks.copy()
        await self._publish_batch(chunks)

        # Update stats
        if reason == "time":
            self._stats["flushes_by_time"] += 1
        elif reason == "size" or reason == "max_size":
            self._stats["flushes_by_size"] += 1
        else:
            self._stats["flushes_manual"] += 1

        # Clear batch
        self.batch.clear()

    async def _publish_batch(self, chunks: list[str]) -> None:
        """Publish a batch of chunks as single event."""
        combined_content = ''.join(chunks)

        try:
            # Call publish_func (could be sync or async)
            result = self.publish_func(
                execution_id=self.execution_id,
                event_type="message_chunk",
                data={
                    "role": "assistant",
                    "content": combined_content,
                    "is_chunk": True,
                    "message_id": self.message_id,
                    # Metadata for debugging
                    "batch_info": {
                        "chunk_count": len(chunks),
                        "batch_size": len(combined_content.encode('utf-8')),
                    } if len(chunks) > 1 else None,
                }
            )

            # Await if it's a coroutine or awaitable (Future, Task, etc.)
            # Check in order of reliability: __await__ method > coroutine > future
            if hasattr(result, '__await__'):
                await result
            elif asyncio.iscoroutine(result):
                await result
            elif asyncio.isfuture(result):
                await result

            self._stats["batches_sent"] += 1
            self._stats["bytes_sent"] += len(combined_content.encode('utf-8'))

        except Exception as e:
            logger.warning(
                "Failed to publish batch",
                execution_id=self.execution_id[:8],
                error=str(e),
                chunk_count=len(chunks),
            )

    def get_stats(self) -> Dict[str, Any]:
        """
        Get batching statistics.

        Returns:
            Dict with stats about batching performance
        """
        chunks_received = self._stats["chunks_received"]
        batches_sent = self._stats["batches_sent"]

        return {
            **self._stats,
            "reduction_percent": round(
                (1 - batches_sent / max(chunks_received, 1)) * 100, 1
            ) if chunks_received > 0 else 0,
            "avg_batch_size": round(
                chunks_received / max(batches_sent, 1), 1
            ) if batches_sent > 0 else 0,
        }

    async def close(self) -> None:
        """
        Close the batcher and flush remaining chunks.

        Call this when streaming is complete.
        """
        await self.flush(reason="close")

        # Log stats
        stats = self.get_stats()
        logger.info(
            "Chunk batching stats",
            execution_id=self.execution_id[:8],
            **stats
        )
