#!/usr/bin/env python3
"""Mock AI Agent: 加入房间, 发布 440Hz 间歇音(模拟 TTS 输出),
订阅 SIP 参与者音频并统计 RMS(模拟 ASR 输入), 验证双向语音。
依赖: pip install livekit
"""
import argparse, array, asyncio, math, struct, time
from livekit import rtc
from livekit.api import AccessToken, VideoGrants

def rms(data) -> float:
    buf = bytes(data)[:4096]
    a = array.array("h"); a.frombytes(buf)
    if not a: return 0.0
    return (sum(x * x for x in a) / len(a)) ** 0.5

async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="ws://127.0.0.1:7880")
    ap.add_argument("--api-key", default="devkey")
    ap.add_argument("--api-secret", default="f4f6c1a9e2b74d5893a0c8e17d2b6a45b7c93f1d6e284a05c9b3d7f2e6a81c04")
    ap.add_argument("--room", required=True)
    ap.add_argument("--identity", default="mock-agent")
    ap.add_argument("--duration", type=float, default=20)
    args = ap.parse_args()

    token = (AccessToken(args.api_key, args.api_secret)
             .with_identity(args.identity).with_name("MockAgent")
             .with_grants(VideoGrants(room_join=True, room=args.room,
                                       can_publish=True, can_subscribe=True)))

    room = rtc.Room()
    stats = {"frames": 0, "rms_sum": 0.0, "src": None}
    done = asyncio.Event()

    def on_sub(track, pub, participant):
        if track.kind != rtc.TrackKind.KIND_AUDIO:
            return
        print(f"[agent] subscribed: {participant.identity} track={track.sid}", flush=True)
        async def consume():
            stream = rtc.AudioStream(track)
            async for ev in stream:
                f = ev.frame
                stats["frames"] += 1
                stats["rms_sum"] += rms(f.data)
                if stats["frames"] % 250 == 0:
                    print(f"[agent] RX frames={stats['frames']} avg_rms={stats['rms_sum']/stats['frames']:.0f}", flush=True)
        stats["src"] = asyncio.create_task(consume())

    room.on("track_subscribed", on_sub)
    await room.connect(args.url, token.to_jwt())
    peers = [p.identity for p in room.remote_participants.values()]
    print(f"[agent] joined room={args.room} peers={peers}", flush=True)

    src = rtc.AudioSource(48000, 1)
    track = rtc.LocalAudioTrack.create_audio_track("agent-tts", src)
    await room.local_participant.publish_track(track)
    print("[agent] publishing 440Hz tone (500ms on / 500ms off)", flush=True)

    async def tone():
        t, end = 0.0, time.time() + args.duration
        while time.time() < end:
            samples = []
            on = (t % 1.0) < 0.5
            for _ in range(960):  # 20ms @ 48kHz
                v = int(6000 * math.sin(2 * math.pi * 440 * t)) if on else 0
                samples.append(v); t += 1 / 48000
            buf = struct.pack("<960h", *samples)
            await src.capture_frame(rtc.AudioFrame(buf, 48000, 1, 960))
            await asyncio.sleep(0.02)

    await tone()
    n, r = stats["frames"], (stats["rms_sum"] / stats["frames"]) if stats["frames"] else 0.0
    print(f"RESULT room={args.room} rx_frames={n} rx_avg_rms={r:.0f} "
          f"bidirectional={'YES' if n > 50 and r > 100 else 'NO'}", flush=True)
    await room.disconnect()

if __name__ == "__main__":
    asyncio.run(main())