#!/usr/bin/env python3
"""
Soniox Collaborative Bi-Directional Input Synchronization Gateway (v5.3)
- Handles all message types: sync_state, insert_speech, speech_insert, update_input, paste.
- Sub-15ms WebSocket routing between Android and Mac.
"""

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_WS_URL = (
    "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"
)
SONIOX_HEADERS = {
    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36",
    "Origin": "https://translate.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):
        try:
            return await websockets.connect(
                SONIOX_WS_URL,
                additional_headers=SONIOX_HEADERS,
                open_timeout=3.5,
                ping_interval=15,
            )
        except Exception as e:
            logger.error("Failed to connect to upstream Soniox: %s", 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_and_flatten_text(text: str) -> str:
    if not text:
        return ""
    flattened = re.sub(r"[\r\n\t]+", " ", text)
    flattened = re.sub(r"\s+", " ", flattened).strip()
    return flattened

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=10.0, autoping=True)
    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
    stop_event = asyncio.Event()
    current_session_id = ""

    full_final_tokens = []
    current_non_final = ""

    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)
                    if data.get("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_and_flatten_text("".join(full_final_tokens) + 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
                            }))

                        if got_fin or data.get("session_ended") or data.get("type") == "session_done":
                            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:
                # Live PCM audio chunk
                if active_soniox_ws and is_ws_open(active_soniox_ws):
                    await active_soniox_ws.send(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"):
                    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()

                    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()
                    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":
                    if active_soniox_ws and is_ws_open(active_soniox_ws):
                        await active_soniox_ws.send(json.dumps({"type": "finalize"}))
                        try:
                            await asyncio.wait_for(stop_event.wait(), timeout=0.35)
                        except asyncio.TimeoutError:
                            pass

                    raw_final = "".join(full_final_tokens) + current_non_final
                    clean_final = sanitize_and_flatten_text(raw_final)
                    logger.info("⚡ Session %s final text: '%s'", sid, 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)
                        }))

                    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 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=10.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":
                        # 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.3",
        "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()
        text = data.get("text", "")
        cursor = data.get("cursor_pos") or data.get("cursor")
        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)