#!/usr/bin/env python3
"""
Soniox Ultra-Low Latency Duplex Gateway (v3.0)
- Single persistent duplex WebSocket to Android.
- Pre-warmed upstream Soniox pool.
- Strict newline stripping & hallucination cleanup.
- Sub-50ms Cmd+V Mac paste dispatch.
"""

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("SonioxRelay")

connected_mac_websockets = set()
recent_pasted_sessions = OrderedDict()  # session_id -> (timestamp, text)

MAC_DIRECT_HTTP_URL = "http://116.16.16.20:8999/paste"
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=3):
        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=20,
            )
        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:
    """
    1. Removes all line breaks (\\r, \\n) and collapses whitespace into single spaces.
    2. Strips English hallucination stop words during Persian speech.
    3. Guarantees zero trailing/leading enters or spaces.
    """
    if not text:
        return ""

    # Flatten all newlines, carriage returns, tabs into a single line
    flattened = re.sub(r"[\r\n\t]+", " ", text)
    flattened = re.sub(r"\s+", " ", flattened).strip()

    if not flattened:
        return ""

    words = flattened.split()
    fa_pattern = re.compile(r"[\u0600-\u06FF\u0750-\u077F\uFB50-\uFDFF\uFE70-\uFEFF]")
    en_pattern = re.compile(r"[a-zA-Z]")

    fa_count = sum(1 for w in words if fa_pattern.search(w))
    en_count = sum(1 for w in words if en_pattern.search(w))
    total = fa_count + en_count

    if total == 0:
        return flattened

    fa_ratio = fa_count / total
    stop_words = {"sex", "from", "no", "oh", "you", "all", "know", "that", "god", "loves", "men", "for", "day", "night", "mankind", "wanted", "moments"}

    cleaned = []
    if fa_ratio >= 0.25:
        for w in words:
            if en_pattern.search(w) and not fa_pattern.search(w):
                clean_w = re.sub(r"[.,!?:;،؛؟\"'()\[\]{}«»–—-]", "", w.lower())
                if clean_w in stop_words or fa_ratio >= 0.70:
                    continue
            cleaned.append(w)
    else:
        cleaned = words

    result = " ".join(cleaned)
    return re.sub(r"\s+", " ", result).strip()

async def broadcast_paste_to_macs(text: str, session_id: str = "") -> bool:
    """Pushes clean single-line paste payload to Mac in <50ms with strict deduplication."""
    clean_text = sanitize_and_flatten_text(text)
    if not clean_text:
        return False

    now = time.time()
    if session_id:
        if session_id in recent_pasted_sessions:
            logger.info("🚫 Suppressed duplicate paste for session_id: %s", session_id)
            return True
        recent_pasted_sessions[session_id] = (now, clean_text)
        while recent_pasted_sessions and (now - next(iter(recent_pasted_sessions.values()))[0] > 60):
            recent_pasted_sessions.popitem(last=False)
    else:
        for prev_sid, (prev_time, prev_txt) in list(recent_pasted_sessions.items())[-5:]:
            if prev_txt == clean_text and (now - prev_time) < 2.0:
                logger.info("🚫 Suppressed rapid identical text paste: '%s'", clean_text[:25])
                return True

    delivered = False
    payload = json.dumps({"action": "paste", "text": clean_text}, ensure_ascii=False)

    # 1. Primary: Direct Push via Active Persistent WebSocket (<15ms)
    dead_sockets = set()
    for ws in list(connected_mac_websockets):
        try:
            if is_ws_open(ws):
                await ws.send_str(payload)
                delivered = True
                logger.info("⚡ Pushed paste via Mac Persistent WS: '%s'", clean_text[:30])
            else:
                dead_sockets.add(ws)
        except Exception:
            dead_sockets.add(ws)

    for dead in dead_sockets:
        connected_mac_websockets.discard(dead)

    if delivered:
        return True

    # 2. Secondary: Direct LAN HTTP (/paste) (<35ms)
    try:
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=0.35)) as session:
            async with session.post(MAC_DIRECT_HTTP_URL, json={"text": clean_text}) as resp:
                if resp.status == 200:
                    logger.info("✅ Pasted to Mac via Direct LAN HTTP: '%s'", clean_text[:30])
                    return True
    except Exception:
        pass

    # 3. Tertiary: SSH Tunnel (Port 2222) - FIXED: Zero-Newline printf '%s' pipe
    try:
        escaped_text = clean_text.replace("'", "'\\''")
        remote_cmd = (
            f"printf '%s' '{escaped_text}' | pbcopy && "
            f"/usr/bin/osascript -e 'tell application \"System Events\" to keystroke \"v\" using command down'"
        )
        proc = await asyncio.create_subprocess_exec(
            "ssh", "-p", "2222", "-o", "BatchMode=yes", "-o", "ConnectTimeout=2", "alig@127.0.0.1",
            remote_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        stdout, stderr = await proc.communicate()
        if proc.returncode == 0:
            logger.info("✅ Pasted to Mac via SSH Tunnel (Zero-Newline printf): '%s'", clean_text[:30])
            return True
    except Exception as e:
        logger.warning("SSH tunnel fallback error: %s", e)

    return False

async def handle_phone_stream_ws(request):
    """
    ⚡ Single Persistent Duplex WebSocket Handler for Android.
    Supports multiple sequential sessions without tearing down the connection.
    """
    ws = web.WebSocketResponse(heartbeat=15.0)
    await ws.prepare(request)
    client_ip = request.remote
    logger.info("📱 Android Client connected to Persistent Stream: %s", client_ip)

    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 (2048 bytes / 64ms)
                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)}")

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

                    # Acquire pre-warmed Soniox session
                    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 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.55)
                        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)

                    # Instant <50ms broadcast to Mac
                    mac_ok = False
                    if clean_final:
                        mac_ok = await broadcast_paste_to_macs(clean_final, session_id=sid)

                    if is_ws_open(ws):
                        await ws.send_str(json.dumps({
                            "type": "final",
                            "session_id": sid,
                            "text": clean_final,
                            "mac_delivered": mac_ok
                        }))

                    # Cleanup Soniox session for this turn
                    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:
        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."""
    ws = web.WebSocketResponse(heartbeat=10.0)
    await ws.prepare(request)
    client_ip = request.remote
    logger.info("🖥️ Mac client connected to persistent WebSocket: %s", client_ip)
    connected_mac_websockets.add(ws)

    try:
        await ws.send_str(json.dumps({"type": "welcome", "message": "Connected to Soniox Ultra Gateway"}))
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                    if data.get("type") == "ping":
                        await ws.send_str(json.dumps({"type": "pong"}))
                except Exception:
                    pass
            elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR):
                break
    finally:
        connected_mac_websockets.discard(ws)
        if is_ws_open(ws):
            await ws.close()
        logger.info("🖥️ Mac client disconnected: %s", client_ip)

    return ws

async def handle_health(request):
    return web.json_response({
        "status": "ok",
        "service": "Soniox Pre-Warmed Duplex Gateway v3.0",
        "connected_macs": len(connected_mac_websockets),
        "recent_sessions": len(recent_pasted_sessions)
    })

async def handle_paste(request):
    try:
        data = await request.json()
        text = data.get("text", "")
        session_id = data.get("session_id", "")
        success = await broadcast_paste_to_macs(text, session_id=session_id)
        return web.json_response({"status": "pasted" if success else "failed", "mac_delivered": success})
    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)