# Copyright (c) Metis. All rights reserved.

"""RedisStreamLightningStore — distributed rollout coordination via Redis Streams.

This store wraps InMemoryLightningStore for all metadata (rollouts, attempts,
spans, resources) and replaces only the rollout queue with Redis Streams.

Redis Stream protocol:
  - Stream key: ``mantis:run:{run_id}:rollouts``
  - Consumer group: ``mantis-runners``
  - Enqueue: ``XADD`` with rollout_id, input_json, resources_id, mode
  - Dequeue: ``XREADGROUP GROUP mantis-runners {worker_id} COUNT 1 BLOCK {timeout_ms}``
  - Ack: ``XACK`` after rollout reaches terminal status

Usage::

    store = RedisStreamLightningStore(
        redis_url="redis://:myredissecret@localhost:6379",
        run_id="run-abc123",
    )
    await store.initialize()

    # Enqueue rollouts (algorithm side)
    rollout = await store.enqueue_rollout(input={"task": "solve"})

    # Dequeue rollouts (runner side)
    attempted = await store.dequeue_rollout(worker_id="runner-1")

    await store.close()
"""

from __future__ import annotations

import json
import logging
import time
from typing import Any, Dict, List, Optional, Sequence

from mantisdk.store.base import (
    UNSET,
    LightningStore,
    LightningStoreCapabilities,
    LightningStoreStatistics,
    Unset,
)
from mantisdk.types import (
    Attempt,
    AttemptedRollout,
    AttemptStatus,
    EnqueueRolloutRequest,
    NamedResources,
    PaginatedResult,
    ResourcesUpdate,
    Rollout,
    RolloutConfig,
    RolloutMode,
    RolloutStatus,
    Span,
    TaskInput,
    Worker,
    WorkerStatus,
)

logger = logging.getLogger(__name__)


class RedisStreamLightningStore(LightningStore):
    """DEPRECATED: Use LightningStoreServer + LightningStoreClient instead.

    This store wraps InMemoryLightningStore with Redis Streams for distributed
    rollout coordination. It has been superseded by the StoreServer/Client
    architecture which avoids the rollout desync bug (separate memory stores
    on algorithm and agent sides can't see each other's state).

    For platform runs, the worker now creates an InMemoryLightningStore
    wrapped in a LightningStoreServer. Agent containers connect via
    LightningStoreClient (HTTP). See platform/worker.py AlgorithmContext.
    """

    @classmethod
    def from_env(cls) -> "RedisStreamLightningStore":
        """Create a store from environment variables set by DockerOrchestrator.

        Expected env vars:
            ``MANTIS_RUN_ID``: The MantisRun ID.
            ``REDIS_HOST``: Redis hostname (default ``localhost``).
            ``REDIS_PORT``: Redis port (default ``6379``).
            ``REDIS_AUTH``: Redis password (optional).

        Returns:
            A configured but not-yet-initialized store.
            Caller must ``await store.initialize()`` before use.

        Raises:
            EnvironmentError: If ``MANTIS_RUN_ID`` is not set.
        """
        import os

        run_id = os.environ.get("MANTIS_RUN_ID")
        if not run_id:
            raise EnvironmentError(
                "MANTIS_RUN_ID environment variable is required. "
                "This is set automatically by the Mantis platform orchestrator."
            )
        host = os.environ.get("REDIS_HOST", "localhost")
        port = os.environ.get("REDIS_PORT", "6379")
        auth = os.environ.get("REDIS_AUTH", "")
        redis_url = f"redis://:{auth}@{host}:{port}" if auth else f"redis://{host}:{port}"
        return cls(redis_url=redis_url, run_id=run_id)

    def __init__(
        self,
        *,
        redis_url: str,
        run_id: str,
        consumer_group: str = "mantis-runners",
        thread_safe: bool = False,
        listeners: Optional[Sequence] = None,
        block_timeout_ms: int = 5000,
    ):
        """
        Args:
            redis_url: Redis connection string (e.g., redis://:password@host:6379/0).
            run_id: The MantisRun ID. Used to namespace the Redis stream key.
            consumer_group: Consumer group name for XREADGROUP.
            thread_safe: Whether the underlying InMemoryLightningStore is thread-safe.
            listeners: StorageListener instances to attach.
            block_timeout_ms: Timeout for XREADGROUP BLOCK (milliseconds).
        """
        self._redis_url = redis_url
        self._run_id = run_id
        self._consumer_group = consumer_group
        self._block_timeout_ms = block_timeout_ms
        self._redis = None
        self._stream_key = f"mantis:run:{run_id}:rollouts"

        # Create the inner in-memory store for all metadata
        from mantisdk.store.memory import InMemoryLightningStore

        self._inner = InMemoryLightningStore(
            thread_safe=thread_safe,
            listeners=list(listeners) if listeners else None,
        )

    @property
    def run_id(self) -> str:
        return self._run_id

    @property
    def capabilities(self) -> LightningStoreCapabilities:
        caps = self._inner.capabilities
        caps["async_safe"] = True
        caps["thread_safe"] = True  # Redis handles concurrency at server level
        return caps

    async def initialize(self) -> None:
        """Create Redis connection and consumer group."""
        try:
            import redis.asyncio as aioredis
        except ImportError:
            raise ImportError(
                "redis[asyncio] is required for RedisStreamLightningStore. " "Install with: pip install redis>=5.0.0"
            )

        try:
            self._redis = aioredis.from_url(self._redis_url, decode_responses=True)
            await self._redis.ping()
        except Exception as e:
            self._redis = None
            from mantisdk.utils.redact import redact_url

            raise ConnectionError(f"Failed to connect to Redis at {redact_url(self._redis_url)}: {e}") from e

        # Create consumer group (MKSTREAM creates the stream if it doesn't exist)
        try:
            await self._redis.xgroup_create(self._stream_key, self._consumer_group, id="0", mkstream=True)
            logger.info(f"Created consumer group '{self._consumer_group}' on stream '{self._stream_key}'")
        except Exception as e:
            if "BUSYGROUP" in str(e):
                logger.debug(f"Consumer group '{self._consumer_group}' already exists")
            else:
                raise

    async def close(self, *, cleanup: bool = True) -> None:
        """Close Redis connection and optionally clean up stream resources.

        Args:
            cleanup: If True (default), delete the stream key and consumer group
                to prevent Redis memory growth from orphaned streams.
        """
        if self._redis:
            if cleanup:
                try:
                    await self._redis.xgroup_destroy(self._stream_key, self._consumer_group)
                    logger.debug(f"Destroyed consumer group '{self._consumer_group}' " f"on '{self._stream_key}'")
                except Exception:
                    pass  # Group may not exist or Redis may be closing

                try:
                    await self._redis.delete(self._stream_key)
                    logger.debug(f"Deleted stream key '{self._stream_key}'")
                except Exception:
                    pass

                # Clean up cancellation signal key if it exists
                try:
                    await self._redis.delete(f"mantis:run:{self._run_id}:cancel")
                except Exception:
                    pass

            await self._redis.aclose()
            self._redis = None

    # =========================================================================
    # Rollout Queue — Redis Stream backed
    # =========================================================================

    async def enqueue_rollout(
        self,
        input: TaskInput,
        *,
        mode: Optional[RolloutMode] = None,
        resources_id: Optional[str] = None,
        config: Optional[RolloutConfig] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Rollout:
        """Create rollout in memory, then XADD to Redis stream."""
        # Create rollout in the inner store (status="queuing")
        rollout = await self._inner.enqueue_rollout(
            input=input,
            mode=mode,
            resources_id=resources_id,
            config=config,
            metadata=metadata,
        )

        # Add to Redis stream for distributed consumption
        if self._redis:
            try:
                await self._redis.xadd(
                    self._stream_key,
                    {
                        "rollout_id": rollout.rollout_id,
                        "input_json": json.dumps(input) if not isinstance(input, str) else input,
                        "resources_id": resources_id or "",
                        "mode": mode or "",
                    },
                )
                logger.debug(f"XADD rollout {rollout.rollout_id} to {self._stream_key}")
            except Exception as e:
                logger.warning(
                    f"Failed to XADD rollout {rollout.rollout_id} to Redis: {e}. "
                    f"Rollout exists in memory but not on Redis stream."
                )
                raise RuntimeError(f"Failed to enqueue rollout {rollout.rollout_id} to Redis stream: {e}") from e

        return rollout

    async def enqueue_rollouts(
        self,
        requests: Sequence[EnqueueRolloutRequest],
    ) -> List[Rollout]:
        """Batch enqueue rollouts with Redis pipelining (F12)."""
        import json as _json

        # 1. Create all rollouts in inner (memory) store
        rollouts = []
        for req in requests:
            rollout = await self._inner.enqueue_rollout(
                input=req.input,
                mode=req.mode,
                resources_id=req.resources_id,
                config=req.config,
                metadata=req.metadata,
            )
            rollouts.append(rollout)

        # 2. Pipeline all XADD commands in one round-trip
        if self._redis and rollouts:
            pipe = self._redis.pipeline()
            for rollout, req in zip(rollouts, requests):
                pipe.xadd(
                    self._stream_key,
                    {
                        "rollout_id": rollout.rollout_id,
                        "input_json": _json.dumps(req.input) if not isinstance(req.input, str) else req.input,
                        "resources_id": req.resources_id or "",
                        "mode": req.mode or "",
                    },
                )
            try:
                await pipe.execute()
            except Exception as e:
                logger.error(f"Failed to pipeline enqueue {len(rollouts)} rollouts: {e}")
                raise RuntimeError(f"Failed to pipeline enqueue {len(rollouts)} rollouts: {e}") from e

        return rollouts

    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
        """XREADGROUP from Redis stream, then create attempt via inner store."""
        if not self._redis:
            return await self._inner.dequeue_rollout(worker_id=worker_id)

        consumer_name = worker_id or f"consumer-{id(self)}"

        # Read one message from the stream
        try:
            result = await self._redis.xreadgroup(
                self._consumer_group,
                consumer_name,
                {self._stream_key: ">"},
                count=1,
                block=self._block_timeout_ms,
            )
        except Exception as e:
            logger.error(f"Redis XREADGROUP failed: {e}. No fallback in distributed mode.")
            return None

        if not result:
            return None

        # result format: [[stream_key, [(message_id, {field: value})]]]
        stream_messages = result[0][1]
        if not stream_messages:
            return None

        message_id, fields = stream_messages[0]
        rollout_id = fields.get("rollout_id")

        if not rollout_id:
            logger.warning(f"Dequeued message without rollout_id: {fields}")
            await self._redis.xack(self._stream_key, self._consumer_group, message_id)
            return None

        # Use the inner store's dequeue to create the attempt
        # The rollout is already in memory from enqueue_rollout()
        attempted = await self._inner.dequeue_rollout(worker_id=worker_id)

        if attempted and attempted.rollout.rollout_id != rollout_id:
            logger.warning(
                f"Redis delivered rollout {rollout_id} but inner store "
                f"returned {attempted.rollout.rollout_id}. Possible desync."
            )

        if attempted:
            # Store message_id for later ACK
            attempted.rollout._redis_message_id = message_id  # type: ignore[attr-defined]
            logger.debug(f"Dequeued rollout {rollout_id} for worker {consumer_name}")
        else:
            # Rollout was in Redis but not in memory — ACK it to prevent redelivery
            await self._redis.xack(self._stream_key, self._consumer_group, message_id)
            logger.warning(f"Rollout {rollout_id} in Redis stream but not in memory store")

        return attempted

    async def start_rollout(
        self,
        input: TaskInput,
        *,
        mode: Optional[RolloutMode] = None,
        resources_id: Optional[str] = None,
        config: Optional[RolloutConfig] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> AttemptedRollout:
        """Create rollout + attempt immediately (bypasses queue)."""
        return await self._inner.start_rollout(
            input=input,
            mode=mode,
            resources_id=resources_id,
            config=config,
            metadata=metadata,
        )

    async def wait_for_rollouts(
        self,
        *,
        rollout_ids: Optional[List[str]] = None,
        timeout: Optional[float] = None,
    ) -> List[Rollout]:
        """Wait for rollouts to reach terminal status."""
        return await self._inner.wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)

    # =========================================================================
    # Rollout/Attempt updates — ACK on terminal status
    # =========================================================================

    async def update_rollout(
        self,
        rollout_id: str,
        *,
        status: Optional[RolloutStatus] = None,
        input: TaskInput | Unset = UNSET,
        resources_id: Optional[str] | Unset = UNSET,
        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
    ) -> Rollout:
        """Update rollout, ACK Redis message on terminal status."""
        rollout = await self._inner.update_rollout(
            rollout_id,
            status=status,
            input=input,
            resources_id=resources_id,
            metadata=metadata,
        )

        # ACK the Redis message when rollout reaches terminal status
        if self._redis and status and status in ("succeeded", "failed", "cancelled"):
            msg_id = getattr(rollout, "_redis_message_id", None)
            if msg_id:
                try:
                    await self._redis.xack(self._stream_key, self._consumer_group, msg_id)
                    logger.debug(f"XACK rollout {rollout_id} (status={status})")
                except Exception as e:
                    logger.warning(f"Failed to XACK rollout {rollout_id}: {e}")

        return rollout

    async def update_attempt(
        self,
        rollout_id: str,
        attempt_id: str,
        *,
        status: Optional[AttemptStatus] = None,
        worker_id: Optional[str] | Unset = UNSET,
        heartbeat: bool = False,
    ) -> Attempt:
        """Update attempt via inner store."""
        return await self._inner.update_attempt(
            rollout_id,
            attempt_id,
            status=status,
            worker_id=worker_id,
            heartbeat=heartbeat,
        )

    # =========================================================================
    # All other methods — delegate to InMemoryLightningStore
    # =========================================================================

    async def add_span(self, span: Span) -> Optional[Span]:
        return await self._inner.add_span(span)

    async def add_otel_span(
        self,
        rollout_id: str,
        attempt_id: str,
        readable_span: Any,
    ) -> Optional[Span]:
        return await self._inner.add_otel_span(rollout_id, attempt_id, readable_span)

    async def query_rollouts(self, **kwargs: Any) -> PaginatedResult[Rollout]:
        return await self._inner.query_rollouts(**kwargs)

    async def query_attempts(self, rollout_id: str, **kwargs: Any) -> PaginatedResult[Attempt]:
        return await self._inner.query_attempts(rollout_id, **kwargs)

    async def query_spans(self, rollout_id: str, attempt_id: str, **kwargs: Any) -> PaginatedResult[Span]:
        return await self._inner.query_spans(rollout_id, attempt_id, **kwargs)

    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
        return await self._inner.get_rollout_by_id(rollout_id)

    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
        return await self._inner.get_latest_attempt(rollout_id)

    async def add_resources(
        self,
        resources: ResourcesUpdate,
        *,
        default: bool = False,
    ) -> ResourcesUpdate:
        return await self._inner.add_resources(resources)

    async def update_resources(
        self,
        resources_id: str,
        resources: ResourcesUpdate,
        *,
        default: bool = False,
    ) -> ResourcesUpdate:
        return await self._inner.update_resources(resources_id, resources)

    async def get_latest_resources(
        self,
        resource_type: Optional[str] = None,
    ) -> Optional[NamedResources]:
        return await self._inner.get_latest_resources(resource_type=resource_type)

    async def get_resources_by_id(self, resources_id: str) -> Optional[NamedResources]:
        return await self._inner.get_resources_by_id(resources_id)

    def get_statistics(self) -> LightningStoreStatistics:
        return self._inner.get_statistics()

    async def register_worker(
        self,
        worker_id: str,
        *,
        status: WorkerStatus = "idle",
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Worker:
        return await self._inner.register_worker(worker_id, status=status, metadata=metadata)

    async def update_worker(
        self,
        worker_id: str,
        *,
        status: Optional[WorkerStatus] = None,
        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
    ) -> Worker:
        return await self._inner.update_worker(worker_id, status=status, metadata=metadata)

    async def get_worker(self, worker_id: str) -> Optional[Worker]:
        return await self._inner.get_worker(worker_id)

    async def query_workers(self, **kwargs: Any) -> PaginatedResult[Worker]:
        return await self._inner.query_workers(**kwargs)
