"""
Transport Stations - Bridge between Pipeline and Transport

InputStation: Read from read_queue -> Chunk stream (supports audio + optional video)
OutputStation: Chunk stream -> Write to send_queue

Design principles:
- InputStation/OutputStation are pure format conversion layers
- NOT responsible for: connection management, flow control, turn management, latency tracking
- These common capabilities are unified by TransportBase's read_worker/send_worker

Vision Support:
- InputStation can optionally receive video frames via video_queue
- Visual context is injected into audio chunks via metadata["visual_context"]
- VAD-triggered frames take priority over heartbeat frames
"""

import asyncio
from collections import deque
from typing import AsyncIterator, Optional, Dict, Any, Union
from vixio.core.station import Station
from vixio.core.protocol import ProtocolBase
from vixio.core.chunk import Chunk, ChunkType
from vixio.providers.vision import ImageContent
from loguru import logger


class InputStation(Station):
    """
    Input format conversion Station - Pipeline starting point.
    
    Responsibilities (pure format conversion):
    1. Read raw data from Transport's read_queue
    2. Call Protocol to decode message
    3. Call Protocol to convert to Chunk
    4. Call AudioCodec to decode audio (if needed)
    5. Inject visual context into audio chunks (if video_queue provided)
    6. Output Chunk to Pipeline
    
    Vision Support:
    - Optionally receives video frames via video_queue
    - VAD-triggered frames stored as speech_context (high priority)
    - Heartbeat frames stored as latest_frame (low priority)
    - Visual context injected into audio chunks via metadata["visual_context"]
    
    NOT responsible for:
    - Connection management (handled by Transport's read_worker)
    - Turn management (handled by TransportBase framework)
    - Latency tracking (handled by TransportBase framework)
    """
    
    def __init__(
        self,
        session_id: str,
        read_queue: asyncio.Queue,      # From Transport - audio/control messages
        protocol: ProtocolBase,          # From Transport
        audio_codec: Optional[Any] = None,  # From Transport, can be None
        video_queue: Optional[asyncio.Queue] = None,  # From Transport - video frames
        name: str = "InputStation"
    ):
        """
        Args:
            session_id: Session ID
            read_queue: Transport's read queue (audio/control)
            protocol: Protocol handler
            audio_codec: Audio codec (can be None)
            video_queue: Transport's video queue (optional, for vision support)
            name: Station name
        """
        super().__init__(name=name)
        self._session_id = session_id
        self._read_queue = read_queue
        self._video_queue = video_queue
        self.protocol = protocol
        self.audio_codec = audio_codec
        self._running = True
        
        # Vision context management
        self._speech_context_frame: Optional[ImageContent] = None  # VAD-triggered frame
        self._latest_frame: Optional[ImageContent] = None  # Heartbeat frame
        self._video_task: Optional[asyncio.Task] = None
    
    async def process_chunk(self, chunk: Chunk) -> AsyncIterator[Chunk]:
        """
        InputStation is starting point, doesn't process upstream Chunks.
        
        Actual data generation is in _generate_chunks().
        """
        # If receiving upstream Chunk (shouldn't happen), pass through
        yield chunk
    
    async def _collect_video_frames(self) -> None:
        """
        Background task: Collect video frames from video_queue.
        
        Frame types:
        - trigger="vad_start": Speech context frame (high priority)
        - trigger="heartbeat": Latest frame (low priority, backup)
        """
        self.logger.info(f"Starting video frame collector for session {self._session_id[:8]}")
        
        try:
            while self._running:
                try:
                    frame_data = await asyncio.wait_for(
                        self._video_queue.get(),
                        timeout=1.0
                    )
                except asyncio.TimeoutError:
                    continue
                
                if frame_data is None:
                    break
                
                # Parse frame data to ImageContent
                image = self._parse_video_frame(frame_data)
                if image is None:
                    continue
                
                # Store based on trigger type
                if image.trigger == "vad_start":
                    self._speech_context_frame = image
                    self.logger.debug(f"Stored VAD-triggered frame")
                else:
                    self._latest_frame = image
                    self.logger.debug(f"Updated latest frame (heartbeat)")
        
        except asyncio.CancelledError:
            self.logger.debug("Video frame collector cancelled")
        except Exception as e:
            self.logger.error(f"Error in video frame collector: {e}")
        finally:
            self.logger.info(f"Video frame collector ended for session {self._session_id[:8]}")
    
    def _parse_video_frame(self, frame_data: Dict[str, Any]) -> Optional[ImageContent]:
        """
        Parse video frame data to ImageContent.
        
        Expected format:
        {
            "data": bytes or base64 string,
            "trigger": "vad_start" | "heartbeat",
            "timestamp": float,
            "width": int,
            "height": int,
            "format": "jpeg" | "png"
        }
        """
        try:
            return ImageContent(
                image=frame_data.get("data", b""),
                mime_type=f"image/{frame_data.get('format', 'jpeg')}",
                source="camera",
                trigger=frame_data.get("trigger", "heartbeat"),
                capture_time=frame_data.get("timestamp"),
                width=frame_data.get("width", 0),
                height=frame_data.get("height", 0)
            )
        except Exception as e:
            self.logger.error(f"Failed to parse video frame: {e}")
            return None
    
    def get_visual_context(self) -> Optional[ImageContent]:
        """
        Get current visual context.
        
        Priority:
        1. VAD-triggered speech context frame
        2. Latest heartbeat frame
        
        Returns:
            ImageContent or None
        """
        return self._speech_context_frame or self._latest_frame
    
    def clear_speech_context(self) -> None:
        """
        Clear speech context frame.
        
        Called after each turn to reset VAD-triggered frame.
        """
        self._speech_context_frame = None
    
    async def _generate_chunks(self) -> AsyncIterator[Chunk]:
        """
        Generate Chunk stream (true input source).
        
        Flow:
        1. Start video collector (if video_queue provided)
        2. Get raw data from read_queue
        3. protocol.decode_message() decode
        4. protocol.message_to_chunk() convert
        5. audio_codec.decode() decode audio (if needed)
        6. Inject visual context into audio chunks
        7. yield Chunk
        """
        self.logger.info(f"Starting input stream for session {self._session_id[:8]}")
        
        # Start video frame collector if video_queue is provided
        if self._video_queue:
            self._video_task = asyncio.create_task(self._collect_video_frames())
            self.logger.info("Video frame collector started")
        
        try:
            while self._running:
                # 1. Read raw data from queue
                try:
                    raw_data = await asyncio.wait_for(
                        self._read_queue.get(), 
                        timeout=30.0  # Prevent infinite wait
                    )
                except asyncio.TimeoutError:
                    # Timeout but connection may still be alive, continue waiting
                    continue
                
                if raw_data is None:
                    # End of stream signal
                    self.logger.debug(f"End of stream signal received")
                    break
                
                # 2. Protocol decode
                try:
                    message = self.protocol.decode_message(raw_data)
                except Exception as e:
                    self.logger.error(f"Failed to decode message: {e}")
                    continue
                
                # 3. Convert to Chunk
                turn_id = self._get_current_turn_id()
                chunk = self.protocol.message_to_chunk(message, self._session_id, turn_id)
                
                if chunk is None:
                    continue
                
                # 4. Audio decode (Opus -> PCM)
                if chunk.type == ChunkType.AUDIO_RAW and self.audio_codec and chunk.data:
                    try:
                        pcm_data = self.audio_codec.decode(chunk.data)
                        chunk.data = pcm_data
                    except Exception as e:
                        self.logger.error(f"Failed to decode audio: {e}")
                        continue
                
                # 5. Inject visual context into audio chunks
                if chunk.type == ChunkType.AUDIO_RAW and self._video_queue:
                    visual_ctx = self.get_visual_context()
                    if visual_ctx:
                        chunk.metadata["visual_context"] = visual_ctx
                
                # 6. Output Chunk
                self.logger.debug(f"Input chunk: {chunk.type.value} from {chunk.source}")
                yield chunk
        
        except asyncio.CancelledError:
            self.logger.debug(f"Input stream cancelled")
        except Exception as e:
            self.logger.error(f"Error in input stream: {e}", exc_info=True)
        finally:
            # Cancel video collector task
            if self._video_task and not self._video_task.done():
                self._video_task.cancel()
                try:
                    await self._video_task
                except asyncio.CancelledError:
                    pass
            self.logger.info(f"Input stream ended for session {self._session_id[:8]}")
    
    def _get_current_turn_id(self) -> int:
        """Get current turn_id"""
        if self.control_bus:
            return self.control_bus.get_current_turn_id()
        return self.current_turn_id
    
    def stop(self) -> None:
        """Stop generating Chunks"""
        self._running = False


class OutputStation(Station):
    """
    Output format conversion Station - Pipeline endpoint.
    
    Responsibilities (pure format conversion):
    1. Receive Chunk from Pipeline
    2. Call AudioCodec to encode audio (if needed)
    3. Call Protocol to convert to message
    4. Call Protocol to encode message
    5. Write to Transport's queues (priority_queue for control, send_queue for data)
    
    Dual queue design:
    - priority_queue: Control messages (TTS_STOP, STATE changes) - sent immediately
    - send_queue: Normal data (audio) - may be delayed by flow control
    
    NOT responsible for:
    - Flow control (handled by TransportBase.send_worker)
    - Turn management (handled by TransportBase.send_worker)
    - Latency tracking (handled by TransportBase.send_worker)
    """
    
    
    def __init__(
        self,
        session_id: str,
        send_queue: asyncio.Queue,      # From Transport - for normal data (audio)
        priority_queue: asyncio.Queue,  # From Transport - for control messages
        protocol: ProtocolBase,          # From Transport
        audio_codec: Optional[Any] = None,  # From Transport, can be None
        name: str = "OutputStation"
    ):
        """
        Args:
            session_id: Session ID
            send_queue: Transport's send queue for normal data
            priority_queue: Transport's priority queue for control messages
            protocol: Protocol handler
            audio_codec: Audio codec (can be None)
            name: Station name
            
        Note: Queue data format for send_queue is (data, metadata) tuple.
        metadata contains turn_id for latency tracking by send_worker.
        """
        super().__init__(name=name)
        self._session_id = session_id
        self._send_queue = send_queue
        self._priority_queue = priority_queue
        self.protocol = protocol
        self.audio_codec = audio_codec
    
    async def process_chunk(self, chunk: Chunk) -> AsyncIterator[Chunk]:
        """
        Process Chunk: convert to protocol message and put into appropriate queue.
        
        Queue selection based on ChunkType.is_high_priority():
        - High priority (CONTROL_INTERRUPT, STATE changes) -> priority_queue (immediate)
        - Normal (audio, TTS events) -> send_queue (may be delayed by flow control)
        
        Special handling:
        - CONTROL_INTERRUPT: converted to TTS_STOP + STATE_LISTENING messages
        - AUDIO_RAW: split into frames and encoded
        
        OutputStation is endpoint, doesn't produce new Chunks.
        """
        # Special handling for AUDIO_RAW: split into frames, encode, and send
        if chunk.type == ChunkType.AUDIO_RAW and "tts" in chunk.source.lower() and chunk.data:
            await self._send_audio_frames(chunk.data, chunk.turn_id)
            return
            yield  # Keep as generator
        
        # Special handling for CONTROL_INTERRUPT: send TTS_STOP + STATE_LISTENING
        if chunk.type == ChunkType.CONTROL_ABORT:
            await self._handle_abort(chunk)
            return
            yield  # Keep as generator
        
        # 1. Call Protocol business methods based on Chunk type
        message = await self._chunk_to_message(chunk)
        
        if message is None:
            return
            yield  # Keep as generator
        
        # 2. Encode message
        try:
            encoded_data = self.protocol.encode_message(message)
        except Exception as e:
            self.logger.error(f"Failed to encode message: {e}")
            return
            yield  # Keep as generator
        
        # 3. Choose queue based on chunk type priority
        use_priority = chunk.type.is_high_priority()
        
        # Create metadata for tracking
        metadata = {
            "type": chunk.type.value,
            "turn_id": chunk.turn_id,
            "is_tts_stop": chunk.type == ChunkType.EVENT_TTS_STOP,
        }
        
        if use_priority:
            # Priority queue: just data (no metadata needed for control messages)
            await self._priority_queue.put(encoded_data)
            self.logger.debug(f"Queued priority message: {chunk.type.value}")
        else:
            # Send queue: (data, metadata) tuple for latency tracking
            await self._send_queue.put((encoded_data, metadata))
        
        # OutputStation doesn't produce new Chunks
        return
        yield  # Keep as generator
    
    async def _handle_abort(self, chunk: Chunk) -> None:
        """
        Handle CONTROL_ABORT: send TTS_STOP and STATE_LISTENING messages.
        
        This converts the semantic interrupt signal into protocol-specific messages.
        Both messages are sent via priority_queue for immediate delivery.
        
        Args:
            chunk: CONTROL_ABORT chunk
        """
        # 1. Send TTS_STOP to stop client playback
        tts_stop_msg = self.protocol.send_tts_event(self._session_id, "stop")
        if tts_stop_msg:
            try:
                encoded = self.protocol.encode_message(tts_stop_msg)
                await self._priority_queue.put(encoded)
                self.logger.debug("Sent TTS_STOP via priority queue (abort)")
            except Exception as e:
                self.logger.error(f"Failed to send TTS_STOP: {e}")
        
        # 2. Send STATE_LISTENING to activate client microphone
        listen_msg = self.protocol.start_listen(self._session_id)
        if listen_msg:
            try:
                encoded = self.protocol.encode_message(listen_msg)
                await self._priority_queue.put(encoded)
                self.logger.debug("Sent STATE_LISTENING via priority queue (abort)")
            except Exception as e:
                self.logger.error(f"Failed to send STATE_LISTENING: {e}")
    
    async def _chunk_to_message(self, chunk: Chunk) -> Optional[Dict[str, Any]]:
        """
        Convert Chunk to protocol message (calling Protocol business methods).
        """
        # TEXT chunks - determine type based on source
        if chunk.type == ChunkType.TEXT:
            text = chunk.data
            
            if "asr" in chunk.source.lower():
                return self.protocol.send_stt(self._session_id, text)
            
            elif "agent" in chunk.source.lower():
                return self.protocol.send_llm(self._session_id, text)
        
        # AUDIO chunks - skip here, handled separately in process_chunk
        elif chunk.type == ChunkType.AUDIO_RAW:
            return None
        
        # TTS events
        elif chunk.type == ChunkType.EVENT_TTS_START:
            text = None
            if hasattr(chunk, 'event_data') and isinstance(chunk.event_data, dict):
                text = chunk.event_data.get("text")
            return self.protocol.send_tts_event(self._session_id, "start", text)
        
        elif chunk.type == ChunkType.EVENT_TTS_SENTENCE_START:
            text = None
            if hasattr(chunk, 'event_data') and isinstance(chunk.event_data, dict):
                text = chunk.event_data.get("text")
            return self.protocol.send_tts_event(self._session_id, "sentence_start", text)
        
        elif chunk.type == ChunkType.EVENT_TTS_SENTENCE_END:
            return self.protocol.send_tts_event(self._session_id, "sentence_end")
        
        elif chunk.type == ChunkType.EVENT_TTS_STOP:
            return self.protocol.send_tts_event(self._session_id, "stop")
        
        # State events - notify client of state changes
        elif chunk.type == ChunkType.EVENT_STATE_LISTENING:
            return self.protocol.start_listen(self._session_id)
        
        elif chunk.type == ChunkType.EVENT_STATE_IDLE:
            return self.protocol.stop_listen(self._session_id)
        
        elif chunk.type == ChunkType.EVENT_STATE_SPEAKING:
            return self.protocol.start_speaker(self._session_id)
        
        # Other Chunks - try default conversion
        else:
            return self.protocol.chunk_to_message(chunk)
    
    async def _send_audio_frames(self, pcm_data: bytes, turn_id: int) -> None:
        """
        Split audio into frames and put into send queue.
        
        Queue data format: (encoded_message, metadata)
        - metadata = {"type": "audio", "turn_id": turn_id, "is_first": bool}
        
        Args:
            pcm_data: PCM audio data (may be large, from TTS)
            turn_id: Turn ID for latency tracking
        """
        # Split PCM into 60ms frames (16kHz, mono, 16-bit = 1920 bytes per 60ms)
        frame_size = 1920  # 60ms at 16kHz mono
        frames = []
        for i in range(0, len(pcm_data), frame_size):
            frames.append(pcm_data[i:i + frame_size])
        
        self.logger.debug(f"Sending {len(frames)} audio frames for turn {turn_id}")
        
        # Encode and put into queue
        for frame_idx, frame in enumerate(frames):
            if not frame:
                continue
            
            # Encode frame (PCM -> ex:Opus)
            if self.audio_codec:
                try:
                    encoded_data = self.audio_codec.encode(frame)
                except Exception as e:
                    self.logger.error(f"Failed to encode audio frame: {e}")
                    continue
            else:
                # No codec, send raw PCM
                encoded_data = frame
            
            # Create audio message
            message = self.protocol.send_tts_audio(self._session_id, encoded_data)
            if not message:
                continue
            
            # Encode message
            try:
                encoded_message = self.protocol.encode_message(message)
            except Exception as e:
                self.logger.error(f"Failed to encode message: {e}")
                continue
            
            # Put into send queue with metadata for latency tracking
            # Format: (data, metadata) where metadata contains turn_id and frame info
            metadata = {
                "type": "audio",
                "turn_id": turn_id,
                "is_first_frame": frame_idx == 0,
            }
            await self._send_queue.put((encoded_message, metadata))
    
    def _get_current_turn_id(self) -> int:
        """Get current turn_id"""
        if self.control_bus:
            return self.control_bus.get_current_turn_id()
        return self.current_turn_id
