"""
Signal Bridge — Redis Stream Reader

Runs as an asyncio task within the workspace agent process.
Reads signals from a customer Redis stream using XREADGROUP,
buffers them, and sends batched signal_batch messages over
the existing WebSocket connection.

Features:
- Consumer group (resume from last ACK on reconnect)
- Buffer: flush every 500ms or 100 items or 1MB
- Reconnect burst backpressure: graduated ramp-up
- Throttle support (backend can send signal_throttle)
"""

import asyncio
import json
import logging
import time
from typing import Optional, Callable, Awaitable

logger = logging.getLogger(__name__)

# Buffer limits
FLUSH_INTERVAL_MS = 500
FLUSH_MAX_ITEMS = 100
FLUSH_MAX_BYTES = 1_048_576  # 1MB

# Reconnect burst ramp
RAMP_BATCH_SIZE_START = 50
RAMP_BATCH_SIZE_NORMAL = 100
RAMP_DELAY_START_MS = 200
RAMP_DELAY_NORMAL_MS = 100
RAMP_TRANSITION_BATCH = 5


class SignalBridge:
    """
    Reads from a Redis stream and forwards signals over WebSocket.
    Lifecycle managed by WorkspaceConnection via start/stop commands.
    """

    def __init__(
        self,
        redis_url: str,
        stream_key: str,
        consumer_group: str,
        agent_id: str,
        send_fn: Callable[[str], Awaitable[None]],  # ws.send
    ):
        self._redis_url = redis_url
        self._stream_key = stream_key
        self._consumer_group = consumer_group
        self._consumer_name = f"bridge_{agent_id}"
        self._agent_id = agent_id
        self._send_fn = send_fn

        self._running = False
        self._task: Optional[asyncio.Task] = None
        self._buffer: list = []
        self._buffer_bytes = 0
        self._throttle_ms = 0  # Backend-requested delay
        self._batch_count = 0  # For ramp-up tracking
        self._messages_processed = 0
        self._started_at: Optional[float] = None
        self._last_error: Optional[str] = None

    async def start(self):
        """Start the bridge as an asyncio task."""
        if self._running:
            logger.warning(f"[BRIDGE] Already running for agent {self._agent_id}")
            return

        self._running = True
        self._batch_count = 0
        self._started_at = time.time()
        self._last_error = None
        self._task = asyncio.create_task(self._run())
        logger.info(f"[BRIDGE] Started for agent {self._agent_id} | stream={self._stream_key}")

    async def stop(self):
        """Stop the bridge gracefully."""
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        logger.info(f"[BRIDGE] Stopped for agent {self._agent_id}")

    def set_throttle(self, delay_ms: int):
        """Backend requests throttling."""
        self._throttle_ms = max(0, delay_ms)
        logger.info(f"[BRIDGE] Throttle set to {delay_ms}ms")

    @property
    def is_running(self) -> bool:
        return self._running

    def get_status(self) -> dict:
        """Return real-time status for process manager UI."""
        uptime = int(time.time() - self._started_at) if self._started_at else 0
        return {
            'agent_id': self._agent_id,
            'stream_key': self._stream_key,
            'running': self._running,
            'messages_processed': self._messages_processed,
            'buffer_size': len(self._buffer),
            'uptime_seconds': uptime,
            'throttle_ms': self._throttle_ms,
            'last_error': self._last_error,
        }

    async def _run(self):
        """Main loop: read from Redis stream, buffer, flush to WebSocket."""
        try:
            import redis.asyncio as aioredis
        except ImportError:
            logger.error("[BRIDGE] redis package not installed. pip install redis")
            self._running = False
            return

        try:
            # Handle SSL for rediss:// URLs (cloud Redis providers)
            kwargs = {'decode_responses': True}
            if self._redis_url.startswith('rediss://'):
                kwargs['ssl_cert_reqs'] = None  # Skip cert verification

            client = aioredis.from_url(self._redis_url, **kwargs)

            # Ensure consumer group exists
            try:
                await client.xgroup_create(
                    self._stream_key, self._consumer_group, id='0', mkstream=True
                )
                logger.info(f"[BRIDGE] Created consumer group: {self._consumer_group}")
            except Exception:
                pass  # Group already exists

            logger.info(f"[BRIDGE] Listening on {self._stream_key}...")

            flush_task = asyncio.create_task(self._flush_loop())

            try:
                while self._running:
                    # Apply throttle
                    if self._throttle_ms > 0:
                        await asyncio.sleep(self._throttle_ms / 1000)

                    # Graduated batch size for reconnect burst
                    if self._batch_count < RAMP_TRANSITION_BATCH:
                        count = RAMP_BATCH_SIZE_START
                    else:
                        count = RAMP_BATCH_SIZE_NORMAL

                    # Read from stream with consumer group
                    results = await client.xreadgroup(
                        groupname=self._consumer_group,
                        consumername=self._consumer_name,
                        streams={self._stream_key: '>'},
                        count=count,
                        block=500,  # 500ms block timeout
                    )

                    if not results:
                        continue

                    for stream_name, messages in results:
                        for msg_id, fields in messages:
                            # Parse payload first so we can extract metadata from it
                            payload = self._parse_payload(fields)

                            # Extract signal_type: check top-level fields first, then inside payload
                            signal_type = (
                                fields.get('type')
                                or fields.get('signal_type')
                                or payload.get('signal_type')
                                or payload.get('type')
                                or payload.get('event')
                                or 'unknown'
                            )

                            signal_data = {
                                'source_id': msg_id,
                                'signal_type': signal_type,
                                'payload': payload,
                                'flow_key': fields.get('flow_key') or payload.get('flow_key'),
                                'priority': int(fields.get('priority', payload.get('priority', 0))),
                            }
                            signal_json = json.dumps(signal_data)
                            self._buffer.append(signal_data)
                            self._buffer_bytes += len(signal_json)

                            # Ack the message
                            await client.xack(self._stream_key, self._consumer_group, msg_id)
                            self._messages_processed += 1

                    # Force flush if buffer limits hit
                    if len(self._buffer) >= FLUSH_MAX_ITEMS or self._buffer_bytes >= FLUSH_MAX_BYTES:
                        await self._flush()

            finally:
                flush_task.cancel()
                # Final flush
                if self._buffer:
                    await self._flush()
                await client.aclose()

        except asyncio.CancelledError:
            logger.info(f"[BRIDGE] Task cancelled for agent {self._agent_id}")
        except Exception as e:
            self._last_error = str(e)
            logger.error(f"[BRIDGE] Error: {e}", exc_info=True)
        finally:
            self._running = False

    async def _flush_loop(self):
        """Periodic flush timer (every 500ms)."""
        try:
            while self._running:
                await asyncio.sleep(FLUSH_INTERVAL_MS / 1000)
                if self._buffer:
                    await self._flush()
        except asyncio.CancelledError:
            pass

    async def _flush(self):
        """Send buffered signals to backend via WebSocket."""
        if not self._buffer:
            return

        batch = self._buffer[:FLUSH_MAX_ITEMS]
        self._buffer = self._buffer[FLUSH_MAX_ITEMS:]

        message = json.dumps({
            'type': 'signal_batch',
            'agent_id': self._agent_id,
            'bridge_id': f"bridge_{self._agent_id}",
            'signals': batch,
        })

        try:
            await self._send_fn(message)
            self._batch_count += 1
            self._buffer_bytes = sum(len(json.dumps(s)) for s in self._buffer)
            self._flush_failures = 0  # Reset on success

            # Apply reconnect burst delay
            if self._batch_count <= RAMP_TRANSITION_BATCH:
                await asyncio.sleep(RAMP_DELAY_START_MS / 1000)
            elif self._batch_count <= 20:
                await asyncio.sleep(RAMP_DELAY_NORMAL_MS / 1000)

            logger.debug(f"[BRIDGE] Flushed {len(batch)} signals (batch #{self._batch_count})")
        except Exception as e:
            self._flush_failures = getattr(self, '_flush_failures', 0) + 1
            if self._flush_failures >= 3:
                # WebSocket is dead — stop the bridge, let reconnect restart it
                logger.warning(f"[BRIDGE] WebSocket dead ({e}), stopping bridge for agent {self._agent_id}")
                self._buffer.clear()
                self._buffer_bytes = 0
                self._running = False
            else:
                # Put signals back into buffer for retry
                self._buffer = batch + self._buffer
                logger.error(f"[BRIDGE] Flush failed ({self._flush_failures}/3): {e}")

    @staticmethod
    def _parse_payload(fields: dict) -> dict:
        """Parse payload from Redis stream fields."""
        # Try 'payload' field first (JSON string)
        if 'payload' in fields:
            try:
                return json.loads(fields['payload'])
            except (json.JSONDecodeError, TypeError):
                return {'raw': fields['payload']}

        # Try 'data' field (common pattern in Redis streams)
        if 'data' in fields:
            try:
                return json.loads(fields['data'])
            except (json.JSONDecodeError, TypeError):
                return {'raw': fields['data']}

        # Filter out known meta fields, return rest as payload
        meta_keys = {'type', 'signal_type', 'flow_key', 'priority', 'source_id'}
        return {k: v for k, v in fields.items() if k not in meta_keys}
