import asyncio
import logging
import time
from typing import Any, Dict, List, Optional, Tuple, AsyncGenerator
import valkey.asyncio as valkey
from valkey.exceptions import ResponseError

from pspf.utils.logging import get_logger

logger = get_logger("ValkeyBackend")
class ValkeyConnector:
    """
    Manages the connection pool to a Valkey (or Redis) server.
    """
    def __init__(self, host: Optional[str] = None, port: Optional[int] = None, password: Optional[str] = None, db: Optional[int] = None, valkey_settings: Optional[Any] = None):
        # Fallback to global settings if nothing provided
        from pspf.settings import settings as global_settings
        self._s = valkey_settings or global_settings.valkey
        
        self.host = host or self._s.HOST
        self.port = port or self._s.PORT
        self.password = password if password is not None else self._s.PASSWORD
        self.db = db if db is not None else self._s.DB
        self._pool: Optional[valkey.Valkey] = None

    async def connect(self) -> None:
        """
        Establish a connection pool to the database.
        
        Raises:
             ConnectionError: If connectivity fails.
        """
        if not self._pool:
            self._pool = valkey.Valkey(
                host=self.host,
                port=self.port,
                password=self.password,
                db=self.db,
                ssl=self._s.SSL,
                ssl_ca_certs=self._s.SSL_CA_CERTS,
                ssl_cert_reqs=self._s.SSL_CERT_REQS,
                decode_responses=True,
                socket_timeout=self._s.SOCKET_TIMEOUT,
                socket_connect_timeout=self._s.SOCKET_TIMEOUT
            )
            # Fail fast if connection is bad
            await self._pool.ping()
            logger.info(f"Connected to Valkey at {self.host}:{self.port} (SSL={self._s.SSL})")

    async def close(self) -> None:
        """
        Close the connection pool.
        """
        if self._pool:
            await self._pool.aclose()
            logger.info("Closed Valkey connection")
            self._pool = None

    def get_client(self) -> valkey.Valkey:
        """
        Get the active client instance.

        Returns:
            valkey.Valkey: The client instance.

        Raises:
            RuntimeError: If connect() has not been called.
        """
        if not self._pool:
            raise RuntimeError("ValkeyConnector is not connected. Call connect() first.")
        return self._pool

    async def __aenter__(self) -> "ValkeyConnector":
        await self.connect()
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        await self.close()


from .base import StreamingBackend

class ValkeyStreamBackend(StreamingBackend):
    """
    Handles all Stream-related operations on top of a ValkeyConnector.

    Attributes:
        connector (ValkeyConnector): The connection manager.
        stream_key (str): Name of the stream.
        group_name (str): Consumer group name.
        consumer_name (str): Unique consumer identifier.
    """
    def __init__(self, connector: ValkeyConnector, stream_key: str, group_name: str, consumer_name: str):
        self.connector = connector
        self._stream_key = stream_key
        self._group_name = group_name
        self.consumer_name = consumer_name
        self.dlq_stream_key = f"{stream_key}-dlq"
        self.retry_tracker_key = f"pspf:retries:{group_name}:{stream_key}"

    @property
    def stream_key(self) -> str:
        return self._stream_key

    @property
    def group_name(self) -> str:
        return self._group_name

    async def connect(self) -> None:
        await self.connector.connect()

    async def close(self) -> None:
        await self.connector.close()

    async def ping(self) -> bool:
        client = self.connector.get_client()
        await client.ping()
        return True

    async def ensure_group_exists(self, start_id: str = "0") -> None:
        """
        Idempotent creation of the consumer group.
        
        Args:
            start_id (str): ID to start from. Default is "0" (process all).
        """
        client = self.connector.get_client()
        try:
            # MKSTREAM ensures the stream is created if it doesn't exist
            await client.xgroup_create(self.stream_key, self.group_name, id=start_id, mkstream=True)
            logger.info(f"Created consumer group '{self.group_name}' on stream '{self.stream_key}'")
        except ResponseError as e:
            if "BUSYGROUP" in str(e):
                logger.debug(f"Consumer group '{self.group_name}' already exists.")
            else:
                raise e

    async def read_batch(self, count: int = 10, block_ms: int = 1000) -> List[Tuple[str, Dict[str, Any]]]:
        """
        Reads a batch of new messages ('&gt;') for this consumer group.
        
        Args:
            count (int): Max messages.
            block_ms (int): Block timeout in milliseconds.

        Returns:
            List[Tuple[str, Dict]]: List of (msg_id, payload).
        """
        client = self.connector.get_client()
        
        try:
            # Structure: [[stream_name, [(msg_id, data), ...]], ...]
            streams = await client.xreadgroup(
                groupname=self.group_name,
                consumername=self.consumer_name,
                streams={self.stream_key: ">"}, # ">" means new messages
                count=count,
                block=block_ms
            )

            if not streams:
                return []

            # We only read one stream, so take the first result
            stream_name, messages = streams[0]
            if not messages:
                return []
            
            # Deserialize nested JSON strings
            import json
            parsed_messages = []
            for msg_id, data in messages:
                parsed_data = {}
                for k, v in data.items():
                    if isinstance(v, str):
                        try:
                            parsed_data[k] = json.loads(v)
                        except (json.JSONDecodeError, TypeError):
                            parsed_data[k] = v
                    else:
                        parsed_data[k] = v
                parsed_messages.append((msg_id, parsed_data))
                
            return parsed_messages
        except Exception as e:
            logger.error(f"Error reading batch: {e}")
            raise

    async def get_retry_count(self, message_id: str) -> int:
        """
        Gets the current retry count for a message ID.
        """
        client = self.connector.get_client()
        count: Any = await client.hget(self.retry_tracker_key, message_id) # type: ignore
        if count:
            return int(count)
        return 0

    async def increment_retry_count(self, message_id: str) -> int:
        """
        Increments retry count for a message by 1.
        
        Args:
            message_id (str): The unique message ID.
            
        Returns:
             int: The new retry count.
        """
        client = self.connector.get_client()
        res: Any = await client.hincrby(self.retry_tracker_key, message_id, 1) # type: ignore
        return int(res)

    async def move_to_dlq(self, message_id: str, data: Dict[str, Any], error: str) -> None:
        """
        Moves a message to the DLQ stream and ACKs it in the main stream.
        
        Args:
             message_id (str): Message ID.
             data (Dict): Message payload.
             error (str): Error description.
        """
        client = self.connector.get_client()
        logger.warning(f"Moving message {message_id} to DLQ {self.dlq_stream_key}. Error: {error}")
        
        # Enrich data with error metadata
        dlq_data = data.copy()
        dlq_data["_error"] = str(error)
        dlq_data["_original_stream"] = self.stream_key
        dlq_data["_original_msg_id"] = message_id
        dlq_data["_moved_timestamp"] = str(time.time())
        
        # Add to DLQ
        await client.xadd(self.dlq_stream_key, dlq_data)
        
        # ACK in original stream to remove from Pending Entries List (PEL)
        await client.xack(self.stream_key, self.group_name, message_id)
        
        # Clean up retry tracker
        await client.hdel(self.retry_tracker_key, message_id) # type: ignore

    async def ack_batch(self, message_ids: List[str]) -> None:
        """
        Acknowledges a batch of message IDs and cleans up retry counters.
        
        Args:
            message_ids (List[str]): IDs to ACK.
        """
        if not message_ids:
            return
        
        client = self.connector.get_client()
        
        # Pipeline the XACK and retry cleanup
        async with client.pipeline() as pipe:
            pipe.xack(self.stream_key, self.group_name, *message_ids)
            pipe.hdel(self.retry_tracker_key, *message_ids)
            await pipe.execute()
        
        logger.debug(f"ACKed {len(message_ids)} messages")

    async def add_event(self, data: Dict[str, Any], max_len: Optional[int] = None) -> str:
        """
        Appends an event to the stream.
        
        Args:
             data (Dict): The payload.
             max_len (Optional[int]): Max stream length.
             
        Returns:
             str: The generated message ID.
        """
        # valkey-python requires primitive types (bytes, str, int, float)
        # We need to flatten or serialize the values if they are complex
        safe_data = {}
        for k, v in data.items():
            if isinstance(v, (dict, list, bool)) or v is None:
                import json
                safe_data[k] = json.dumps(v)
            else:
                safe_data[k] = v

        client = self.connector.get_client()
        msg_id: Any = await client.xadd(self.stream_key, safe_data, maxlen=max_len)
        return str(msg_id)

    async def claim_stuck_messages(self, min_idle_time_ms: int = 60000, count: int = 10) -> List[Tuple[str, Dict[str, Any]]]:
        """
        Auto-claims pending messages from crashed consumers.
        Using XAUTOCLAIM (Redis 6.2+ / Valkey) for efficiency.
        
        Args:
            min_idle_time_ms (int): Min idle time to consider a message 'stuck'.
            count (int): Max messages to claim.
            
        Returns:
             List[Tuple[str, Dict]]: Claimed messages.
        """
        client = self.connector.get_client()
        
        try:
            # xautoclaim(name, group, consumer, min_idle_time, start_id='0-0', count=None, justid=False)
            # Returns: (next_start_id, messages)
            # messages is a list of (msg_id, data)
            current_id = "0-0"
            all_messages: List[Tuple[str, Dict[str, Any]]] = []
            
            # We might need to loop if we really want to grab 'count' messages, 
            # but XAUTOCLAIM returns a batch usually. We'll just take one batch for simplicity to avoid implementation complexity.
            # XAUTOCLAIM returns (next_id, messages, [deleted_ids]) depending on version/client
            result = await client.xautoclaim(
                name=self.stream_key,
                groupname=self.group_name,
                consumername=self.consumer_name,
                min_idle_time=min_idle_time_ms,
                start_id=current_id,
                count=count
            )
            next_id = result[0]
            messages = result[1]
            
            if messages:
                logger.warning(f"Consumer {self.consumer_name} claimed {len(messages)} stuck messages.")
                from typing import cast
                # xautoclaim returns list of (msg_id, data)
                # Ensure data is deserialized if it was serialized? 
                # xautoclaim returns raw data just like xreadgroup
                # We need to deserialize it here too!
                raw_messages = cast(List[Tuple[str, Dict[str, Any]]], messages)
                
                # Deserialize nested JSON strings
                import json
                parsed_messages = []
                for msg_id, data in raw_messages:
                    parsed_data = {}
                    for k, v in data.items():
                        if isinstance(v, str):
                            try:
                                parsed_data[k] = json.loads(v)
                            except (json.JSONDecodeError, TypeError):
                                parsed_data[k] = v
                        else:
                            parsed_data[k] = v
                    parsed_messages.append((msg_id, parsed_data))
                return parsed_messages
            return []
            return []
        except Exception as e:
            logger.error(f"Error during XAUTOCLAIM: {e}")
            return []

    async def get_pending_info(self) -> Dict[str, Any]:
        """
        Retrieves lag and pending consumer group info.
        """
        client = self.connector.get_client()
        info = {
            "pending": 0,
            "lag": 0,
            "consumers": 0
        }
        try:
            # 1. Get Pending Count (Fast)
            pending_summary = await client.xpending(self.stream_key, self.group_name)
            info["pending"] = pending_summary["pending"]

            # 2. Get Group Info for Lag (slower, but accurate in Redis 7/Valkey)
            # XINFO GROUPS key
            groups = await client.xinfo_groups(self.stream_key)
            for g in groups:
                 if g["name"] == self.group_name:
                     info["lag"] = g.get("lag", 0)
                     info["consumers"] = g.get("consumers", 0)
                     break
            
            return info
        except Exception as e:
            logger.error(f"Error getting pending info: {e}")
            return info
