"""
Redis-based event storage for Server-Sent Events operations.

This module provides persistent event storage with TTL cleanup for SSE operations,
enabling event replay for late connections and reliable operation monitoring.
"""

import json
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from enum import Enum
from typing import Any, cast

import redis.asyncio as redis_async
from redis import Redis

from robosystems.config.defaults import CacheDefaults
from robosystems.logger import logger


class EventType(str, Enum):
  """Standard event types for SSE operations."""

  OPERATION_STARTED = "operation_started"
  OPERATION_PROGRESS = "operation_progress"
  OPERATION_ERROR = "operation_error"
  OPERATION_COMPLETED = "operation_completed"
  OPERATION_CANCELLED = "operation_cancelled"

  # Custom event types for specific operations
  GRAPH_CREATION_PROGRESS = "graph_creation_progress"
  AGENT_ANALYSIS_PROGRESS = "agent_analysis_progress"
  BACKUP_PROGRESS = "backup_progress"
  SYNC_PROGRESS = "sync_progress"


class OperationStatus(str, Enum):
  """Operation status values."""

  PENDING = "pending"
  RUNNING = "running"
  COMPLETED = "completed"
  FAILED = "failed"
  CANCELLED = "cancelled"


@dataclass
class SSEEvent:
  """
  Structured event for SSE operations.

  All events are JSON-serializable and include standard metadata
  for consistent client handling.
  """

  event_type: EventType
  operation_id: str
  timestamp: str
  data: dict[str, Any]
  sequence_number: int = 0

  def to_sse_format(self) -> str:
    """
    Convert to Server-Sent Events format.

    Note: This method is kept for backwards compatibility.
    New code should use sse-starlette's EventSourceResponse instead.
    """
    lines = []
    lines.append(f"event: {self.event_type.value}")
    # Ensure JSON is on a single line for SSE format
    data_json = json.dumps(
      {
        "operation_id": self.operation_id,
        "timestamp": self.timestamp,
        "sequence_number": self.sequence_number,
        **self.data,
      },
      separators=(",", ":"),
      ensure_ascii=False,
      default=str,
    )
    lines.append(f"data: {data_json}")
    # SSE format requires double newline to terminate event
    return "\n".join(lines) + "\n\n"

  def to_dict(self) -> dict[str, Any]:
    """Convert to dictionary for JSON serialization."""
    return asdict(self)

  @classmethod
  def from_dict(cls, data: dict[str, Any]) -> "SSEEvent":
    """Create SSEEvent from dictionary."""
    return cls(**data)


@dataclass
class OperationMetadata:
  """
  Metadata about an operation for tracking and management.
  """

  operation_id: str
  operation_type: str
  user_id: str
  graph_id: str | None
  status: OperationStatus
  created_at: str
  updated_at: str
  error_message: str | None = None
  result_data: dict[str, Any] | None = None

  def to_dict(self) -> dict[str, Any]:
    """Convert to dictionary for JSON serialization."""
    return asdict(self)


class SSEEventStorage:
  """
  Redis-based storage for SSE events with automatic TTL cleanup.

  Provides event persistence for replay capability and operation tracking.
  Events are stored per operation with automatic expiration.
  """

  def __init__(
    self,
    redis_client: redis_async.Redis | None = None,
    default_ttl: int = CacheDefaults.LONG,
  ):
    """
    Initialize event storage.

    Args:
        redis_client: Async Redis client instance (uses default if None)
        default_ttl: Default TTL for events in seconds
    """
    self._redis_client = redis_client
    self._async_redis = None
    self._sync_redis = None  # For sync methods (background tasks)
    self.default_ttl = default_ttl

    # Redis key prefixes
    self.event_prefix = "sse:operation:events:"
    self.metadata_prefix = "sse:operation:meta:"
    self.sequence_prefix = "sse:operation:seq:"

  async def _get_redis(self) -> redis_async.Redis:
    """Get async Redis client, creating if needed."""
    if self._async_redis is None:
      if self._redis_client:
        self._async_redis = self._redis_client
      else:
        self._async_redis = await self._get_default_async_redis()
    return self._async_redis

  async def _get_default_async_redis(self) -> redis_async.Redis:
    """Get default async Redis client from environment."""
    from robosystems.config.valkey_registry import (
      ValkeyDatabase,
      create_async_redis_client,
    )

    # Use SSE events database from registry with proper ElastiCache support
    client = create_async_redis_client(ValkeyDatabase.SSE)
    # Test connection
    await client.ping()
    return client

  def _get_sync_redis(self) -> Redis:
    """Get synchronous Redis client for background tasks."""
    if self._sync_redis is None:
      from robosystems.config.valkey_registry import ValkeyDatabase, create_redis_client

      # Use SSE events database from registry with proper ElastiCache support
      self._sync_redis = create_redis_client(ValkeyDatabase.SSE)
      # Test connection
      self._sync_redis.ping()
    return self._sync_redis

  def generate_operation_id(self) -> str:
    """Generate a unique operation ID (prefixed ULID)."""
    from robosystems.utils.ulid import generate_prefixed_ulid

    return generate_prefixed_ulid("op")

  async def create_operation(
    self,
    operation_type: str,
    user_id: str,
    graph_id: str | None = None,
    operation_id: str | None = None,
    ttl: int | None = None,
  ) -> str:
    """
    Create a new operation and return its ID.

    Args:
        operation_type: Type of operation (e.g., "graph_creation", "agent_analysis")
        user_id: User who initiated the operation
        graph_id: Graph database ID (if applicable)
        operation_id: Custom operation ID (generates one if None)
        ttl: Custom TTL for this operation (uses default if None)

    Returns:
        str: The operation ID
    """
    if operation_id is None:
      operation_id = self.generate_operation_id()

    ttl = ttl or self.default_ttl
    now = datetime.now(UTC).isoformat()

    # Create operation metadata
    metadata = OperationMetadata(
      operation_id=operation_id,
      operation_type=operation_type,
      user_id=user_id,
      graph_id=graph_id,
      status=OperationStatus.PENDING,
      created_at=now,
      updated_at=now,
    )

    # Store metadata with TTL
    redis = await self._get_redis()
    metadata_key = f"{self.metadata_prefix}{operation_id}"
    await redis.setex(metadata_key, ttl, json.dumps(metadata.to_dict()))

    # Initialize sequence counter
    seq_key = f"{self.sequence_prefix}{operation_id}"
    await redis.setex(seq_key, ttl, "0")

    logger.info(
      f"Created SSE operation {operation_id} of type {operation_type} for user {user_id}"
    )

    return operation_id

  async def store_event(
    self,
    operation_id: str,
    event_type: EventType,
    data: dict[str, Any],
    ttl: int | None = None,
  ) -> SSEEvent:
    """
    Store an event for an operation.

    Args:
        operation_id: Operation identifier
        event_type: Type of event
        data: Event data
        ttl: Custom TTL (uses default if None)

    Returns:
        SSEEvent: The stored event

    Raises:
        ValueError: If operation doesn't exist
    """
    ttl = ttl or self.default_ttl

    # Check if operation exists
    redis = await self._get_redis()
    metadata_key = f"{self.metadata_prefix}{operation_id}"
    if not await redis.exists(metadata_key):
      raise ValueError(f"Operation {operation_id} not found")

    # Get and increment sequence number
    seq_key = f"{self.sequence_prefix}{operation_id}"
    sequence_number = await redis.incr(seq_key)

    # Create event
    event = SSEEvent(
      event_type=event_type,
      operation_id=operation_id,
      timestamp=datetime.now(UTC).isoformat(),
      data=data,
      sequence_number=sequence_number,
    )

    # Store event in a sorted set by sequence number
    events_key = f"{self.event_prefix}{operation_id}"
    await redis.zadd(events_key, {json.dumps(event.to_dict()): sequence_number})

    # Set TTL on events set
    await redis.expire(events_key, ttl)

    # Publish event to Redis pub/sub channel for real-time notifications
    # This allows the API process to receive events from worker processes
    channel = f"sse:events:{operation_id}"
    await redis.publish(channel, json.dumps(event.to_dict()))

    # Update operation metadata timestamp and possibly status
    await self._update_operation_metadata(operation_id, event_type, data)

    logger.debug(
      f"Stored event {event_type} for operation {operation_id} (seq: {sequence_number}), published to {channel}"
    )

    return event

  def store_event_sync(
    self,
    operation_id: str,
    event_type: EventType,
    data: dict[str, Any],
    ttl: int | None = None,
  ) -> SSEEvent:
    """
    Synchronous version of store_event for use in background tasks.

    Args:
        operation_id: Operation identifier
        event_type: Type of event
        data: Event data
        ttl: Custom TTL (uses default if None)

    Returns:
        SSEEvent: The stored event

    Raises:
        ValueError: If operation doesn't exist
    """
    ttl = ttl or self.default_ttl

    # Check if operation exists
    redis = self._get_sync_redis()
    metadata_key = f"{self.metadata_prefix}{operation_id}"
    if not redis.exists(metadata_key):
      # For sync context, we might not have the operation created yet
      # Just log and continue
      logger.warning(
        f"Operation {operation_id} not found in metadata, continuing anyway"
      )

    # Get and increment sequence number
    seq_key = f"{self.sequence_prefix}{operation_id}"
    sequence_number = cast(int, redis.incr(seq_key))

    # Create event
    event = SSEEvent(
      event_type=event_type,
      operation_id=operation_id,
      timestamp=datetime.now(UTC).isoformat(),
      data=data,
      sequence_number=sequence_number,
    )

    # Store event in a sorted set by sequence number
    events_key = f"{self.event_prefix}{operation_id}"
    redis.zadd(events_key, {json.dumps(event.to_dict()): sequence_number})

    # Set TTL on events set
    redis.expire(events_key, ttl)

    # Publish event to Redis pub/sub channel for real-time notifications
    channel = f"sse:events:{operation_id}"
    redis.publish(channel, json.dumps(event.to_dict()))

    # For sync version, update metadata directly
    self._update_operation_metadata_sync(operation_id, event_type, data)

    logger.debug(
      f"[SYNC] Stored event {event_type} for operation {operation_id} (seq: {sequence_number}), published to {channel}"
    )

    return event

  def _update_operation_metadata_sync(
    self, operation_id: str, event_type: EventType, data: dict[str, Any]
  ):
    """Synchronous version to update operation metadata based on event type.

    Uses atomic Redis operations to prevent race conditions when multiple
    events are stored concurrently. Only status-changing events (started,
    completed, error, cancelled) update the full metadata. Progress events
    are skipped to avoid overwriting status changes.
    """
    # Skip progress events - they don't change status and can cause race conditions
    # where a progress event overwrites a completed/failed status
    status_changing_events = {
      EventType.OPERATION_STARTED,
      EventType.OPERATION_COMPLETED,
      EventType.OPERATION_ERROR,
      EventType.OPERATION_CANCELLED,
    }
    if event_type not in status_changing_events:
      return

    redis = self._get_sync_redis()
    metadata_key = f"{self.metadata_prefix}{operation_id}"

    # Use optimistic locking with WATCH to prevent race conditions
    # between concurrent status-changing events
    with redis.pipeline() as pipe:
      try:
        # Watch the key for changes
        pipe.watch(metadata_key)

        # Get current metadata
        metadata_json = redis.get(metadata_key)
        if not metadata_json:
          pipe.unwatch()
          return

        metadata_dict = json.loads(str(metadata_json))
        metadata = OperationMetadata(**metadata_dict)

        # Don't downgrade status (e.g., don't go from COMPLETED back to RUNNING)
        status_priority = {
          OperationStatus.PENDING: 0,
          OperationStatus.RUNNING: 1,
          OperationStatus.COMPLETED: 2,
          OperationStatus.FAILED: 2,
          OperationStatus.CANCELLED: 2,
        }
        new_status = None
        if event_type == EventType.OPERATION_STARTED:
          new_status = OperationStatus.RUNNING
        elif event_type == EventType.OPERATION_COMPLETED:
          new_status = OperationStatus.COMPLETED
        elif event_type == EventType.OPERATION_ERROR:
          new_status = OperationStatus.FAILED
        elif event_type == EventType.OPERATION_CANCELLED:
          new_status = OperationStatus.CANCELLED

        # Only update if new status has higher or equal priority
        if new_status and status_priority.get(new_status, 0) >= status_priority.get(
          metadata.status, 0
        ):
          metadata.status = new_status
          metadata.updated_at = datetime.now(UTC).isoformat()

          # Handle result/error data
          if event_type == EventType.OPERATION_COMPLETED:
            # Merge new data with existing result_data (preserves graph_id from Dagster job)
            new_result = data.get("result") or data
            if metadata.result_data:
              metadata.result_data.update(new_result)
            else:
              metadata.result_data = new_result
          elif event_type == EventType.OPERATION_ERROR:
            metadata.error_message = data.get("error", "Unknown error")

          # Store atomically
          pipe.multi()
          pipe.setex(metadata_key, self.default_ttl, json.dumps(metadata.to_dict()))
          pipe.execute()
        else:
          pipe.unwatch()

      except Exception:
        # Transaction failed (key was modified), which is fine - another
        # status-changing event won. Just log and continue.
        logger.debug(
          f"[SYNC] Metadata update skipped for {operation_id} due to concurrent modification"
        )

  def update_operation_result_sync(
    self, operation_id: str, result: dict[str, Any]
  ) -> None:
    """
    Update operation metadata with result data (sync version).

    This is used by Dagster jobs to store the graph_id and other result data
    in the operation metadata before the job completes. When the monitor later
    emits the OPERATION_COMPLETED event, it will include this result data.

    Args:
        operation_id: Operation identifier
        result: Result data to merge into metadata (e.g., {"graph_id": "kg123"})
    """
    redis = self._get_sync_redis()
    metadata_key = f"{self.metadata_prefix}{operation_id}"
    metadata_json = redis.get(metadata_key)

    if not metadata_json:
      logger.warning(f"Operation {operation_id} not found, cannot update result")
      return

    metadata_dict = json.loads(str(metadata_json))
    metadata = OperationMetadata(**metadata_dict)

    # Update timestamp and merge result data
    metadata.updated_at = datetime.now(UTC).isoformat()

    # Merge new result with existing result_data (if any)
    if metadata.result_data:
      metadata.result_data.update(result)
    else:
      metadata.result_data = result

    # Update graph_id if provided
    if "graph_id" in result:
      metadata.graph_id = result["graph_id"]

    # Store updated metadata
    redis.setex(
      metadata_key,
      self.default_ttl,
      json.dumps(metadata.to_dict()),
    )

    logger.debug(
      f"Updated operation {operation_id} result data with keys: {list(result.keys())}"
    )

  def get_operation_result_sync(self, operation_id: str) -> dict[str, Any] | None:
    """
    Get operation result data (sync version).

    This is used by the Dagster monitor to retrieve stored result data
    when emitting the OPERATION_COMPLETED event.

    Args:
        operation_id: Operation ID to get result for

    Returns:
        Result data dictionary if found, None otherwise
    """
    redis = self._get_sync_redis()
    metadata_key = f"{self.metadata_prefix}{operation_id}"
    metadata_json = redis.get(metadata_key)

    if not metadata_json:
      return None

    try:
      metadata_dict = json.loads(str(metadata_json))
      metadata = OperationMetadata(**metadata_dict)
      return metadata.result_data
    except Exception as e:
      logger.warning(f"Failed to get operation result for {operation_id}: {e}")
      return None

  async def _update_operation_metadata(
    self, operation_id: str, event_type: EventType, data: dict[str, Any]
  ):
    """Update operation metadata based on event type.

    Uses atomic Redis operations to prevent race conditions when multiple
    events are stored concurrently. Only status-changing events (started,
    completed, error, cancelled) update the full metadata. Progress events
    are skipped to avoid overwriting status changes.
    """
    # Skip progress events - they don't change status and can cause race conditions
    # where a progress event overwrites a completed/failed status
    status_changing_events = {
      EventType.OPERATION_STARTED,
      EventType.OPERATION_COMPLETED,
      EventType.OPERATION_ERROR,
      EventType.OPERATION_CANCELLED,
    }
    if event_type not in status_changing_events:
      return

    redis = await self._get_redis()
    metadata_key = f"{self.metadata_prefix}{operation_id}"

    try:
      # Watch the key for changes - this enables optimistic locking
      await redis.watch(metadata_key)

      # Get current metadata
      metadata_json = await redis.get(metadata_key)
      if not metadata_json:
        await redis.unwatch()
        return

      metadata_dict = json.loads(metadata_json)
      metadata = OperationMetadata(**metadata_dict)

      # Don't downgrade status (e.g., don't go from COMPLETED back to RUNNING)
      status_priority = {
        OperationStatus.PENDING: 0,
        OperationStatus.RUNNING: 1,
        OperationStatus.COMPLETED: 2,
        OperationStatus.FAILED: 2,
        OperationStatus.CANCELLED: 2,
      }
      new_status = None
      if event_type == EventType.OPERATION_STARTED:
        new_status = OperationStatus.RUNNING
      elif event_type == EventType.OPERATION_COMPLETED:
        new_status = OperationStatus.COMPLETED
      elif event_type == EventType.OPERATION_ERROR:
        new_status = OperationStatus.FAILED
      elif event_type == EventType.OPERATION_CANCELLED:
        new_status = OperationStatus.CANCELLED

      # Only update if new status has higher or equal priority
      if new_status and status_priority.get(new_status, 0) >= status_priority.get(
        metadata.status, 0
      ):
        metadata.status = new_status
        metadata.updated_at = datetime.now(UTC).isoformat()

        # Handle result/error data
        if event_type == EventType.OPERATION_COMPLETED:
          new_result = data.get("result") or {}
          if metadata.result_data:
            metadata.result_data.update(new_result)
          else:
            metadata.result_data = new_result
        elif event_type == EventType.OPERATION_ERROR:
          metadata.error_message = data.get("error", "Unknown error")

        # Get TTL and store atomically using pipeline
        ttl = await redis.ttl(metadata_key)
        if ttl > 0:
          pipe = redis.pipeline()
          pipe.setex(metadata_key, ttl, json.dumps(metadata.to_dict()))
          await pipe.execute()
      else:
        await redis.unwatch()

    except Exception:
      # Transaction failed (key was modified), which is fine - another
      # status-changing event won. Just log and continue.
      logger.debug(
        f"Metadata update skipped for {operation_id} due to concurrent modification"
      )

  async def get_events(
    self, operation_id: str, from_sequence: int = 0, limit: int | None = None
  ) -> list[SSEEvent]:
    """
    Retrieve events for an operation.

    Args:
        operation_id: Operation identifier
        from_sequence: Start from this sequence number (inclusive)
        limit: Maximum number of events to return

    Returns:
        List[SSEEvent]: Events in sequence order
    """
    redis = await self._get_redis()
    events_key = f"{self.event_prefix}{operation_id}"

    # Get events from sorted set
    if limit:
      raw_events = await redis.zrangebyscore(
        events_key, from_sequence, "+inf", start=0, num=limit
      )
    else:
      raw_events = await redis.zrangebyscore(events_key, from_sequence, "+inf")

    # Parse events
    events = []
    for raw_event in raw_events:
      try:
        event_dict = json.loads(raw_event)
        events.append(SSEEvent(**event_dict))
      except (json.JSONDecodeError, TypeError) as e:
        logger.warning(f"Failed to parse event for operation {operation_id}: {e}")
        continue

    return events

  async def get_operation_metadata(self, operation_id: str) -> OperationMetadata | None:
    """
    Get operation metadata.

    Args:
        operation_id: Operation identifier

    Returns:
        OperationMetadata or None if not found
    """
    redis = await self._get_redis()
    metadata_key = f"{self.metadata_prefix}{operation_id}"
    metadata_json = await redis.get(metadata_key)

    if not metadata_json:
      return None

    try:
      metadata_dict = json.loads(metadata_json)
      return OperationMetadata(**metadata_dict)
    except (json.JSONDecodeError, TypeError) as e:
      logger.warning(f"Failed to parse metadata for operation {operation_id}: {e}")
      return None

  async def cancel_operation(
    self, operation_id: str, reason: str = "Cancelled by user"
  ):
    """
    Cancel an operation by storing a cancellation event.

    Args:
        operation_id: Operation identifier
        reason: Cancellation reason
    """
    await self.store_event(
      operation_id, EventType.OPERATION_CANCELLED, {"reason": reason}
    )

  async def cleanup_expired_operations(self) -> int:
    """
    Clean up expired operations (Redis should handle TTL, but this is a backup).

    Returns:
        int: Number of operations cleaned up
    """
    # This is mainly for manual cleanup or operations that somehow didn't get TTL
    cleaned = 0
    redis = await self._get_redis()

    # Scan for expired metadata keys
    async for key in redis.scan_iter(match=f"{self.metadata_prefix}*"):
      if not await redis.exists(key):
        # Already expired
        continue

      ttl = await redis.ttl(key)
      if ttl == -1:  # No TTL set (shouldn't happen)
        # Set default TTL
        await redis.expire(key, self.default_ttl)
        cleaned += 1

    logger.info(f"Cleaned up {cleaned} expired operations")
    return cleaned


# Global instance (initialized lazily)
_event_storage: SSEEventStorage | None = None


def get_event_storage() -> SSEEventStorage:
  """Get the global event storage instance."""
  global _event_storage
  if _event_storage is None:
    _event_storage = SSEEventStorage()
  return _event_storage
