import asyncio
import json
import threading
import cv2
import numpy as np
import win32gui
import win32ui
import win32con
import win32api

import aiortc
from aiortc import MediaStreamTrack, RTCPeerConnection, RTCSessionDescription
from aiortc.contrib.media import MediaBlackhole, MediaPlayer, MediaRecorder
from aiortc.sdp import SessionDescription
import webview

# -------------------------- 2. 兼容VideoFrame导入 --------------------------
try:from aiortc.contrib.media import VideoFrame,AudioFrame
except ImportError:from aiortc import VideoFrame,AudioFrame


# -------------------------- 1. Windows 11高DPI适配 --------------------------
ctypes.windll.shcore.SetProcessDpiAwareness(2)  # 解决高分屏采集偏移问题

# -------------------------- 2. 兼容VideoFrame导入 --------------------------
try:from aiortc.contrib.media import VideoFrame,AudioFrame
except ImportError:from aiortc import VideoFrame,AudioFrame

# -------------------------- 2. 屏幕视频采集轨道 --------------------------
class ScreenVideoTrack(MediaStreamTrack):
    kind = "video"
    def __init__(self, fps=30, scale_factor=1.0):
        super().__init__()
        self.fps = fps
        self.scale_factor = scale_factor
        self.stop_flag = False
        self.pts = 0
        self.frame_interval = 1.0 / fps

        # 获取真实屏幕分辨率（适配高DPI）
        self.user32 = ctypes.windll.user32
        self.screen_width = self.user32.GetSystemMetrics(0)
        self.screen_height = self.user32.GetSystemMetrics(1)

    def capture_screen(self):
        """Windows 11原生API采集屏幕（低延迟）"""
        left, top, width, height = 0, 0, self.screen_width, self.screen_height

        # 1. 创建设备上下文（DC）
        hdesktop = win32gui.GetDesktopWindow()
        hwnd_dc = win32gui.GetWindowDC(hdesktop)
        mfc_dc = win32ui.CreateDCFromHandle(hwnd_dc)
        save_dc = mfc_dc.CreateCompatibleDC()

        # 2. 复制屏幕内容到位图
        save_bitmap = win32ui.CreateBitmap()
        save_bitmap.CreateCompatibleBitmap(mfc_dc, width, height)
        save_dc.SelectObject(save_bitmap)
        save_dc.BitBlt((0, 0), (width, height), mfc_dc, (left, top), win32con.SRCCOPY)

        # 3. 转换为numpy数组（BGR格式）
        bmp_data = save_bitmap.GetBitmapBits(True)
        frame = np.frombuffer(bmp_data, dtype=np.uint8).reshape((height, width, 4))
        frame = frame[:, :, :3]  # 去掉Alpha通道
        frame = frame[:, :, ::-1]  # BGRA → BGR（适配aiortc）

        # 4. 缩放（可选，降低文件大小）
        if self.scale_factor != 1.0:
            new_width = int(width * self.scale_factor)
            new_height = int(height * self.scale_factor)
            frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_AREA)

        # 5. 释放资源（避免内存泄漏）
        win32gui.DeleteObject(save_bitmap.GetHandle())
        save_dc.DeleteDC()
        mfc_dc.DeleteDC()
        win32gui.ReleaseDC(hdesktop, hwnd_dc)

        return frame

    async def recv(self):
        """aiortc核心方法：持续返回视频帧"""
        if self.stop_flag:
            raise StopAsyncIteration
        # 控制帧率
        await asyncio.sleep(self.frame_interval)  
        # 获取windows11 原始的视频帧
        frame_data = self.capture_screen()

        # 转换为aiortc VideoFrame
        video_frame = self._convert_to_video_frame(frame_data)

        video_frame.pts = int(self.pts)
        video_frame.time_base = np.array([1, self.fps])
        self.pts += 1

        return video_frame

    def _convert_to_video_frame(self, img):
        """将numpy数组转换为aiortc的VideoFrame"""  
        return VideoFrame.from_ndarray(img, format="bgr24")

    def stop(self):
        """停止屏幕采集"""
        self.stop_flag = True

# -------------------------- 3. 系统音频采集轨道 --------------------------
class SystemAudioTrack(MediaStreamTrack):
    kind = "audio"

    def __init__(self, sample_rate=48000, channels=2):
        super().__init__()
        self.sample_rate = sample_rate
        self.channels = channels
        self.stop_flag = False
        self.pts = 0
        self.audio_queue = asyncio.Queue(maxsize=10)  # 音频数据队列

        # 启动麦克风/系统音频采集（这里默认采集麦克风，系统音频见备注）
        self._start_audio_capture()

    def _start_audio_capture(self):
        """启动音频采集（麦克风）"""
        def audio_callback(indata, frames, time, status):
            if status or self.stop_flag:
                return
            # 将音频数据转为float32格式（aiortc要求）
            audio_data = indata.astype(np.float32)
            self.audio_queue.put_nowait(audio_data)

        # 打开音频输入流
        self.audio_stream = sd.InputStream(
            samplerate=self.sample_rate,
            channels=self.channels,
            callback=audio_callback,
            blocksize=1024
        )
        self.audio_stream.start()

    async def recv(self):
        """aiortc核心方法：持续返回音频帧"""
        if self.stop_flag and self.audio_queue.empty():
            raise StopAsyncIteration

        # 从队列获取音频数据
        audio_data = await self.audio_queue.get()
        sample_count = len(audio_data)

        # 转换为aiortc AudioFrame
        audio_frame = AudioFrame(
            samples=audio_data.T,  # 转置为 (channels, samples)
            sample_rate=self.sample_rate,
            channels=self.channels
        )
        audio_frame.pts = int(self.pts)
        audio_frame.time_base = np.array([1, self.sample_rate])
        self.pts += sample_count

        return audio_frame

    def stop(self):
        """停止音频采集"""
        self.stop_flag = True
        self.audio_stream.stop()
        self.audio_stream.close()

 
class WebRTCService:
    def __init__(self):
        super().__init__()
        self.peer=None
        self.channel = None
        self.screen_track = None
        self.loop = asyncio.new_event_loop()
        self.thread = None
    
    def start_async_loop(self):
        """在子线程运行asyncio循环"""
        asyncio.set_event_loop(self.loop)
        self.loop.run_forever()
    
    # 1. 初始化 RTCPeerConnection
    async def init_connection(self):
        # 初始化 RTCPeerConnection
        if self.peer :self.peer =None
        # 创建 RTCPeerConnection 实例
        self.peer = RTCPeerConnection()
        
        # 添加事件监听器
        @self.peer.on("connectionstatechange")
        async def on_connectionstatechange():
            print(f"连接状态: {self.peer.connectionState}")
    
    async def recv(self):
        """捕获屏幕帧并转换为 aiortc 可处理的 VideoFrame"""
        if not self.running:
            raise RuntimeError("Track stopped")

        # 1. 捕获屏幕（返回 numpy 数组）
        img = np.array(self.sct.grab(self.monitor))
        # 2. 转换格式：mss 捕获的是 BGRA，转为 BGR（OpenCV 格式）
        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        # 3. 转换为 VideoFrame（aiortc 要求的格式）
        frame = VideoFrame.from_ndarray(img, format="bgr24")
        
        # 调整时间戳（保证流的时序）
        frame.pts = self.pts
        frame.time_base = self.time_base
        self.pts += 1
        
        return frame

    # 2. 获取用户媒体
    async def setup_media(self):
        # 使用摄像头和麦克风
        # player = MediaPlayer('/dev/video0')  # Linux 摄像头
        # 或者使用文件
        media_flag = True
        if media_flag:
            player = MediaPlayer('screen_60fps.mp4')
            return player.audio, player.video
        else:
            player_video = ScreenVideoTrack(fps=30)
            player_audio = SystemAudioTrack()
            return player_audio, player_video
    
    # 3. 创建 Offer
    async def create_offer(self):
        # 添加音视频轨道
        audio_track, video_track = await self.setup_media()
        
        if audio_track:
            self.peer.addTrack(audio_track)
        if video_track:
            self.peer.addTrack(video_track)
        
        # 创建 offer
        offer = await self.peer.createOffer()
        await self.peer.setLocalDescription(offer)
        
        return offer
    
    # 3. 创建 接收 offer 并创建 Answer
    async def handle_remote_offer(self, offer_sdp):
        # 设置远程描述
        offer = RTCSessionDescription(sdp=offer_sdp, type="offer")
        await self.peer.setRemoteDescription(offer)
        
        # 创建 answer
        answer = await self.peer.createAnswer()
        await self.peer.setLocalDescription(answer)
        return answer

    async def handle_remote_answer(self,answer_sdp):
        await self.peer.setRemoteDescription(obj)
        
        # 监听ICE候选，发送到信令服务器
        @self.peer.on("icecandidate")
        async def on_icecandidate(candidate):
            if candidate:
                # await signaling.send(candidate)
                pass

    # 4. 数据通道使用
    async def setup_data_channel(self,pc):
        # 初始化 数据通道
        if self.channel : self.channel = None
        # 创建数据通道
        self.channel = self.peer.createDataChannel("chat")
        
        @self.channel.on("open")
        def on_open():
            print("数据通道已打开")
            # 发送消息
            self.channel.send("Hello from aiortc!")
        
        @self.channel.on("message")
        def on_message(message):
            print(f"收到消息: {message}")
        
        return self.channel
    
if __name__ == "__main__":
    rtc=rtc()
    rtc.setup_media()
    pass