import asyncio
import audioop
from array import array
import json
import sys
import time
import wave
import websockets

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8")

WS_URL = "ws://localhost:10095"
TARGET_SAMPLE_RATE = 16000
TARGET_SAMPLE_WIDTH = 2
CHUNK_SECONDS = 0.06
FINAL_RESULT_TIMEOUT = 60
WAV_PATH = r"E:\ZKYNLP\gonwgendoc\fastapi_wss\fastapi_wss\src\多人会议90分钟.wav"


def downmix_to_mono_pcm16(chunk: bytes, channels: int) -> bytes:
    if channels == 1:
        return chunk

    samples = array("h")
    samples.frombytes(chunk)
    if sys.byteorder != "little":
        samples.byteswap()

    mono = array("h")
    for index in range(0, len(samples), channels):
        frame = samples[index:index + channels]
        if len(frame) == channels:
            mono.append(int(sum(frame) / channels))

    if sys.byteorder != "little":
        mono.byteswap()
    return mono.tobytes()


def validate_wav_for_streaming(wf: wave.Wave_read):
    sample_rate = wf.getframerate()
    sample_width = wf.getsampwidth()
    channels = wf.getnchannels()

    if sample_width != TARGET_SAMPLE_WIDTH:
        raise ValueError(f"WAV sample width must be {TARGET_SAMPLE_WIDTH} bytes, got {sample_width} bytes")
    if channels < 1:
        raise ValueError(f"WAV channels must be >= 1, got {channels}")

    if sample_rate != TARGET_SAMPLE_RATE:
        print(f"Input WAV is {sample_rate}Hz; resampling to {TARGET_SAMPLE_RATE}Hz before sending.")
    if channels > 1:
        print(f"Input WAV has {channels} channels; downmixing to mono before sending.")
    source_frames_per_chunk = max(1, int(sample_rate * CHUNK_SECONDS))
    return sample_rate, channels, source_frames_per_chunk


class Pcm16Mono16kConverter:
    def __init__(self, source_rate: int, channels: int):
        self.source_rate = source_rate
        self.channels = channels
        self.rate_state = None

    def convert(self, chunk: bytes) -> bytes:
        chunk = downmix_to_mono_pcm16(chunk, self.channels)
        if self.source_rate != TARGET_SAMPLE_RATE:
            chunk, self.rate_state = audioop.ratecv(
                chunk,
                TARGET_SAMPLE_WIDTH,
                1,
                self.source_rate,
                TARGET_SAMPLE_RATE,
                self.rate_state,
            )
        return chunk

def is_final_asr_message(data: dict) -> bool:
    mode = data.get("mode")
    if mode not in {"2pass-offline", "offline-speaker"}:
        return False
    return bool(data.get("is_final", True))


async def recv_loop(ws, audio_sent_event: asyncio.Event, final_result_event: asyncio.Event, stats: dict):
    async for msg in ws:
        print("RECV:", msg)
        if not isinstance(msg, str):
            continue
        try:
            data = json.loads(msg)
        except json.JSONDecodeError:
            continue
        if not is_final_asr_message(data):
            continue

        now = time.perf_counter()
        stats["last_final_time"] = now
        if audio_sent_event.is_set() and "first_final_after_send_time" not in stats:
            stats["first_final_after_send_time"] = now
            final_result_event.set()


def print_timing_stats(stats: dict):
    end_time = stats.get("first_final_after_send_time") or stats.get("end_time") or time.perf_counter()
    total_elapsed = end_time - stats["start_time"]
    send_elapsed = stats.get("send_end_time", end_time) - stats.get("send_start_time", stats["start_time"])
    tail_elapsed = None
    if "send_end_time" in stats and "first_final_after_send_time" in stats:
        tail_elapsed = stats["first_final_after_send_time"] - stats["send_end_time"]

    print("\n===== ASR Timing =====")
    print(f"Audio duration: {stats.get('audio_duration_sec', 0.0):.2f}s")
    print(f"Audio send time: {send_elapsed:.2f}s")
    if tail_elapsed is not None:
        print(f"Post-send ASR wait: {tail_elapsed:.2f}s")
    else:
        print("Post-send ASR wait: timeout/no final result")
    print(f"Total elapsed: {total_elapsed:.2f}s")
    if stats.get("audio_duration_sec"):
        rtf = total_elapsed / stats["audio_duration_sec"]
        print(f"RTF: {rtf:.3f}x")
    print("======================\n")

async def main():
    stats = {"start_time": time.perf_counter()}
    audio_sent_event = asyncio.Event()
    final_result_event = asyncio.Event()

    async with websockets.connect(
        WS_URL,
        max_size=None,
        ping_interval=None,
        ping_timeout=None,
    ) as ws:
        recv_task = asyncio.create_task(recv_loop(ws, audio_sent_event, final_result_event, stats))

        init = {
            "mode": "2pass",
            "chunk_size": [10, 10, 10],
            "chunk_interval": 10,
            "wav_name": "test",
            "is_speaking": True,
            "wav_format": "pcm",
            "itn": True
        }
        await ws.send(json.dumps(init, ensure_ascii=False))

        connection_closed = False

        with wave.open(WAV_PATH, "rb") as wf:
            source_rate, channels, source_frames_per_chunk = validate_wav_for_streaming(wf)
            stats["audio_duration_sec"] = wf.getnframes() / source_rate
            converter = Pcm16Mono16kConverter(source_rate, channels)

            stats["send_start_time"] = time.perf_counter()
            while True:
                chunk = wf.readframes(source_frames_per_chunk)
                if not chunk:
                    break
                chunk = converter.convert(chunk)
                try:
                    await ws.send(chunk)
                except websockets.exceptions.ConnectionClosed as e:
                    connection_closed = True
                    print(f"WebSocket closed while sending audio: {type(e).__name__}: {e}")
                    break
                await asyncio.sleep(CHUNK_SECONDS)

        if not connection_closed:
            try:
                await ws.send(json.dumps({"is_speaking": False}, ensure_ascii=False))
            except websockets.exceptions.ConnectionClosed as e:
                connection_closed = True
                print(f"WebSocket closed before sending end marker: {type(e).__name__}: {e}")

        stats["send_end_time"] = time.perf_counter()
        audio_sent_event.set()

        if connection_closed and not final_result_event.is_set():
            print("Connection closed before final ASR result was received.")
            stats["end_time"] = time.perf_counter()
            print_timing_stats(stats)
        else:
            try:
                await asyncio.wait_for(final_result_event.wait(), timeout=FINAL_RESULT_TIMEOUT)
            except asyncio.TimeoutError:
                print(f"Timed out waiting for final ASR result after {FINAL_RESULT_TIMEOUT}s.")
            finally:
                stats["end_time"] = time.perf_counter()
                print_timing_stats(stats)

        recv_task.cancel()
        await asyncio.gather(recv_task, return_exceptions=True)

asyncio.run(main())
 
