"""Implements a Redis Connection"""

import asyncio
import functools
import logging
import os
import typing as tp
import urllib.parse
import uuid
from concurrent.futures import ThreadPoolExecutor

import can_ada
import redis.asyncio as redis
from redis import RedisError, ResponseError

from ....conf import config
from ....exceptions import ConnectionError, PublishError, RequeueMessage
from ....models import Envelope, IncomingMessageProtocol
from .. import BaseAsyncConnection, is_task

log = logging.getLogger(__name__)

VHOST = os.environ.get("REDIS_VHOST") or os.environ.get("REDIS_DB", "0")


class Connection(BaseAsyncConnection):
    """Async Redis Connection wrapper."""

    _lock: asyncio.Lock | None = None
    instance_by_vhost: dict[str, "Connection | None"] = {}

    def __init__(
        self,
        username: str | None = None,
        password: str | None = None,
        host: str | None = None,
        port: int | None = None,
        db: int | None = None,
        vhost: str = "",
        url: str | None = None,
        worker_threads: int = 2,
        prefetch_count: int = 1,
        requeue_delay: int = 3,
        **kwargs,
    ):
        """Initialize Redis connection.

        Args:
            username: Redis username
            password: Redis password
            host: Redis host
            port: Redis port
            url: Redis URL. It will override username, password, host and port
            vhost: Redis virtual host (it's used as db number string if db not present)
            db: Redis database number
            worker_threads: number of concurrent callback workers to use
            prefetch_count: how many messages to prefetch from the queue
            requeue_delay: how long to wait before re-queueing a message (seconds)
        """
        super().__init__()
        uname = username or os.environ.get(
            "REDIS_USERNAME", os.environ.get("REDIS_DEFAULT_USER", "")
        )
        passwd = password or os.environ.get(
            "REDIS_PASSWORD", os.environ.get("REDIS_DEFAULT_PASS", "")
        )
        host = host or os.environ.get("REDIS_HOST", "127.0.0.1")
        port = port or int(os.environ.get("REDIS_PORT", "6379"))
        db = db or int(os.environ.get("REDIS_DB", "0"))
        # URL encode credentials and vhost to prevent injection
        vhost = vhost or VHOST or str(db)
        self.vhost = vhost
        username = urllib.parse.quote(uname, safe="")
        password = urllib.parse.quote(passwd, safe="")
        host = urllib.parse.quote(host, safe="")
        # URL for connection
        url = url or os.environ.get("REDIS_URL")
        if not url:
            # Build the URL based on what is available
            if username and password:
                url = f"redis://{username}:{password}@{host}:{port}/{vhost}?protocol=3"
            elif password:
                url = f"redis://:{password}@{host}:{port}/{vhost}?protocol=3"
            elif username:
                url = f"redis://{username}@{host}:{port}/{vhost}?protocol=3"
            else:
                url = f"redis://{host}:{port}/{vhost}?protocol=3"
        else:
            parsed = can_ada.parse(url)
            url = f"redis://{parsed.username}:{parsed.password}@{parsed.host}{parsed.pathname}{parsed.search}"

        self._url = url
        self._connection: redis.Redis | None = None
        self.prefetch_count = prefetch_count
        self.requeue_delay = requeue_delay
        self.queues: dict[str, dict] = {}
        self.consumers: dict[str, dict[str, tp.Any]] = {}
        self.executor = ThreadPoolExecutor(max_workers=worker_threads)
        self._instance_lock: asyncio.Lock | None = None

        self._delimiter = config.backend_config.topic_delimiter
        self._exchange = config.backend_config.namespace

    async def __aenter__(self, **kwargs) -> "Connection":
        await self.connect(**kwargs)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
        await self.disconnect()
        return False

    @property
    def lock(self) -> asyncio.Lock:
        """Lazy instance lock."""
        if self._instance_lock is None:
            self._instance_lock = asyncio.Lock()
        return self._instance_lock

    @classmethod
    def _get_class_lock(cls) -> asyncio.Lock:
        """Ensure the class lock is bound to the current running loop."""
        if cls._lock is None:
            cls._lock = asyncio.Lock()
        return cls._lock

    def build_topic_key(self, topic: str) -> str:
        return f"{self._exchange}:{topic}"

    @property
    def is_connected_event(self) -> asyncio.Event:
        """Lazily create the event in the current running loop."""
        if self._is_connected_event is None:
            self._is_connected_event = asyncio.Event()
        return self._is_connected_event

    @property
    def connection(self) -> "redis.Redis":
        """Get the connection object.

        Raises:
            ConnectionError: If not connected
        """
        if not self._connection:
            raise ConnectionError("Connection not initialized. Call connect() first.")
        return self._connection

    async def connect(self, **kwargs) -> "Connection":
        """Establish Redis connection.

        Args:

        Raises:
            ConnectionError: If connection fails
            asyncio.TimeoutError: If connection times out
        """
        async with self.lock:
            if self.instance_by_vhost.get(self.vhost) and self.is_connected():
                return self.instance_by_vhost[self.vhost]
            try:
                # Parsing URL for logging (removing credentials)
                log.info("Establishing Redis connection to %s", self._url.split("@")[-1])
                # protobunny sends raw bytes with protobuf serialized payloads
                kwargs.pop("decode_responses", None)
                # Using from_url handles connection pooling automatically
                self._connection = redis.from_url(
                    self._url,
                    decode_responses=False,
                    **kwargs,
                )

                await asyncio.wait_for(self._connection.ping(), timeout=30)
                self.is_connected_event.set()
                log.info("Successfully connected to Redis")
                self.instance_by_vhost[self.vhost] = self
                return self

            except asyncio.TimeoutError:
                log.error("Redis connection timeout after %.1f seconds", 30)
                self.is_connected_event.clear()
                self._connection = None
                raise
            except Exception as e:
                if self._connection:
                    await self._connection.aclose()

                self.is_connected_event.clear()
                self._connection = None
                log.exception("Failed to establish Redis connection")
                raise ConnectionError(f"Failed to connect to Redis: {e}") from e

    async def disconnect(self, timeout: float = 10.0) -> None:
        """Close Redis connection and cleanup resources.

        Args:
            timeout: Maximum time to wait for cleanup (seconds)
        """
        async with self.lock:
            if not self.is_connected():
                log.debug("Already disconnected from Redis")
                return

            try:
                log.info("Closing Redis connection")
                # In Redis, consumers are local asyncio Tasks.
                # We cancel them here. Note: Redis doesn't have "exclusive queues"
                # that auto-delete, so we just clear our local registry.
                for tag, consumer in self.consumers.items():
                    task = consumer["task"]
                    try:
                        task.cancel()
                        # We give the task a moment to wrap up if needed
                        await asyncio.sleep(0)  # force context switching
                        await asyncio.wait([task], timeout=2.0)
                    except Exception as e:
                        log.warning("Error stopping Redis consumer %s: %s", tag, e)

                # Shutdown Thread Executor (if used for sync callbacks)
                self.executor.shutdown(wait=False, cancel_futures=True)

                # Close the Redis Connection Pool
                if self._connection:
                    # aclose() closes the connection pool and all underlying connections
                    await asyncio.wait_for(self._connection.aclose(), timeout=timeout)

            except asyncio.TimeoutError:
                log.warning("Redis connection close timeout after %.1f seconds", timeout)
            except Exception:
                log.exception("Error during Redis disconnect")
            finally:
                # Reset state
                self._connection = None
                self.queues.clear()  # (Local queue metadata)
                self.consumers.clear()
                self.is_connected_event.clear()
                # Remove from registry
                Connection.instance_by_vhost.pop(self.vhost, None)
                log.info("Redis connection closed")

    async def subscribe(self, topic: str, callback: tp.Callable, shared: bool = False) -> str:
        """Subscribe to Redis.

        Args:
            topic: The stream key/topic to subscribe to
            callback: Function to handle incoming messages.
            shared: If True, uses a shared consumer group (round-robin).
                    If False, uses a unique group for this instance and let redis pubsub manage the routing.

        Returns:
            A unique subscription tag used to stop the consumer later.
        """
        async with self.lock:
            if not self.is_connected():
                raise ConnectionError("Not connected to Redis")
            queue = await self.setup_queue(topic, shared)
            ready_event = asyncio.Event()
            stop_event = asyncio.Event()
            if queue["is_shared"]:
                assert (
                    queue["mechanism"] == "stream"
                ), f"Invalid queue mechanism: {queue['mechanism']}"

                # create the asyncio task for the consumer loop
                log.debug("Subscribing callback to Redis PubSub topic %s", queue["key"])
                task = asyncio.create_task(
                    self._tasks_consumer_loop(
                        queue["key"],
                        queue["group_name"],
                        queue["tag"],
                        callback,
                        ready_event,
                        stop_event,
                    )
                )

                # Wait for the loop to signal it has performed the first check
                await asyncio.wait_for(ready_event.wait(), timeout=1.0)
                # register the topic in the registry
                # We store the task so we can cancel it in 'disconnect'
                self.consumers[queue["tag"]] = {
                    "task": task,
                    "key": queue["key"],
                    "topic": topic,
                    "stop_event": stop_event,
                }
                log.info("Redis consumer %s subscribed to %s", queue["tag"], queue["key"])
                return queue["tag"]
            else:
                # Using redis.PubSub
                assert (
                    queue["mechanism"] == "pubsub"
                ), f"Invalid queue mechanism: {queue['mechanism']}"
                # create the asyncio task for the consumer loop
                task = asyncio.create_task(
                    self._pubsub_consumer_loop(ready_event, callback, queue, stop_event)
                )
                # Wait for the loop to signal it has performed the first check
                await asyncio.wait_for(ready_event.wait(), timeout=1.0)
                # register the topic in the registry
                # We store the task so we can cancel it in 'disconnect'
                self.consumers[queue["tag"]] = {
                    "task": task,
                    "key": queue["key"],
                    "topic": topic,
                    "stop_event": stop_event,
                }
                log.info("Redis consumer %s subscribed to %s", queue["tag"], topic)
                return queue["tag"]

    async def _pubsub_consumer_loop(
        self,
        ready_event: asyncio.Event,
        callback: tp.Callable,
        queue: dict,
        stop_event: asyncio.Event,
    ):
        pubsub = self._connection.pubsub()
        callback = functools.partial(self._on_message_pubsub, queue["key"], callback)
        log.debug("Subscribing callback to Redis PubSub topic %s", queue["key"])
        if any(char in queue["key"] for char in ("*", "[", "?")):
            await pubsub.psubscribe(queue["key"])
        else:
            await pubsub.subscribe(queue["key"])
        ready_event.set()

        async for message in pubsub.listen():
            if message.get("type") not in ("message", "pmessage"):
                continue
            routing_key = (
                message["channel"].decode().removeprefix(f"{self._exchange}{self._delimiter}")
            )
            envelope = Envelope(
                body=message["data"],
                routing_key=routing_key,
            )
            if asyncio.iscoroutinefunction(callback):
                await callback(envelope)
            else:
                asyncio.run_coroutine_threadsafe(callback(envelope), self._loop)
            if stop_event.is_set():
                await pubsub.unsubscribe()
                await pubsub.close()
                break

    async def setup_queue(self, topic: str, shared: bool = False) -> dict:
        """Set up a Redis Stream and Consumer Group for tasks if shared, otherwise use Pub/Sub.

        Args:
            topic: The stream key / routing key (queue name)
            shared: If True, uses a fixed group name for round-robin.
                    If False, just creates metadata to use to subscribe the callback

        Returns:
            The name of the consumer group to use
        """
        if not self.is_connected():
            raise ConnectionError("Not connected to Redis")

        key = self.build_topic_key(topic)
        tag = f"consumer_{uuid.uuid4().hex[:8]}"

        if shared:
            group_name = "shared_group"
            log.debug(
                "Setting up Redis Stream %s with group %s (shared=%s)", key, group_name, shared
            )
            try:
                # Create the Consumer Group
                # MKSTREAM: Creates the stream key if it doesn't exist
                # id='$': Only read new messages arriving from now on
                await self._connection.xgroup_create(
                    name=key, groupname=group_name, id="$", mkstream=True
                )
            except Exception as e:
                if "BUSYGROUP" not in str(e):
                    log.error("Failed to setup Redis Stream group: %s", e)
                    raise

            queue_meta = {
                "mechanism": "stream",
                "key": key,
                "group_name": group_name,
                "tag": tag,
                "topic": topic,
                "is_shared": True,
            }
        else:
            # NORMAL QUEUE LOGIC (Pub/Sub)
            # No server-side 'declaration' needed for Pub/Sub.
            # We just prepare the metadata to use PSUBSCRIBE.
            queue_meta = {
                "mechanism": "pubsub",
                "tag": tag,
                "topic": topic,
                "is_shared": False,
                "group_name": f"fanout_{uuid.uuid4().hex[:8]}",
                "key": key,
            }
        self.queues[key] = queue_meta
        return queue_meta

    async def _tasks_consumer_loop(
        self,
        key: str,
        group_name: str,
        consumer_id: str,
        callback: tp.Callable,
        ready_event: asyncio.Event,
        stop_event: asyncio.Event,
    ):
        """Internal loop to read messages from Redis Stream."""
        queue_meta = self.queues.get(key)
        if not queue_meta:
            log.warning("No queue found for key %s", key)
            return

        # Signal that the loop has started before entering the blocking section
        ready_event.set()
        await asyncio.sleep(0)
        is_mock = hasattr(self._connection, "server")  # Simple way to detect fakeredis

        while not stop_event.is_set():
            try:
                # XREADGROUP: block=0 means wait indefinitely for new messages
                # ">" means "read only new messages never delivered to others"

                log.debug("Consumer loop for %s started", consumer_id)
                response = await self._connection.xreadgroup(
                    groupname=group_name,
                    consumername=consumer_id,
                    streams={key: ">"},
                    count=self.prefetch_count or 10,
                    block=None
                    if is_mock
                    else 1000,  # Block for 1 second then loop (allows for clean shutdown)
                )
                if not response:
                    continue
                messages = response.get(key.encode())[0]
                for msg_id, payload in messages:
                    await self._on_message_task(key, group_name, msg_id, payload, callback)
            except asyncio.CancelledError:
                break
            except Exception as e:
                if "UNBLOCKED" not in str(e):
                    log.error("Error in Redis consumer loop: %s", e)
                    await asyncio.sleep(1)  # Backoff on error

    async def _on_message_pubsub(
        self, topic: str, callback, envelope: IncomingMessageProtocol
    ) -> None:
        if isinstance(envelope.body, str):
            envelope.body.encode()  # Ensure it's bytes if it was auto-decoded

        key = self.build_topic_key(topic)
        try:
            log.debug("Running callback for topic %s", key)
            if asyncio.iscoroutinefunction(callback):
                await callback(envelope)
            else:
                # Run sync callback in the executor
                await asyncio.get_event_loop().run_in_executor(self.executor, callback, envelope)
        except RequeueMessage:
            log.warning("Requeuing message on topic '%s' after RequeueMessage exception", key)
            await asyncio.sleep(self.requeue_delay)
            await self._connection.publish(key, envelope.body)
        except Exception as e:
            log.exception("Callback failed for topic %s", key)
            raise PublishError(f"Failed to publish message to topic {key}: {e}") from e

    async def publish(
        self,
        topic: str,
        message: "IncomingMessageProtocol",
        **kwargs,
    ) -> None:
        """
        Simulates Topic Exchange routing.
        """
        if not self.is_connected():
            raise ConnectionError("Not connected")

        topic_key = self.build_topic_key(topic)
        is_shared = is_task(topic)

        if is_shared:
            log.debug("Publishing message to Redis Stream: %s", topic_key)
            payload = {
                "body": message.body,
                "correlation_id": message.correlation_id,
                "topic": topic,  # add the topic here to implement topic exchange patterns
            }
            await self._connection.xadd(name=topic_key, fields=payload, maxlen=1000)
            if config.log_task_in_redis:
                # Tasks messages go to streams but the logger do a simple pubsub psubscription to <prefix>.*
                # Send the message to the same topic with redis.publish so it appears there
                # Note: this should be used carefully (e.g. only for debugging)
                # as it doubles the network calls when publishing tasks
                log.debug("Publishing message to topic: %s for logger", topic_key)
                await self._connection.publish(topic_key, message.body)
        else:
            # For fan-out, we use native Redis Pub/Sub.
            # Redis will automatically match this to any PSUBSCRIBE patterns.
            log.debug("Publishing message to topic: %s", topic_key)
            await self._connection.publish(topic_key, message.body)

    async def _on_message_task(
        self, stream_key: str, group_name: str, msg_id: str, payload: dict, callback: tp.Callable
    ):
        """Wraps the user callback."""
        # Response is not decoded because we use bytes. But the keys will be bytes as well
        normalized_payload = {
            k.decode() if isinstance(k, bytes) else k: v for k, v in payload.items()
        }
        body = normalized_payload.get("body", b"")
        topic = normalized_payload.get("topic", b"").decode()
        correlation_id = normalized_payload.get("correlation_id", b"").decode()
        if isinstance(body, str):
            body = body.encode()  # Ensure it's bytes if it was auto-decoded
        # Create the Envelope
        envelope = Envelope(
            body=body,
            correlation_id=correlation_id,
            routing_key=topic.decode() if isinstance(topic, bytes) else topic,
        )
        try:
            if asyncio.iscoroutinefunction(callback):
                await callback(envelope)
            else:
                # Run sync callback in the executor
                await asyncio.get_event_loop().run_in_executor(self.executor, callback, envelope)

            # Manual ACK: Message is processed, tell Redis to remove from Pending List
            await self._connection.xack(stream_key, group_name, msg_id)
        except RequeueMessage:
            log.warning("Requeuing message on topic '%s' after RequeueMessage exception", topic)
            # In Redis, to "requeue" so it's processed again:
            # XADD again and
            # XACK the current ID
            await asyncio.sleep(self.requeue_delay)
            await self._connection.xadd(name=stream_key, fields=payload)
            await self._connection.xack(stream_key, group_name, msg_id)
        except Exception:
            log.exception("Callback failed for message %s", msg_id)
            # Avoid poisoning messages
            # Note: In Redis, if you don't XACK, the message stays in the
            # Pending Entry List (PEL) for retry logic.
            await self._connection.xack(stream_key, group_name, msg_id)

    async def unsubscribe(self, tag: str, if_unused: bool = True, if_empty: bool = True) -> None:
        task_to_cancel = None
        async with self.lock:
            if tag not in self.consumers:
                return
            consumer_info = self.consumers.pop(tag)
            if consumer_info:
                task_to_cancel = consumer_info["task"]
                stop_event = consumer_info["stop_event"]
                key = consumer_info["key"]
                # Stop the local asyncio loop
                stop_event.set()

        if task_to_cancel:
            try:
                # wait for the task to stop (outside the lock otherwise will deadlock for python 3.10/3.11
                task_to_cancel.cancel()
                await asyncio.sleep(0)
                await asyncio.wait_for(task_to_cancel, timeout=1.0)
            except (asyncio.TimeoutError, asyncio.CancelledError):
                pass

        async with self.lock:
            queue_meta = self.queues.get(key)
            if not queue_meta:
                return

            group_name = queue_meta["group_name"]
            is_shared = queue_meta["is_shared"]

            # Tell Redis to remove THIS specific consumer from the group
            if is_shared:
                try:
                    await self._connection.xgroup_delconsumer(key, group_name, tag)
                    # Optionally delete the actual stream data
                    if not if_empty or (await self._connection.xlen(key) == 0):
                        await self._connection.delete(key)
                        self.queues.pop(key, None)
                    log.debug("Deleted Redis consumer %s from group %s", tag, group_name)
                except Exception as e:
                    log.warning("Could not delete Redis consumer %s: %s", tag, e)

            # Handle non shared queues
            else:
                try:
                    log.info("Unsubscribe from %s (group=%s, consumer=%s)", key, group_name, tag)
                    assert (
                        queue_meta["mechanism"] == "pubsub"
                    ), f"Invalid queue mechanism: {queue_meta['mechanism']}"
                    pubsub = self._connection.pubsub()
                    # Use psubscribe to support wildcards like 'user.*'
                    if any(char in key for char in ("*", "[", "?")):
                        await pubsub.punsubscribe(queue_meta["key"])
                    else:
                        await pubsub.unsubscribe(queue_meta["key"])
                    self.queues.pop(key, None)

                    # Delete stream if_empty=False
                    if not if_empty:
                        log.info("Deleting subscription %s even if not empty", key)
                        await self._connection.delete(key)
                except ResponseError as e:
                    log.error("Error checking during unsubscribe: %s", e)
                    raise
                except RedisError as e:
                    log.error("Redis error during unsubscribe: %s", e)

    async def purge(self, topic: str, reset_groups: bool = False) -> None:
        """Empty a Redis Stream and optionally clear all consumer groups.

        Args:
            topic: The stream/topic name to purge
            reset_groups: If True, deletes all consumer groups (resets consumer count to 0)
        """
        async with self.lock:
            if not self.is_connected():
                raise ConnectionError("Not connected to Redis")

            stream_key = self.build_topic_key(topic)
            log.info("Purging Redis stream '%s' (reset_groups=%s)", stream_key, reset_groups)

            try:
                # Clear all messages
                await self._connection.xtrim(stream_key, maxlen=0, approximate=False)
                if reset_groups:
                    try:
                        await self.reset_stream_groups(stream_key)
                        await self._connection.delete(stream_key)
                    except ResponseError as e:
                        # Ignore error if the stream key doesn't exist yet
                        if "no such key" not in str(e).lower():
                            raise e

                # Clear local metadata cache
                if stream_key in self.queues:
                    self.queues.pop(stream_key)

            except Exception as e:
                log.error("Failed to purge Redis stream %s: %s", stream_key, e)
                raise ConnectionError(f"Failed to purge topic: {e}") from e

    async def get_message_count(self, topic: str) -> int:
        """Get the number of messages in the Redis Stream.

        Args:
            topic: The stream topic name

        Returns:
            Number of entries currently in the stream.
        """
        if not self.is_connected():
            raise ConnectionError("Not connected to Redis")

        stream_key = self.build_topic_key(topic)
        log.debug("Getting message count for shared topic (stream) '%s'", stream_key)

        try:
            # XLEN returns the number of entries in a stream
            return await self._connection.xlen(stream_key)
        except Exception as e:
            log.error("Failed to get message count for %s: %s", stream_key, e)
            return 0

    async def get_consumer_count(self, topic: str) -> int:
        """Get the total number of consumers across all groups for a topic.

        Args:
            topic: The stream topic

        Returns:
            Total number of consumers
        """
        async with self.lock:
            if not self.is_connected():
                raise ConnectionError("Not connected to Redis")

            stream_key = self.build_topic_key(topic)
            total_consumers = 0

            try:
                # XINFO GROUPS returns a list of all consumer groups for this stream
                groups = await self._connection.xinfo_groups(stream_key)

                for group in groups:
                    # Each group dictionary contains a 'consumers' count
                    total_consumers += group.get("consumers", 0)

                return total_consumers

            except Exception as e:
                # If the stream doesn't exist, xinfo will raise an error
                if "no such key" in str(e).lower():
                    return 0
                log.error("Failed to get consumer count for %s: %s", stream_key, e)
                return 0

    async def reset_stream_groups(self, stream_key: str) -> None:
        """Hard reset: Deletes all consumer groups for a topic. To be used with caution."""
        if not self.is_connected():
            return
        try:
            # Get all groups for this stream
            groups = await self._connection.xinfo_groups(stream_key)
            for group in groups:
                group_name = group["name"]
                # Destroy the group (this removes all consumers inside it)
                await self._connection.xgroup_destroy(stream_key, group_name)
                log.info("Destroyed group %s on %s", group_name, stream_key)
        except Exception as e:
            if "no such key" in str(e).lower():
                return
            log.error("Failed to reset groups for %s: %s", stream_key, e)
