tai/tai_pkg/edge.py

"""Edge voice synthesis and local audio capture/playback."""

import base64
import hashlib
import html
import os
import random
import re
import shutil
import socket
import ssl
import struct
import subprocess
import sys
import time
import uuid

from tai_pkg.consts import Consts
from tai_pkg.errors import BackendError
from tai_pkg.text import Text

__all__ = ["EdgeTts", "Audio"]


class EdgeTts:
    """EdgeTts."""

    @staticmethod
    def edge_token(skew=0):
        ticks = int(time.time()) + skew + 11644473600
        ticks -= ticks % 300
        ticks *= 10000000
        return hashlib.sha256(("%d%s" % (ticks, Consts.EDGE_TRUSTED_TOKEN)).encode("ascii")).hexdigest().upper()

    @staticmethod
    def edge_timestamp():
        return time.strftime("%a %b %d %Y %H:%M:%S GMT+0000 (Coordinated Universal Time)", time.gmtime())

    @staticmethod
    def edge_ssml(text, voice):
        safe = html.escape(text, quote=False).replace("\r", " ").replace("\n", " ")
        return "<speak version='1.0' xmlns='http://www.w3.org/2001/10/synthesis' xml:lang='en-US'><voice name='%s'><prosody pitch='+0Hz' rate='+0%%' volume='+0%%'>%s</prosody></voice></speak>" % (voice, safe)

    @staticmethod
    def edge_chunks(text, limit=4000):
        sentences = re.split(r"(?<=[.!?])\s+", text.strip())
        parts = []
        current = ""
        for sentence in sentences:
            trial = (current + " " + sentence).strip()
            if len(trial.encode("utf-8")) > limit and current:
                parts.append(current)
                current = sentence
            else:
                current = trial
        if current:
            parts.append(current)
        return parts or [text]

    @staticmethod
    def ws_send(sock, payload):
        if isinstance(payload, str):
            payload = payload.encode("utf-8")
        mask = os.urandom(4)
        header = bytes([0x81])
        length = len(payload)
        if length < 126:
            header += bytes([0x80 | length])
        elif length < 65536:
            header += bytes([0x80 | 126]) + struct.pack(">H", length)
        else:
            header += bytes([0x80 | 127]) + struct.pack(">Q", length)
        masked = bytes(piece ^ mask[pos % 4] for pos, piece in enumerate(payload))
        sock.sendall(header + mask + masked)

    @staticmethod
    def ws_recv_exact(sock, count):
        data = b""
        while len(data) < count:
            piece = sock.recv(count - len(data))
            if not piece:
                raise BackendError("edge connection closed")
            data += piece
        return data

    @staticmethod
    def ws_recv(sock):
        opcode = 0
        payload = b""
        first = True
        while True:
            head = EdgeTts.ws_recv_exact(sock, 2)
            done = head[0] & 0x80
            part = head[0] & 0x0F
            masked = head[1] & 0x80
            length = head[1] & 0x7F
            if length == 126:
                length = struct.unpack(">H", EdgeTts.ws_recv_exact(sock, 2))[0]
            elif length == 127:
                length = struct.unpack(">Q", EdgeTts.ws_recv_exact(sock, 8))[0]
            key = EdgeTts.ws_recv_exact(sock, 4) if masked else b""
            body = EdgeTts.ws_recv_exact(sock, length) if length else b""
            if masked:
                body = bytes(piece ^ key[pos % 4] for pos, piece in enumerate(body))
            if part == 0x8:
                raise BackendError("edge closed stream")
            if part == 0x9:
                sock.sendall(bytes([0x8A, 0x00]))
                continue
            if part in (0x1, 0x2) and first:
                opcode = part
                first = False
            payload += body
            if done:
                return opcode, payload

    @staticmethod
    def edge_skew(header_text):
        match = re.search(r"Date:\s*(.+?)\r?$", header_text, re.MULTILINE | re.IGNORECASE)
        if not match:
            return 0
        try:
            stamp = time.mktime(time.strptime(match.group(1).strip(), "%a, %d %b %Y %H:%M:%S %Z"))
            return int(stamp - time.time())
        except (ValueError, OverflowError):
            return 0

    @staticmethod
    def edge_connect():
        skew = 0
        for attempt in range(2):
            path = "%s?TrustedClientToken=%s&ConnectionId=%s&Sec-MS-GEC=%s&Sec-MS-GEC-Version=1-%s" % (Consts.EDGE_WS_PATH, Consts.EDGE_TRUSTED_TOKEN, uuid.uuid4().hex, EdgeTts.edge_token(skew), Consts.EDGE_CHROMIUM)
            raw = socket.create_connection((Consts.EDGE_HOST, 443), timeout=20)
            sock = ssl.create_default_context().wrap_socket(raw, server_hostname=Consts.EDGE_HOST)
            sock.settimeout(30)
            key = base64.b64encode(os.urandom(16)).decode("ascii")
            muid = "".join(random.choice("0123456789ABCDEF") for _ in range(32))
            handshake = "\r\n".join([
                "GET %s HTTP/1.1" % path,
                "Host: %s" % Consts.EDGE_HOST,
                "Upgrade: websocket",
                "Connection: Upgrade",
                "Sec-WebSocket-Version: 13",
                "Sec-WebSocket-Key: %s" % key,
                "Pragma: no-cache",
                "Cache-Control: no-cache",
                "Origin: chrome-extension://jdiccldimpdaibmpdkjnbmckianbfold",
                "User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36 Edg/143.0.0.0",
                "Accept-Encoding: gzip, deflate, br, zstd",
                "Accept-Language: en-US,en;q=0.9",
                "Cookie: muid=%s;" % muid,
                "", "",
            ])
            sock.sendall(handshake.encode("ascii"))
            answer = b""
            while b"\r\n\r\n" not in answer:
                piece = sock.recv(4096)
                if not piece:
                    break
                answer += piece
            header_text = answer.decode("latin-1", "replace")
            if " 101 " in header_text.split("\r\n", 1)[0]:
                return sock
            sock.close()
            if " 403 " in header_text and attempt == 0:
                skew = EdgeTts.edge_skew(header_text)
                continue
            raise BackendError("edge handshake failed: %s" % header_text.split("\r\n", 1)[0][:200])
        raise BackendError("edge handshake failed after retry")

    @staticmethod
    def edge_synthesize(text, voice):
        sock = EdgeTts.edge_connect()
        audio = bytearray()
        try:
            config = "X-Timestamp:%s\r\nContent-Type:application/json; charset=utf-8\r\nPath:speech.config\r\n\r\n{\"context\":{\"synthesis\":{\"audio\":{\"metadataoptions\":{\"sentenceBoundaryEnabled\":\"false\",\"wordBoundaryEnabled\":\"true\"},\"outputFormat\":\"audio-24khz-48kbitrate-mono-mp3\"}}}}" % EdgeTts.edge_timestamp()
            EdgeTts.ws_send(sock, config)
            for chunk in EdgeTts.edge_chunks(text):
                message = "X-RequestId:%s\r\nContent-Type:application/ssml+xml\r\nX-Timestamp:%sZ\r\nPath:ssml\r\n\r\n%s" % (uuid.uuid4().hex, EdgeTts.edge_timestamp(), EdgeTts.edge_ssml(chunk, voice))
                EdgeTts.ws_send(sock, message)
                while True:
                    opcode, payload = EdgeTts.ws_recv(sock)
                    if opcode == 0x2 and len(payload) > 2:
                        head_len = int.from_bytes(payload[0:2], "big")
                        audio += payload[2 + head_len:]
                    elif opcode == 0x1 and "Path:turn.end" in payload.decode("utf-8", "replace"):
                        break
            return bytes(audio)
        finally:
            try:
                sock.close()
            except OSError:
                pass


class Audio:
    """Audio."""

    @staticmethod
    def play_audio(path):
        if sys.platform == "darwin" and shutil.which("afplay"):
            runner = ["afplay", path]
        elif shutil.which("paplay"):
            runner = ["paplay", path]
        elif shutil.which("aplay"):
            runner = ["aplay", "-q", path]
        elif shutil.which("pw-play"):
            runner = ["pw-play", path]
        elif sys.platform == "win32":
            runner = ["powershell", "-c", "(New-Object Media.SoundPlayer '%s').PlaySync();" % path]
        else:
            return False
        try:
            subprocess.run(runner, timeout=300, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False)
            return True
        except (OSError, subprocess.SubprocessError):
            return False

    @staticmethod
    def listen_audio(seconds, workdir):
        recorder = next((item for item in Consts.RECORDERS if shutil.which(item[0])), None)
        transcriber = next((name for name in Consts.TRANSCRIBERS if shutil.which(name)), "")
        if recorder is None or not transcriber:
            missing = []
            if recorder is None:
                missing.append("audio recorder (install arecord, sox, or ffmpeg)")
            if not transcriber:
                missing.append("transcriber (install whisper-cpp or whisper)")
            return "listening unavailable, missing: " + ", ".join(missing)
        path = os.path.join(workdir, "listen-%d.wav" % int(time.time()))
        command = [part.replace("{seconds}", str(seconds)).replace("{path}", path) for part in recorder[1]]
        try:
            subprocess.run(command, timeout=seconds + 15, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False)
            done = subprocess.run([transcriber, path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=180)
        except (OSError, subprocess.SubprocessError) as exc:
            return "listening failed: " + Text.short_error(exc)
        finally:
            try:
                os.remove(path)
            except OSError:
                pass
        if done.returncode != 0:
            return "transcriber error: " + (done.stderr or "")[:500]
        text = (done.stdout or "").strip()
        return Text.truncate(text, 3000, 500) if text else "heard nothing"