#!/usr/bin/env python3
"""
Soniox Bi-Directional Input Synchronization Gateway (v4.0)
- Mirrors Mac focused input box <--> Android phone in real-time.
- Supports inserting voice text at exact cursor location.
- Instant <15ms WebSocket push to Mac for editing and pasting.
"""

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()
connected_phone_websockets = set()
recent_pasted_sessions = OrderedDict()  # session_id -> (timestamp, text)

latest_mac_input_state = {
    "type": "mac_input_state",
    "app": "Desktop",
    "text": "",
    "cursor": 0,
    "selection": 0,
    "timestamp": time.time()
}

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 ""

    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_to_macs(payload_dict: dict) -> bool:
    """Pushes command payload directly to Mac via persistent WebSocket in <15ms."""
    delivered = False
    payload_str = json.dumps(payload_dict, ensure_ascii=False)

    dead_sockets = set()
    for ws in list(connected_mac_websockets):
        try:
            if is_ws_open(ws):
                await ws.send_str(payload_str)
                delivered = True
                logger.info("⚡ Pushed to Mac WS: %s", payload_dict.get("action") or payload_dict.get("type"))
            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

    # Fallback to SSH script if WebSocket temporarily disconnected
    try:
        text = payload_dict.get("text", "")
        if text:
            clean_text = sanitize_and_flatten_text(text)
            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 fallback: '%s'", clean_text[:30])
                return True
    except Exception as e:
        logger.warning("SSH fallback error: %s", e)

    return False

async def broadcast_to_phones(payload_dict: dict):
    """Pushes Mac input state changes to all connected Android clients."""
    payload_str = json.dumps(payload_dict, ensure_ascii=False)
    dead_phones = set()
    for ws in list(connected_phone_websockets):
        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)

async def handle_phone_stream_ws(request):
    """
    ⚡ Persistent Duplex Channel for Android Client:
    - Receives live Mac input state on connect and continuously.
    - Handles live audio streaming and returns live STT tokens.
    - Receives user manual text edits and pushes them instantly to Mac.
    """
    ws = web.WebSocketResponse(heartbeat=15.0)
    await ws.prepare(request)
    client_ip = request.remote
    logger.info("📱 Android Client connected to Duplex Stream: %s", client_ip)
    connected_phone_websockets.add(ws)

    # Immediately send the latest Mac input state to phone upon connection!
    if latest_mac_input_state:
        await ws.send_str(json.dumps(latest_mac_input_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 (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()

                    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.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)

                    # Return final voice transcription to Android
                    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 == "phone_input_edit" or msg_type == "update_mac_input":
                    # User manually edited text on phone or pressed "Insert in Mac"
                    edit_text = data.get("text", "")
                    cursor = data.get("cursor_pos") or data.get("cursor")
                    is_full_replace = data.get("is_full_replace", True)
                    logger.info("📱 Received phone edit to push to Mac: '%s' (cursor: %s)", edit_text[:30], cursor)

                    mac_ok = await broadcast_to_macs({
                        "action": "update_input",
                        "text": edit_text,
                        "cursor": cursor,
                        "is_full_replace": is_full_replace
                    })

                    # Update our cached latest state
                    latest_mac_input_state["text"] = edit_text
                    latest_mac_input_state["cursor"] = cursor if cursor is not None else len(edit_text)
                    latest_mac_input_state["timestamp"] = time.time()

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

            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 input state & pushes edits)."""
    global latest_mac_input_state
    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 Duplex 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 == "mac_input_state":
                        # Mac reports focused input box text and cursor position
                        latest_mac_input_state = data
                        # Broadcast immediately to phone!
                        await broadcast_to_phones(data)

                    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)
        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 Bi-Directional Duplex Gateway v4.0",
        "connected_macs": len(connected_mac_websockets),
        "connected_phones": len(connected_phone_websockets),
        "latest_mac_input_app": latest_mac_input_state.get("app", ""),
        "latest_mac_input_text_len": len(latest_mac_input_state.get("text", ""))
    })

async def handle_paste(request):
    try:
        data = await request.json()
        text = data.get("text", "")
        session_id = data.get("session_id", "")
        cursor = data.get("cursor_pos") or data.get("cursor")
        success = await broadcast_to_macs({
            "action": "update_input",
            "text": text,
            "cursor": cursor,
            "is_full_replace": True
        })
        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)