# |---------------------------------------------------------|
# |                                                         |
# |                 Give Feedback / Get Help                |
# | https://github.com/Pebbling-ai/pebble/issues/new/choose |
# |                                                         |
# |---------------------------------------------------------|
#
# REDIS SCHEDULER IMPLEMENTATION:
#
# This is a Redis-based implementation of the Scheduler using Upstash Redis
# for distributed task operations. Perfect for production deployments,
# multi-process systems, and cloud environments.
#
# BURGER STORE ANALOGY - REDIS IMPLEMENTATION:
#
# Think of a burger chain with multiple restaurant locations:
#
# 1. CENTRAL ORDER SYSTEM (Redis):
#    - Cloud-based order management system (Upstash Redis)
#    - All restaurants connect to the same order queue
#    - Orders persist even if individual restaurants go offline
#    - Real-time synchronization across all locations
#
# 2. ORDER FLOW:
#    - Customer order comes in → Sent to central system (Redis queue)
#    - Any kitchen staff at any location → Can pick up orders
#    - Special requests: pause, cancel, resume orders across locations
#    - Order history preserved in central database
#
# 3. ADVANTAGES:
#    - Distributed: Multiple processes can share the same queue
#    - Persistent: Orders survive restaurant restarts
#    - Scalable: Add more locations without changing the system
#    - Reliable: Cloud-hosted with automatic backups
#
# 4. FEATURES:
#    - Multi-process support (multiple restaurant locations)
#    - Persistence across restarts (orders don't get lost)
#    - Load balancing (distribute orders fairly)
#    - Real-time updates (instant order synchronization)
#
# WHEN TO USE:
# - Production deployments with multiple workers
# - Distributed agent systems
# - Cloud environments (Docker, Kubernetes)
# - When you need persistence and reliability
#
# UPSTASH REDIS:
# - Serverless Redis with automatic scaling
# - No infrastructure management required
# - Pay-per-request pricing model
# - Global edge locations for low latency
#
#  Thank you users! We ❤️ you! - 🐧

from __future__ import annotations as _annotations

import json
from collections.abc import AsyncIterator
from typing import Any

import redis.asyncio as redis
from opentelemetry.trace import get_current_span

from pebbling.common.protocol.types import TaskIdParams, TaskSendParams

from .base import (
    Scheduler,
    TaskOperation,
    _CancelTask,
    _PauseTask,
    _ResumeTask,
    _RunTask,
)


class RedisScheduler(Scheduler):
    """A Redis-based scheduler using Upstash Redis for distributed task operations."""

    def __init__(
        self,
        redis_url: str,
        queue_name: str = "pebbling:tasks",
        max_connections: int = 10,
        retry_on_timeout: bool = True,
    ):
        """Initialize Redis scheduler.

        Args:
            redis_url: Upstash Redis URL (redis://...)
            queue_name: Redis queue name for task operations
            max_connections: Maximum Redis connection pool size
            retry_on_timeout: Whether to retry on Redis timeout
        """
        self.redis_url = redis_url
        self.queue_name = queue_name
        self.max_connections = max_connections
        self.retry_on_timeout = retry_on_timeout
        self._redis_client: redis.Redis | None = None

    async def __aenter__(self):
        """Initialize Redis connection pool."""
        self._redis_client = redis.from_url(
            self.redis_url,
            encoding="utf-8",
            decode_responses=True,
            max_connections=self.max_connections,
            retry_on_timeout=self.retry_on_timeout,
        )

        # Test connection
        await self._redis_client.ping()
        return self

    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any):
        """Close Redis connection pool."""
        if self._redis_client:
            await self._redis_client.aclose()
            self._redis_client = None

    async def run_task(self, params: TaskSendParams) -> None:
        """Send a run task operation to Redis queue."""
        task_operation = _RunTask(operation="run", params=params, _current_span=get_current_span())
        await self._push_task_operation(task_operation)

    async def cancel_task(self, params: TaskIdParams) -> None:
        """Send a cancel task operation to Redis queue."""
        task_operation = _CancelTask(operation="cancel", params=params, _current_span=get_current_span())
        await self._push_task_operation(task_operation)

    async def pause_task(self, params: TaskIdParams) -> None:
        """Send a pause task operation to Redis queue."""
        task_operation = _PauseTask(operation="pause", params=params, _current_span=get_current_span())
        await self._push_task_operation(task_operation)

    async def resume_task(self, params: TaskIdParams) -> None:
        """Send a resume task operation to Redis queue."""
        task_operation = _ResumeTask(operation="resume", params=params, _current_span=get_current_span())
        await self._push_task_operation(task_operation)

    async def receive_task_operations(self) -> AsyncIterator[TaskOperation]:
        """Receive task operations from Redis queue using blocking pop."""
        if not self._redis_client:
            raise RuntimeError("Redis client not initialized. Use async context manager.")

        while True:
            try:
                # Blocking pop with 1 second timeout
                result = await self._redis_client.blpop(self.queue_name, timeout=1)

                if result:
                    _, task_data = result
                    task_operation = self._deserialize_task_operation(task_data)
                    yield task_operation

            except redis.RedisError as e:
                # Log error and continue (could add exponential backoff here)
                print(f"Redis error in receive_task_operations: {e}")
                continue
            except Exception as e:
                # Log unexpected errors
                print(f"Unexpected error in receive_task_operations: {e}")
                continue

    async def _push_task_operation(self, task_operation: TaskOperation) -> None:
        """Push a task operation to Redis queue."""
        if not self._redis_client:
            raise RuntimeError("Redis client not initialized. Use async context manager.")

        serialized_task = self._serialize_task_operation(task_operation)
        await self._redis_client.rpush(self.queue_name, serialized_task)

    def _serialize_task_operation(self, task_operation: TaskOperation) -> str:
        """Serialize task operation to JSON string for Redis storage."""
        # Convert span to string representation (spans are not JSON serializable)
        serializable_task = {
            "operation": task_operation["operation"],
            "params": task_operation["params"],
            "span_id": str(task_operation["_current_span"].get_span_context().span_id),
            "trace_id": str(task_operation["_current_span"].get_span_context().trace_id),
        }
        return json.dumps(serializable_task)

    def _deserialize_task_operation(self, task_data: str) -> TaskOperation:
        """Deserialize task operation from JSON string."""
        data = json.loads(task_data)

        # Reconstruct the task operation (span will be recreated by the worker)
        if data["operation"] == "run":
            return _RunTask(
                operation="run",
                params=data["params"],
                _current_span=get_current_span(),  # Use current span context
            )
        elif data["operation"] == "cancel":
            return _CancelTask(operation="cancel", params=data["params"], _current_span=get_current_span())
        elif data["operation"] == "pause":
            return _PauseTask(operation="pause", params=data["params"], _current_span=get_current_span())
        elif data["operation"] == "resume":
            return _ResumeTask(operation="resume", params=data["params"], _current_span=get_current_span())
        else:
            raise ValueError(f"Unknown operation type: {data['operation']}")

    async def get_queue_length(self) -> int:
        """Get the current length of the task queue."""
        if not self._redis_client:
            raise RuntimeError("Redis client not initialized. Use async context manager.")

        return await self._redis_client.llen(self.queue_name)

    async def clear_queue(self) -> int:
        """Clear all tasks from the queue. Returns number of tasks removed."""
        if not self._redis_client:
            raise RuntimeError("Redis client not initialized. Use async context manager.")

        return await self._redis_client.delete(self.queue_name)

    async def health_check(self) -> bool:
        """Check if Redis connection is healthy."""
        try:
            if not self._redis_client:
                return False
            await self._redis_client.ping()
            return True
        except Exception:
            return False
