ForcePilot/backend/package/yuxi/channels/adapters/wechat/voice.py

from __future__ import annotations

from typing import Any, Literal

import httpx

from yuxi.channels.models import DeliveryResult

from .bridge import BridgeClient
from .mp.client import MPClient
from .wecom.client import WeComClient

VOICE_FORMATS = ("mp3", "amr", "silk")
DEFAULT_VOICE_FORMAT: Literal["mp3", "amr", "silk"] = "amr"
MAX_VOICE_DURATION_SECONDS = 30


def detect_voice_duration(voice_data: bytes) -> float:
    if voice_data[:3] == b"ID3" or (voice_data[0] == 0xFF and (voice_data[1] & 0xE0) == 0xE0):
        return _mp3_duration(voice_data)
    if voice_data[:6] == b"#!AMR\n":
        return _amr_duration(voice_data)
    return 0.0


def _mp3_duration(data: bytes) -> float:
    frame_count = 0
    i = 0
    while i < len(data) - 4:
        if data[i] == 0xFF and (data[i + 1] & 0xE0) == 0xE0:
            frame_count += 1
            i += 4
        else:
            i += 1
    return frame_count * 0.026


def _amr_duration(data: bytes) -> float:
    frame_count = 0
    i = 6
    while i < len(data):
        header = data[i]
        frame_size_map = {0: 13, 1: 14, 2: 16, 3: 18, 4: 20, 5: 21, 6: 27, 7: 32}
        frame_type = (header >> 3) & 0x0F
        frame_size = frame_size_map.get(frame_type, 0)
        if frame_size == 0:
            i += 1
            frame_count += 1
        else:
            i += frame_size + 1
            frame_count += 1
    return frame_count * 0.020


def validate_voice_data(
    voice_data: bytes,
    config: dict[str, Any] | None = None,
) -> tuple[bool, str]:
    max_duration = (config or {}).get("tts", {}).get("max_duration_seconds", MAX_VOICE_DURATION_SECONDS)
    if not voice_data:
        return False, "Empty voice data"
    duration = detect_voice_duration(voice_data)
    if duration > 0 and duration > max_duration:
        return False, f"Voice duration {duration:.1f}s exceeds limit {max_duration}s"
    return True, ""


def ensure_format(
    voice_data: bytes,
    target_format: str = "amr",
    source_format: str = "mp3",
) -> bytes:
    if source_format == target_format:
        return voice_data

    if target_format not in VOICE_FORMATS:
        raise ValueError(
            f"Unsupported target format: {target_format}, must be one of {VOICE_FORMATS}"
        )

    raise ValueError(
        f"Voice format conversion from {source_format} to {target_format} is not supported. "
        f"Please provide voice data in {target_format} format."
    )


def build_voice_payload(
    to_user: str,
    media_id: str,
    agent_id: str | None = None,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "touser": to_user,
        "msgtype": "voice",
        "voice": {"media_id": media_id},
    }
    if agent_id:
        payload["agentid"] = agent_id
    return payload


async def send_voice_wecom(
    client: WeComClient, http_client: httpx.AsyncClient, to_user: str, voice_data: bytes, agent_id: str
) -> DeliveryResult:
    try:
        voice_data = ensure_format(voice_data, target_format="amr")
        media_id = await client.upload_media(voice_data, "voice.amr", "voice")
    except Exception as e:
        return DeliveryResult(success=False, error=f"WeCom voice upload failed: {e}")

    token = await client.get_access_token()
    url = "https://qyapi.weixin.qq.com/cgi-bin/message/send"
    payload = build_voice_payload(to_user, media_id, agent_id)
    try:
        resp = await http_client.post(url, params={"access_token": token}, json=payload)
        data = resp.json()
        if data.get("errcode") == 0:
            return DeliveryResult(success=True)
        return DeliveryResult(success=False, error=data.get("errmsg", "Unknown"))
    except Exception as e:
        return DeliveryResult(success=False, error=str(e))


async def send_voice_mp(
    client: MPClient, http_client: httpx.AsyncClient, to_user: str, voice_data: bytes
) -> DeliveryResult:
    try:
        voice_data = ensure_format(voice_data, target_format="amr")
        media_id = await client.upload_temp_media(voice_data, "voice.amr", "voice")
    except Exception as e:
        return DeliveryResult(success=False, error=f"MP voice upload failed: {e}")

    token = await client.get_access_token()
    url = "https://api.weixin.qq.com/cgi-bin/message/custom/send"
    payload = build_voice_payload(to_user, media_id)
    try:
        resp = await http_client.post(url, params={"access_token": token}, json=payload)
        data = resp.json()
        if data.get("errcode") == 0:
            return DeliveryResult(success=True)
        return DeliveryResult(success=False, error=data.get("errmsg", "Unknown"))
    except Exception as e:
        return DeliveryResult(success=False, error=str(e))


async def send_voice_bridge(
    bridge_client: BridgeClient, chat_id: str, voice_data: bytes, is_group: bool = False
) -> DeliveryResult:
    import base64

    voice_data = ensure_format(voice_data, target_format="amr")
    payload = {
        "chat_id": chat_id,
        "msg_type": 34,
        "is_group": is_group,
        "voice_data": base64.b64encode(voice_data).decode("ascii"),
    }
    return await bridge_client.send_media_message(payload)