import asyncio
import websockets
import numpy as np
from .whisper_vad import RealtimeWhisperVAD
import json
import logging
import time
import aiohttp

logger = logging.getLogger(__name__)
clients = set()


# Simple transcript assembly utilities
class TranscriptBuffer:
    def __init__(self, silence_threshold=1.0):
        self.words = []
        self.last_word_time = None
        self.silence_threshold = silence_threshold

    def add_word(self, word_obj):
        self.words.append(word_obj.word)
        self.last_word_time = time.time()

    def should_flush(self):
        if not self.words:
            return False
        return (time.time() - (self.last_word_time or 0)) > self.silence_threshold

    def flush(self):
        txt = " ".join(self.words).strip()
        self.words = []
        self.last_word_time = None
        return txt


async def post_transcript_to_rag(transcript: str):
    url = "http://localhost:5001/evaluate"
    payload = {"transcript": transcript}
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(url, json=payload, timeout=10) as resp:
                text = await resp.text()
                logger.info(f"Posted transcript to RAG; response status={resp.status}")
                return resp.status, text
    except Exception as e:
        logger.error(f"Failed to post transcript to RAG: {e}")
        return None, str(e)


# Callback function for RealtimeWhisperVAD to call per word
transcript_buffer = TranscriptBuffer(silence_threshold=1.0)

def word_callback(word):
    # This callback runs in a thread; schedule async broadcast and flush checks
    try:
        # Add to transcript buffer
        transcript_buffer.add_word(word)
        # Broadcast word to connected clients asynchronously
        msg = json.dumps({"word": word.word, "start": word.start, "end": word.end})
        asyncio.run_coroutine_threadsafe(broadcast_message(msg), asyncio.get_event_loop())
    except Exception as e:
        logger.error(f"Error in word_callback: {e}")


async def broadcast_message(message: str):
    dead = []
    for client in clients:
        try:
            await client.send(message)
        except Exception:
            dead.append(client)
    for d in dead:
        clients.discard(d)


# Configuration for RealtimeWhisperVAD
config = {
    "whisper_model": "tiny",
    "sample_rate": 16000,
    "chunk_duration": 0.9,
    "vad_threshold": 0.5,
    "min_speech_duration": 0.5,
    "max_speech_duration": 5,
    "silence_duration": 0.5,
    "transcription_callback": word_callback,
}

# Initialize RealtimeWhisperVAD
transcriber = RealtimeWhisperVAD(**config)


async def handler(websocket, path):
    logger.info(f"Client connected from {websocket.remote_address}")
    clients.add(websocket)
    try:
        # If STT backend is unavailable, inform client and close
        if not getattr(transcriber, 'use_faster_whisper', False):
            await websocket.send(json.dumps({"error": "STT backend unavailable"}))
            await websocket.close()
            return

        # Start the transcriber when a client connects
        if not transcriber.running:
            transcriber.start()

        # Polling loop: receive audio chunks and periodically check buffer flush
        async for message in websocket:
            try:
                audio_float32 = np.frombuffer(message, dtype=np.float32)
                transcriber.add_audio_chunk(audio_float32)
            except Exception as e:
                logger.error(f"Error processing audio chunk: {e}")

            # If buffer indicates silence, flush transcript and POST to RAG
            if transcript_buffer.should_flush():
                transcript = transcript_buffer.flush()
                if transcript:
                    status, resp_text = await post_transcript_to_rag(transcript)
                    # Send evaluation result back to clients
                    payload = {"event": "evaluation_result", "status": status, "resp": resp_text}
                    await broadcast_message(json.dumps(payload))

    except websockets.exceptions.ConnectionClosed:
        logger.info(f"Client disconnected from {websocket.remote_address}")
    finally:
        clients.discard(websocket)


async def main():
    # Use 0.0.0.0 to be reachable; keep localhost if you want local-only
    async with websockets.serve(handler, "0.0.0.0", 5002):
        logger.info("WebSocket server started on ws://0.0.0.0:5002")
        await asyncio.Future()  # run forever


if __name__ == "__main__":
    asyncio.run(main())