# stt_client.py
# Upstream WebSocket client to STT Space

import os
import json
import logging
import asyncio
import websockets
from typing import Optional

logger = logging.getLogger("stt_client")

class SttClient:
    def __init__(self, ws_url: Optional[str] = None):
        self.ws_url = ws_url or os.getenv("STT_SPACE_WS_URL", "wss://abid-official-stt.hf.space/ws/stt")
        self.ws = None
        self.connected = False
        self.send_queue = asyncio.Queue()
        self.receive_queue = asyncio.Queue()
        self._running = False
        self._tasks = []
        self._run_task = None

    def start(self):
        """Starts the upstream connection management task."""
        if self._running or (self._run_task and not self._run_task.done()):
            logger.info("STT client is already running or start task is active.")
            return self
        self._running = True
        self._run_task = asyncio.create_task(self.connect_and_run())
        return self

    async def connect_and_run(self):
        """Main connection loop managing auto-reconnect backoff and read/write loops."""
        backoff = 1.0
        while self._running:
            try:
                logger.info(f"Connecting to upstream STT Space: {self.ws_url}")
                async with websockets.connect(self.ws_url) as ws:
                    self.ws = ws
                    self.connected = True
                    backoff = 1.0
                    logger.info("Connected to upstream STT Space ✓")

                    # Spawn read and write loop tasks
                    read_task = asyncio.create_task(self._read_loop())
                    write_task = asyncio.create_task(self._write_loop())
                    self._tasks = [read_task, write_task]

                    await asyncio.gather(read_task, write_task)
            except asyncio.CancelledError:
                logger.info("STT client run task cancelled.")
                break
            except Exception as e:
                logger.error(f"STT upstream connection error: {e}")
                self.connected = False
                for t in self._tasks:
                    if not t.done():
                        t.cancel()
                self._tasks = []
                self.ws = None
                
                # Exponential backoff
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, 30.0)
        
        self.connected = False

    async def _write_loop(self):
        try:
            while self._running:
                data = await self.send_queue.get()
                if self.ws and self.connected:
                    try:
                        await self.ws.send(data)
                    except Exception as e:
                        logger.error(f"STT client write error: {e}")
                        break
                self.send_queue.task_done()
        except asyncio.CancelledError:
            pass

    async def _read_loop(self):
        try:
            async for message in self.ws:
                try:
                    event = json.loads(message)
                    await self.receive_queue.put(event)
                except Exception as e:
                    logger.warning(f"Failed to parse STT message: {e} | Raw: {message[:120]}")
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"Upstream STT read loop error: {e}")
            self.connected = False

    def send_pcm(self, pcm_bytes: bytes):
        """Queue raw PCM bytes to send to STT Space."""
        if self._running:
            self.send_queue.put_nowait(pcm_bytes)

    async def get_event(self) -> dict:
        """Fetch next transcription event from queue."""
        return await self.receive_queue.get()

    async def close(self):
        """Closes all queues, loops, and websocket connections."""
        self._running = False
        self.connected = False
        
        if self._run_task and not self._run_task.done():
            self._run_task.cancel()
            
        for t in self._tasks:
            if not t.done():
                t.cancel()
        self._tasks = []
        
        if self.ws:
            try:
                await self.ws.close()
            except Exception:
                pass
            self.ws = None
        logger.info("STT client connections closed.")