#!/usr/bin/env python3
"""
Soniox Collaborative Bi-Directional Input Synchronization Gateway (v5.5)
- Ultra-low latency streaming speech recognition with pre-warmed Soniox pool.
- Dedicated asynchronous audio forwarder queue to prevent receiver blocking & ping timeouts.
- Extended 4.0s finalize timeout for flawless long-speech transcriptions.
- Strict artifact & lone-punctuation suppression on speech stream.
- Full multiline (\n) preservation for typing & keyboard sync.
"""

import asyncio
import json
import logging
import re
import subprocess
import time
from collections import OrderedDict
import aiohttp
from aiohttp import web
import websockets
from websockets.protocol import State

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("SyncGateway")

connected_mac_websockets = set()
connected_phone_websockets = set()

# Global Room State Snapshot (Single Source of Truth)
current_room_state = {
    "type": "sync_state",
    "source": "server_init",
    "app": "Desktop",
    "text": "",
    "cursor": 0,
    "selection": 0,
    "revision": 0,
    "timestamp": time.time()
}
room_lock = asyncio.Lock()

SONIOX_ENDPOINTS = [
    (
        "wss://translate.compare.soniox.com/compare/api/compare-websocket"
        "?language_hints=fa&language_hints=en&language_hints=ar"
        "&enable_speaker_diarization=false&enable_language_identification=true"
        "&enable_endpoint_detection=false&providers=soniox",
        {"Origin": "https://translate.compare.soniox.com"}
    ),
    (
        "wss://stt.compare.soniox.com/compare/api/compare-websocket"
        "?language_hints=fa&language_hints=en&language_hints=ar"
        "&enable_speaker_diarization=false&enable_language_identification=true"
        "&enable_endpoint_detection=false&providers=soniox",
        {"Origin": "https://stt.compare.soniox.com"}
    )
]

def is_ws_open(ws) -> bool:
    if ws is None:
        return False
    if hasattr(ws, 'closed'):
        return not ws.closed
    if hasattr(ws, 'state'):
        return ws.state == State.OPEN
    return True

class SonioxPool:
    """Pre-warms upstream WebSockets to Soniox for 0ms speech start delay."""
    def __init__(self, size=4):
        self._pool = asyncio.Queue(maxsize=size)
        self._refilling = False

    async def get_session(self):
        while not self._pool.empty():
            try:
                ws = self._pool.get_nowait()
                if is_ws_open(ws):
                    asyncio.create_task(self.refill())
                    return ws
            except asyncio.QueueEmpty:
                break

        logger.info("Pool empty, connecting fresh Soniox session...")
        asyncio.create_task(self.refill())
        return await self._create_ws()

    async def _create_ws(self):
        for url, headers in SONIOX_ENDPOINTS:
            try:
                req_headers = {
                    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36",
                    **headers
                }
                ws = await websockets.connect(
                    url,
                    additional_headers=req_headers,
                    open_timeout=7.0,
                    ping_interval=20,
                    ping_timeout=20,
                    max_size=10 * 1024 * 1024
                )
                return ws
            except Exception as e:
                logger.warning("Soniox endpoint %s failed: %s", url[:40], e)
        return None

    async def refill(self):
        if self._refilling or self._pool.full():
            return
        self._refilling = True
        try:
            while not self._pool.full():
                ws = await self._create_ws()
                if ws and is_ws_open(ws):
                    await self._pool.put(ws)
                else:
                    break
        finally:
            self._refilling = False

soniox_pool = SonioxPool()

def sanitize_speech_text(text: str) -> str:
    if not text:
        return ""
    flattened = re.sub(r"[\r\n\t]+", " ", text)
    flattened = re.sub(r"\s+", " ", flattened).strip()
    
    # Suppress lone punctuation artifacts (e.g. «, », ., ,, !, ?, etc.)
    if re.fullmatch(r"[\s«»\.\,\،\؛\؟\!\?\:\;\-\–—\"\'\(\)\[\]\{\}]+", flattened):
        return ""
        
    return flattened

def normalize_sync_text(text: str) -> str:
    if not text:
        return ""
    # Normalize line endings and preserve intentional multiline text (\n)
    normalized = text.replace("\r\n", "\n").replace("\r", "\n")
    return normalized

async def broadcast_state(payload_dict: dict, exclude_ws=None):
    """Broadcasts state snapshot to all connected clients (Mac and Phone) except sender."""
    global current_room_state
    async with room_lock:
        current_room_state["revision"] += 1
        payload_dict["revision"] = current_room_state["revision"]
        payload_dict["timestamp"] = time.time()
        
        # Update our cached authoritative state
        current_room_state.update(payload_dict)
        payload_str = json.dumps(payload_dict, ensure_ascii=False)

    logger.info("📡 Broadcasting %s (len: %d) to %d Macs, %d Phones", 
                payload_dict.get("type"), len(payload_dict.get("text", "")), 
                len(connected_mac_websockets), len(connected_phone_websockets))

    # 1. Send to Phone clients
    dead_phones = set()
    for ws in list(connected_phone_websockets):
        if ws is exclude_ws:
            continue
        try:
            if is_ws_open(ws):
                await ws.send_str(payload_str)
            else:
                dead_phones.add(ws)
        except Exception:
            dead_phones.add(ws)
    for dead in dead_phones:
        connected_phone_websockets.discard(dead)

    # 2. Send to Mac clients
    dead_macs = set()
    for ws in list(connected_mac_websockets):
        if ws is exclude_ws:
            continue
        try:
            if is_ws_open(ws):
                await ws.send_str(payload_str)
            else:
                dead_macs.add(ws)
        except Exception:
            dead_macs.add(ws)
    for dead in dead_macs:
        connected_mac_websockets.discard(dead)

async def handle_phone_stream_ws(request):
    """Duplex WebSocket for Android client (Real-time Live Sync + Audio Dictation)."""
    ws = web.WebSocketResponse(heartbeat=30.0, autoping=True, max_msg_size=15 * 1024 * 1024)
    await ws.prepare(request)
    client_ip = request.remote
    logger.info("📱 Android Client Connected: %s", client_ip)
    connected_phone_websockets.add(ws)

    # Immediately hydrate phone with full current state!
    async with room_lock:
        await ws.send_str(json.dumps(current_room_state, ensure_ascii=False))

    active_soniox_ws = None
    reader_task = None
    forwarder_task = None
    stop_event = asyncio.Event()
    current_session_id = ""
    audio_queue = asyncio.Queue(maxsize=1000)

    full_final_tokens = []
    current_non_final = ""

    async def audio_forwarder(soniox_ws):
        try:
            while True:
                chunk = await audio_queue.get()
                if chunk is None:
                    audio_queue.task_done()
                    break
                if is_ws_open(soniox_ws):
                    await soniox_ws.send(chunk)
                audio_queue.task_done()
        except Exception as e:
            logger.warning("Audio forwarder error: %s", e)

    async def soniox_reader(soniox_ws, sid):
        nonlocal current_non_final
        try:
            async for s_msg in soniox_ws:
                if not is_ws_open(ws):
                    break
                try:
                    data = json.loads(s_msg)
                    msg_type = data.get("type")

                    if msg_type == "session_done" or data.get("session_ended"):
                        logger.info("⚡ Soniox session_done received for %s", sid)
                        stop_event.set()
                        break

                    if msg_type == "data" and "parts" in data:
                        new_finals = []
                        new_non_finals = []
                        got_fin = False
                        for p in data["parts"]:
                            if p.get("translation_status") == "translation":
                                continue
                            txt = p.get("text", "")
                            is_final = p.get("is_final", False)
                            if "<fin>" in txt:
                                got_fin = True
                                clean = txt.replace("<fin>", "")
                                if clean: new_finals.append(clean)
                            elif is_final:
                                if txt: new_finals.append(txt)
                            else:
                                if txt: new_non_finals.append(txt)

                        if new_finals:
                            full_final_tokens.extend(new_finals)
                        current_non_final = "".join(new_non_finals)

                        live_text = sanitize_speech_text("".join(full_final_tokens) + current_non_final)
                        final_clean = sanitize_speech_text("".join(full_final_tokens))
                        partial_clean = sanitize_speech_text(current_non_final)

                        if live_text and is_ws_open(ws):
                            await ws.send_str(json.dumps({
                                "type": "live",
                                "session_id": sid,
                                "text": live_text,
                                "final_text": final_clean,
                                "partial_text": partial_clean
                            }))

                        if got_fin:
                            stop_event.set()
                            break
                except Exception as e:
                    logger.warning("Soniox parse error: %s", e)
        except Exception as e:
            logger.warning("Soniox reader error: %s", e)
        finally:
            stop_event.set()

    try:
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.BINARY:
                # Non-blocking audio chunk enqueue
                if not audio_queue.full():
                    audio_queue.put_nowait(msg.data)

            elif msg.type == aiohttp.WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                except Exception:
                    continue

                msg_type = data.get("type") or data.get("action")
                sid = data.get("session_id", f"sess_{int(time.time()*1000)}")

                # ALL text / sync / insert operations must be broadcast to Mac!
                if msg_type in ("sync_state", "insert_speech", "speech_insert", "phone_input_edit", "update_input", "paste"):
                    if msg_type in ("insert_speech", "speech_insert"):
                        clean_text = sanitize_speech_text(data.get("text", ""))
                        if not clean_text:
                            continue  # Do not broadcast empty speech or lone quotes
                    else:
                        # Direct typing / keyboard / next lines: preserve multiline \n exactly!
                        raw_text = data.get("text", "")
                        clean_text = normalize_sync_text(raw_text) if raw_text else ""

                    data["text"] = clean_text
                    data["source"] = "android"
                    data["type"] = msg_type
                    await broadcast_state(data, exclude_ws=ws)

                elif msg_type == "start":
                    current_session_id = sid
                    full_final_tokens.clear()
                    current_non_final = ""
                    stop_event.clear()

                    # Clear existing queue
                    while not audio_queue.empty():
                        try: audio_queue.get_nowait(); audio_queue.task_done()
                        except Exception: break

                    active_soniox_ws = await soniox_pool.get_session()
                    if not active_soniox_ws or not is_ws_open(active_soniox_ws):
                        await ws.send_str(json.dumps({"type": "error", "message": "Upstream Soniox unavailable"}))
                        continue

                    if reader_task and not reader_task.done():
                        reader_task.cancel()
                    if forwarder_task and not forwarder_task.done():
                        forwarder_task.cancel()

                    forwarder_task = asyncio.create_task(audio_forwarder(active_soniox_ws))
                    reader_task = asyncio.create_task(soniox_reader(active_soniox_ws, current_session_id))
                    logger.info("🎙️ Started Live Speech Session: %s", current_session_id)

                elif msg_type == "stop":
                    logger.info("⏹️ Stop received for %s. Draining audio queue...", sid)
                    
                    # 1. Drain audio queue to Soniox with timeout
                    try:
                        await asyncio.wait_for(audio_queue.join(), timeout=2.0)
                    except Exception:
                        pass

                    # 2. Send finalize
                    if active_soniox_ws and is_ws_open(active_soniox_ws):
                        await active_soniox_ws.send(json.dumps({"type": "finalize"}))
                        try:
                            # 4.0s timeout to allow long speech finalization without cutoff
                            await asyncio.wait_for(stop_event.wait(), timeout=4.0)
                        except asyncio.TimeoutError:
                            logger.info("Finalize wait completed (timeout reached)")

                    raw_final = "".join(full_final_tokens) + current_non_final
                    clean_final = sanitize_speech_text(raw_final)
                    logger.info("⚡ Session %s final text (%d chars): '%s'", sid, len(clean_final), clean_final)

                    # Return final speech text to phone
                    if is_ws_open(ws):
                        await ws.send_str(json.dumps({
                            "type": "final",
                            "session_id": sid,
                            "text": clean_final,
                            "cursor_pos": data.get("cursor_pos", -1)
                        }))

                    # Cleanup tasks
                    if forwarder_task and not forwarder_task.done():
                        forwarder_task.cancel()
                    if reader_task and not reader_task.done():
                        reader_task.cancel()

                    if active_soniox_ws:
                        try:
                            await active_soniox_ws.close()
                        except Exception:
                            pass
                        active_soniox_ws = None

                    asyncio.create_task(soniox_pool.refill())

            elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR):
                break

    finally:
        connected_phone_websockets.discard(ws)
        if forwarder_task and not forwarder_task.done():
            forwarder_task.cancel()
        if reader_task and not reader_task.done():
            reader_task.cancel()
        if active_soniox_ws:
            try:
                await active_soniox_ws.close()
            except Exception:
                pass
        asyncio.create_task(soniox_pool.refill())
        logger.info("📱 Android Client disconnected: %s", client_ip)

    return ws

async def handle_mac_ws(request):
    """Persistent WebSocket for Mac Bridge (receives live Mac typing & pushes edits)."""
    ws = web.WebSocketResponse(heartbeat=30.0, autoping=True)
    await ws.prepare(request)
    client_ip = request.remote
    logger.info("🖥️ Mac client connected: %s", client_ip)
    connected_mac_websockets.add(ws)

    try:
        await ws.send_str(json.dumps({"type": "welcome", "message": "Connected to Soniox Sync Gateway"}))
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                    msg_type = data.get("type")

                    if msg_type == "sync_state" or msg_type == "mac_input_state":
                        # If server recently processed an active edit from android (< 1.5s), suppress stale Mac echoes
                        if current_room_state.get("source") == "android":
                            time_since_android = time.time() - current_room_state.get("timestamp", 0)
                            if time_since_android < 1.5 and data.get("text") != current_room_state.get("text"):
                                logger.info("Suppressing stale Mac sync_state echo (%0.2fs after android edit)", time_since_android)
                                continue

                        # Mac reports typing / cursor change: broadcast to all phones immediately!
                        data["source"] = "mac"
                        data["type"] = "sync_state"
                        await broadcast_state(data, exclude_ws=ws)

                    elif msg_type == "ping":
                        await ws.send_str(json.dumps({"type": "pong"}))
                except Exception as e:
                    logger.warning("Mac message error: %s", e)
            elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR):
                break
    finally:
        connected_mac_websockets.discard(ws)
        logger.info("🖥️ Mac client disconnected: %s", client_ip)

    return ws

async def handle_health(request):
    async with room_lock:
        state_copy = dict(current_room_state)
    return web.json_response({
        "status": "ok",
        "service": "Soniox Collaborative Sync Gateway v5.5",
        "connected_macs": len(connected_mac_websockets),
        "connected_phones": len(connected_phone_websockets),
        "current_app": state_copy.get("app", ""),
        "current_revision": state_copy.get("revision", 0),
        "current_text_len": len(state_copy.get("text", ""))
    })

async def handle_paste(request):
    try:
        data = await request.json()
        raw_text = data.get("text", "")
        clean_text = normalize_sync_text(raw_text) if raw_text else ""
        data["text"] = clean_text
        data["source"] = "http_post"
        data["type"] = "update_input"
        await broadcast_state(data)
        return web.json_response({"status": "synced", "revision": current_room_state["revision"]})
    except Exception as e:
        return web.json_response({"error": str(e)}, status=400)

async def start_background_tasks(app):
    asyncio.create_task(soniox_pool.refill())

def create_app():
    app = web.Application(client_max_size=25 * 1024 * 1024)
    app.on_startup.append(start_background_tasks)
    app.router.add_get("/health", handle_health)
    app.router.add_get("/status", handle_health)
    app.router.add_post("/paste", handle_paste)
    app.router.add_get("/ws/stream", handle_phone_stream_ws)
    app.router.add_get("/ws/mac", handle_mac_ws)
    return app

if __name__ == "__main__":
    app = create_app()
    web.run_app(app, host="0.0.0.0", port=8999)