import json
import logging
import gzip
from abc import ABC, abstractmethod
from typing import Union

import redis.asyncio as aioredis
from sqlmodel import SQLModel

from ..models import TelemetryEvent, UsageEvent

log = logging.getLogger(__name__)


class EventStreamTransport(ABC):
    """Abstract base class for event transport mechanisms."""

    @abstractmethod
    async def send_event(self, event: Union[UsageEvent, TelemetryEvent]):
        """
        Sends a structured event to the underlying transport stream.

        Args:
            event: The event object to send, either a UsageEvent or TelemetryEvent.
        """
        pass


class RedisStreamTransport(EventStreamTransport):
    """
    An event transport that uses Redis Streams (XADD) to publish events.
    This is durable and supports multiple consumer groups.
    """

    def __init__(self, redis_client: aioredis.Redis, max_stream_length: int = 100_000):
        """
        Initializes the Redis Stream transport.

        Args:
            redis_client: An asynchronous Redis client instance.
            max_stream_length: The approximate maximum number of entries in the stream.
        """
        self.redis = redis_client
        self.max_len = max_stream_length

    async def ensure_stream(self, stream_name: str, group: str = "ingesters"):
        """Create stream and consumer group if they don't exist."""
        try:
            # Try to create consumer group, will fail if it already exists
            await self.redis.xgroup_create(stream_name, group, id="0", mkstream=True)
            log.info(f"Created consumer group '{group}' for stream '{stream_name}'")
        except aioredis.ResponseError as e:
            if "BUSYGROUP Consumer Group name already exists" not in str(e):
                log.warning(f"Could not create consumer group: {e}")

    async def send_event(self, event: Union[UsageEvent, TelemetryEvent]):
        """
        Serializes the event and adds it to a Redis Stream.

        The stream name is derived from the event's table name (e.g., 'events:usage_events').
        """
        if not isinstance(event, SQLModel):
            raise TypeError("Event must be a SQLModel instance.")

        topic = f"events:{event.__tablename__}"

        # Ensure stream and consumer group exist
        await self.ensure_stream(topic)

        # model_dump_json handles UUIDs, datetimes, etc., correctly.
        event_json = event.model_dump_json()

        # Include trace_id if available
        event_dict = json.loads(event_json)
        trace_id = event_dict.get("payload", {}).get("trace_id", "")

        # Compress payload if large (> 64KB)
        payload = {"payload": event_json}
        if len(event_json) > 65536:
            compressed = gzip.compress(event_json.encode("utf-8"))
            payload = {
                "payload": compressed.hex(),
                "compression": "gzip",
                "original_size": str(len(event_json)),
            }

        if trace_id:
            payload["trace_id"] = trace_id

        try:
            # XADD adds the event to the stream. MAXLEN ~... keeps the stream from growing indefinitely.
            await self.redis.xadd(topic, payload, maxlen=self.max_len, approximate=True)
            log.debug(f"Event published to Redis stream '{topic}'")
        except Exception as e:
            log.error(
                f"Failed to publish event to Redis stream '{topic}': {e}", exc_info=True
            )
            raise


class LogTransport(EventStreamTransport):
    """A simple event transport for local development that just logs events."""

    async def send_event(self, event: Union[UsageEvent, TelemetryEvent]):
        """Logs the event's JSON representation."""
        log.info(
            f"Event Collected ({event.__tablename__}): {event.model_dump_json(indent=2)}"
        )
