import logging
import asyncio
import cv2
import numpy as np
from time import perf_counter
from livekit import rtc, api

# 直接写死你的 API_KEY 和 API_SECRET
API_KEY = "devkey"
API_SECRET = "secret"
URL = "ws://192.168.2.236:7880"  # 直接使用你的 LIVEKIT_URL
IDENTITY = "win-client"
ROOM_NAME = "temi-room"

# 生成 token
def generate_token():
    token = api.AccessToken(API_KEY, API_SECRET) \
        .with_identity(IDENTITY) \
        .with_name("Monitor Client1") \
        .with_grants(api.VideoGrants(
            room_join=True,
            room=ROOM_NAME,
        )).to_jwt()
    return token

async def main():
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)
    room = rtc.Room()

    try:
        # 生成连接房间的 token
        token = generate_token()

        # 连接到 LiveKit 房间
        await room.connect(URL, token)
        logger.info("Connected to room: %s", room.name)

        # 监听参与者连接事件
        @room.on("participant_connected")
        def on_participant_connected(participant: rtc.RemoteParticipant):
            logger.info("Participant connected: %s %s", participant.sid, participant.identity)

        async def receive_frames(stream: rtc.VideoStream, track_sid: str):
            window_name = f"LiveKit Video Stream ({track_sid})"
            cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE)
            last_time = perf_counter()
            frame_count = 0
            fps = 0.0
            
            try:
                async for event in stream:
                    frame = event.frame
                    rgb_frame = frame.convert(rtc.VideoBufferType.RGB24)
                    arr = np.frombuffer(rgb_frame.data, dtype=np.uint8)
                    try:
                        img = arr.reshape((rgb_frame.height, rgb_frame.width, 3))
                        img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
                        frame_count += 1
                        now = perf_counter()
                        elapsed = now - last_time
                        if elapsed >= 1.0:
                            fps = frame_count / elapsed
                            frame_count = 0
                            last_time = now
                        cv2.putText(
                            img_bgr,
                            f"FPS: {fps:.1f}",
                            (10, 30),
                            cv2.FONT_HERSHEY_SIMPLEX,
                            1.0,
                            (0, 255, 0),
                            2,
                            cv2.LINE_AA,
                        )
                        cv2.imshow(window_name, img_bgr)
                        if cv2.waitKey(1) & 0xFF == ord('q'):
                            break
                    except Exception as e:
                        logger.error("Error processing frame: %s", e)
            finally:
                cv2.destroyWindow(window_name)

        # 监听 track 订阅事件
        @room.on("track_subscribed")
        def on_track_subscribed(track: rtc.Track, publication: rtc.RemoteTrackPublication, participant: rtc.RemoteParticipant):
            logger.info("Track subscribed: %s", publication.sid)
            if track.kind == rtc.TrackKind.KIND_VIDEO:
                video_stream = rtc.VideoStream(track)
                asyncio.ensure_future(receive_frames(video_stream, publication.sid))

        # 保持连接并处理事件
        await asyncio.Event().wait()

    except Exception as e:
        logger.error("Failed to connect or handle events: %s", e)

    finally:
        # 确保断开连接
        await room.disconnect()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("连接已被用户中断")