"""独立调用 VUILabs 实时同传 API；默认连接国内生产环境。"""
import argparse
import asyncio
import base64
import getpass
import http.client
import json
import math
import os
import ssl
import struct
import sys
import wave
from pathlib import Path
from urllib.parse import urlencode, urlsplit

DEFAULT_BASE_URL = "https://api.vuilabs.cn"
SESSION_TIMEOUT_SECONDS = 15 * 60 + 30
MAX_REFERENCE_BYTES = 4 * 1024 * 1024
REFERENCE_LANGUAGES = {"zh", "en", "ja", "ko", "pt", "de", "fr", "ru", "es", "it", "ms", "yue", "id", "th"}


def base_url(value):
    """只接受 HTTPS origin；凭据始终放在请求头。"""
    try:
        parsed = urlsplit(value)
        if (parsed.scheme != "https" or not parsed.hostname
                or parsed.username is not None or parsed.password is not None
                or parsed.path not in ("", "/") or parsed.query or parsed.fragment
                or any(char.isspace() for char in value)):
            raise ValueError
        parsed.port  # 同时验证显式端口。
    except ValueError:
        raise argparse.ArgumentTypeError("--base-url 须为不含凭据、路径、查询或片段的 HTTPS 地址") from None
    return "https://" + parsed.netloc


def list_voices(args, key, tls):
    params = {"source": args.source}
    if args.cursor:
        params["cursor"] = args.cursor
    connection = http.client.HTTPSConnection(urlsplit(args.base_url).netloc, timeout=20, context=tls)
    try:
        connection.request("GET", "/v1/interpret/voices?" + urlencode(params),
                           headers={"X-API-Key": key})
        response = connection.getresponse()
        result = json.loads(response.read())
        if response.status != 200 or result.get("succeed") is not True or result.get("error_code") != 0:
            raise ValueError(f"音色查询失败（HTTP {response.status}）：{result}")
        print(json.dumps(result["data"], ensure_ascii=False, indent=2))
    finally:
        connection.close()


async def send_audio(ws, wav):
    while chunk := wav.readframes(1600):
        await ws.send(chunk)
        await asyncio.sleep(len(chunk) / 32000)
    await ws.send('{"action":"finish"}')


async def receive_frame(ws, sender, deadline):
    """检测发送任务失败，但不把没有译文事件当成连接失活。"""
    receiver = asyncio.create_task(ws.recv())
    try:
        waiting = {receiver}
        if sender is not None:
            if sender.done():
                await sender
            else:
                waiting.add(sender)
        remaining = deadline - asyncio.get_running_loop().time()
        if remaining <= 0:
            raise TimeoutError("同传会话超过 15 分钟及 30 秒收尾等待上限")
        done, _ = await asyncio.wait(waiting, timeout=remaining, return_when=asyncio.FIRST_COMPLETED)
        if not done:
            raise TimeoutError("同传会话超过 15 分钟及 30 秒收尾等待上限")
        if sender is not None and sender in done:
            await sender
        remaining = deadline - asyncio.get_running_loop().time()
        return await asyncio.wait_for(receiver, timeout=max(0, remaining))
    except asyncio.TimeoutError as error:
        raise TimeoutError("同传会话超过 15 分钟及 30 秒收尾等待上限") from error
    finally:
        if not receiver.done():
            receiver.cancel()
        await asyncio.gather(receiver, return_exceptions=True)


def validate_reference_wav(data):
    if len(data) < 44 or data[:4] != b"RIFF" or data[8:12] != b"WAVE":
        raise ValueError("参考音频必须为普通 WAV")
    offset = 12
    metadata = None
    samples = None
    while offset + 8 <= len(data):
        size = struct.unpack_from("<I", data, offset + 4)[0]
        end = offset + 8 + size
        if end > len(data):
            raise ValueError("参考 WAV 数据不完整")
        chunk = data[offset + 8:end]
        if data[offset:offset + 4] == b"fmt ":
            if len(chunk) < 16:
                raise ValueError("参考 WAV 格式信息不完整")
            metadata = struct.unpack_from("<HHIIHH", chunk)
        elif data[offset:offset + 4] == b"data":
            samples = chunk
        offset = end + size % 2
    if metadata is None or not samples:
        raise ValueError("参考 WAV 缺少格式或采样数据")
    encoding, channels, rate, _, _, bits = metadata
    if encoding == 0xFFFE:
        raise ValueError(
            "暂不支持 WAVE_FORMAT_EXTENSIBLE（格式标记 0xFFFE）参考 WAV；"
            "请转换为普通 PCM16 WAV，例如："
            "ffmpeg -i original.wav -ac 1 -ar 16000 -c:a pcm_s16le reference.wav")
    valid_bits = (encoding == 1 and bits in (8, 16, 24, 32)) or (encoding == 3 and bits in (32, 64))
    if channels != 1 or not 8000 <= rate <= 48000 or not valid_bits:
        raise ValueError("参考 WAV 须为单声道、8–48 kHz PCM 或 float")
    width = bits // 8
    if len(samples) % width or not rate <= len(samples) // width <= rate * 30:
        raise ValueError("参考 WAV 时长须为 1–30 秒，且采样数据完整")
    if encoding == 3:
        non_silent = False
        for (value,) in struct.iter_unpack("<f" if bits == 32 else "<d", samples):
            if not math.isfinite(value):
                raise ValueError("参考 WAV 包含非有限采样")
            non_silent = non_silent or value != 0
    else:
        non_silent = any(value != (128 if bits == 8 else 0) for value in samples)
    if not non_silent:
        raise ValueError("参考音频不能为静音")


def load_voice_prompt(args):
    if not args.reference_audio:
        if args.reference_text is not None or args.reference_text_file or args.reference_language:
            raise ValueError("参考转写和参考语言须与 --reference-audio 一起使用")
        return None
    if (args.reference_text is None) == (args.reference_text_file is None):
        raise ValueError("直接参考音频须选择 --reference-text 或 --reference-text-file")
    if args.reference_text_file:
        with args.reference_text_file.open("rb") as source:
            raw_text = source.read(16004)
        text = raw_text.decode("utf-8-sig").strip()
        if len(raw_text) >= 16004:
            raise ValueError("参考转写超过 4000 字符")
    else:
        text = args.reference_text.strip()
    if not 1 <= len(text) <= 4000 or "\x00" in text:
        raise ValueError("参考转写须为匹配录音的 1–4000 个字符，且不能含 NUL")
    with args.reference_audio.open("rb") as source:
        data = source.read(MAX_REFERENCE_BYTES + 1)
    if len(data) > MAX_REFERENCE_BYTES:
        raise ValueError("参考音频不能超过 4 MiB")
    validate_reference_wav(data)
    prompt = {"audio": base64.b64encode(data).decode("ascii"), "text": text}
    if args.reference_language:
        prompt["language"] = args.reference_language
    return prompt


async def translate(args, key, tls, wav, output):
    from websockets.asyncio.client import connect

    start = {
        "action": "start", "source_lang": args.source_lang,
        "target_lang": args.target_lang, "enable_tts": not args.text_only,
        "enable_vad": args.enable_vad, "tts_backend": args.tts_backend,
        "voice_mode": "clone",
    }
    if args.voice:
        start.update(voice_mode="preset", voice=args.voice)
    if args.platform_voice_id:
        start.update(voice_mode="prompt", platform_voice_id=args.platform_voice_id)
    if getattr(args, "voice_prompt", None):
        start.update(voice_mode="prompt", voice_prompt=args.voice_prompt)
    async with connect(f"wss://{urlsplit(args.base_url).netloc}/v1/interpret/stream",
                       additional_headers={"X-API-Key": key},
                       ssl=tls, proxy=None, open_timeout=15,
                       ping_interval=20, ping_timeout=None, close_timeout=10) as ws:
        await ws.send(json.dumps(start))
        sender = None
        pending = None
        output_rate = None
        request_id = None
        tts_failed = False
        deadline = asyncio.get_running_loop().time() + SESSION_TIMEOUT_SECONDS
        try:
            while True:
                frame = await receive_frame(ws, sender, deadline)
                if pending is not None:
                    rate, samples = pending
                    if not isinstance(frame, bytes) or len(frame) != samples * 4:
                        raise ValueError("输出音频与元信息不匹配")
                    if output_rate is None:
                        output.setparams((1, 2, rate, 0, "NONE", "not compressed"))
                        output_rate = rate
                    if output_rate != rate:
                        raise ValueError("同一会话的输出采样率发生变化")
                    pcm = bytearray()
                    for (sample,) in struct.iter_unpack("<f", frame):
                        if not math.isfinite(sample):
                            raise ValueError("输出含无效音频采样")
                        pcm.extend(struct.pack("<h", round(max(-1, min(1, sample)) * 32767)))
                    output.writeframesraw(pcm)
                    pending = None
                    continue
                if not isinstance(frame, str):
                    raise ValueError("收到没有元信息的音频帧")
                event = json.loads(frame)
                kind = event.get("type")
                request_id = event.get("request_id") or request_id
                if kind == "error":
                    if event.get("code") == "voice_unavailable":
                        raise ValueError(f"同传失败：{event}。请先查询可用音色，再更换 --platform-voice-id。")
                    raise ValueError(f"同传失败：{event}")
                if kind == "started" and sender is None:
                    sender = asyncio.create_task(send_audio(ws, wav))
                if kind == "tts_error" or (kind == "speech_done" and event.get("code") == "tts_failed"):
                    tts_failed = True
                if kind in {"speech", "speech_chunk"}:
                    audio = event["audio"]
                    rate, samples = audio["sample_rate"], audio["samples"]
                    if audio["encoding"] != "pcm_f32le" or rate not in (16000, 24000) or type(samples) is not int or samples <= 0:
                        raise ValueError("输出音频格式不受支持")
                    pending = rate, samples
                if kind in {"session_opened", "started", "translation", "translation_complete",
                            "tts_error", "speech_done", "usage", "completed"}:
                    print(json.dumps(event, ensure_ascii=False), flush=True)
                if kind == "completed":
                    if sender is None or not sender.done():
                        raise ValueError("服务在输入结束前关闭了会话")
                    await sender
                    if tts_failed:
                        raise ValueError("部分语音合成失败，已交付的译文与用量请按 request_id 核对")
                    if output_rate is None and output:
                        raise ValueError("同传完成但未返回语音")
                    return
        except Exception as error:
            detail = str(error) or "未返回错误描述"
            suffix = f"；request_id={request_id}" if request_id else ""
            raise RuntimeError(f"{type(error).__name__}: {detail}{suffix}") from error
        finally:
            if sender is not None:
                sender.cancel()
                await asyncio.gather(sender, return_exceptions=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", type=base_url, default=DEFAULT_BASE_URL,
                        help="HTTPS 接入地址，放在子命令前；默认 %(default)s")
    commands = parser.add_subparsers(dest="command", required=True)
    voices = commands.add_parser("voices", help="查询音色")
    voices.add_argument("--source", choices=("model", "system", "personal"), default="model")
    voices.add_argument("--cursor", help="个人音色分页游标")
    stream = commands.add_parser(
        "stream",
        help="用 WAV 调用实时同传",
        epilog="预置音色示例：python3 interpret_client.py stream input.wav --target-lang en --voice michael --output michael-en.wav",
    )
    stream.add_argument("audio", type=Path, help="PCM16、16 kHz、单声道 WAV")
    stream.add_argument("--output", type=Path, default=Path("interpret-output.wav"))
    stream.add_argument("--source-lang", default="zh")
    stream.add_argument("--target-lang", default="en")
    stream.add_argument("--text-only", action="store_true")
    stream.add_argument("--enable-vad", action="store_true")
    stream.add_argument("--tts-backend", choices=("lunasi-tts",), default="lunasi-tts")
    voice = stream.add_mutually_exclusive_group()
    voice.add_argument("--voice", help="模型预置音色 ID")
    voice.add_argument("--platform-voice-id", help="平台系统或个人音色 ID")
    voice.add_argument("--reference-audio", type=Path, help="直接参考的普通 WAV，单声道、8–48 kHz、1–30 秒")
    reference_text = stream.add_mutually_exclusive_group()
    reference_text.add_argument("--reference-text", help="与参考录音逐字匹配的转写")
    reference_text.add_argument("--reference-text-file", type=Path, help="UTF-8 参考转写文件")
    stream.add_argument("--reference-language", choices=sorted(REFERENCE_LANGUAGES), help="参考录音实际语言；省略时按接口默认 en")
    args = parser.parse_args()
    if args.command == "stream":
        if args.text_only and (args.voice or args.platform_voice_id or args.reference_audio):
            parser.error("仅译文模式不能指定音色或参考音频")
        try:
            args.voice_prompt = load_voice_prompt(args)
        except (ValueError, OSError) as error:
            parser.error(str(error))
    key = os.environ.get("VUILABS_API_KEY", "").strip() or os.environ.get("INTERPRET_API_KEY", "").strip()
    if not key:
        if not sys.stdin.isatty():
            parser.error("非交互环境请设置 VUILABS_API_KEY 或 INTERPRET_API_KEY")
        try:
            key = getpass.getpass("VUILabs API Key（须与接入环境匹配）: ").strip()
        except (EOFError, KeyboardInterrupt):
            parser.error("未读取到 API Key，请设置 VUILABS_API_KEY 后重试")
    if not key.strip():
        parser.error("请提供所选环境的 VUILabs API Key")
    try:
        import certifi

        tls = ssl.create_default_context(cafile=certifi.where())
        if args.command == "voices":
            list_voices(args, key.strip(), tls)
            return 0
        with wave.open(str(args.audio), "rb") as wav:
            if (wav.getnchannels(), wav.getsampwidth(), wav.getframerate()) != (1, 2, 16000):
                raise ValueError("输入须为 PCM16、16 kHz、单声道 WAV")
            if not 0 < wav.getnframes() < 15 * 60 * 16000:
                raise ValueError("音频时长须大于 0 且小于 15 分钟")
            if args.text_only:
                asyncio.run(translate(args, key.strip(), tls, wav, None))
            else:
                with args.output.open("xb") as file:
                    try:
                        with wave.open(file, "wb") as output:
                            output.setparams((1, 2, 24000, 0, "NONE", "not compressed"))
                            asyncio.run(translate(args, key.strip(), tls, wav, output))
                    except BaseException:
                        args.output.unlink()
                        raise
                print(f"输出音频已保存到 {args.output}")
        return 0
    except Exception as error:
        detail = str(error) or "未返回错误描述"
        print(f"调用失败：{type(error).__name__}: {detail.replace(key.strip(), '***')}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())