import asyncio
import logging
from aiortc import RTCPeerConnection, RTCSessionDescription
import aiohttp
from ocr.ocr_violation_detector import OCRViolationDetector

import logging

# 创建检测器实例
detector = OCRViolationDetector(
        forbidden_words_path=r"D:\Git\bin\video\ocr\forbidden_words.txt",
        ocr_confidence_threshold=0.7,
        log_level=logging.INFO,
        log_file="ocr_detection.log"
)

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("whep_video_puller")


async def whep_pull_video_stream(ip,whep_url):
    """
    通过WHEP从指定URL拉取视频流并在收到每一帧时打印消息

    Args:
        whep_url: WHEP端点的URL
    """
    pc = RTCPeerConnection()

    # 添加连接状态变化监听
    @pc.on("connectionstatechange")
    async def on_connectionstatechange():
        print(f"连接状态: {pc.connectionState}")

    # 添加ICE连接状态变化监听
    @pc.on("iceconnectionstatechange")
    async def on_iceconnectionstatechange():
        print(f"ICE连接状态: {pc.iceConnectionState}")

    # 添加视频接收器
    pc.addTransceiver("video", direction="recvonly")

    # 处理接收到的视频轨道
    @pc.on("track")
    def on_track(track):
        print(f"接收到轨道: {track.kind}")
        if track.kind == "video":
            print(f"轨道ID: {track.id}")
            print(f"轨道就绪状态: {track.readyState}")
            # 创建异步任务来处理视频帧
            asyncio.ensure_future(handle_video_track(track))

    async def handle_video_track(track):
        """处理视频轨道，接收并打印每一帧"""
        frame_count = 0
        print("开始处理视频轨道...")

        while True:
            try:
                # 尝试接收帧
                frame = await track.recv()
                frame_count += 1
                print(f"收到原始帧 (第{frame_count}帧)")

                # 打印帧的基本信息
                if hasattr(frame, 'width') and hasattr(frame, 'height'):
                    print(f"  尺寸: {frame.width}x{frame.height}")
                if hasattr(frame, 'time_base'):
                    print(f"  时间基准: {frame.time_base}")
                if hasattr(frame, 'pts'):
                    print(f"  显示时间戳: {frame.pts}")

                has_violation, violations, confidences = OCRViolationDetector.detect(frame)

                # 输出检测结果
                if has_violation:
                    detector.logger.info(f"在图片中检测到 {len(violations)} 个违禁词:")
                    for word, conf in zip(violations, confidences):
                        detector.logger.info(f"- {word} (置信度: {conf:.4f})")
                else:
                    detector.logger.info("图片中未检测到违禁词")
            except Exception as e:
                print(f"接收帧时出错: {e}")
                # 等待一段时间后重试
                await asyncio.sleep(0.1)
                continue

    # 创建offer
    offer = await pc.createOffer()
    await pc.setLocalDescription(offer)

    print(f"本地SDP信息:\n{offer.sdp}")

    # 通过HTTP POST发送offer到WHEP端点
    async with aiohttp.ClientSession() as session:
        async with session.post(
                whep_url,
                data=offer.sdp,
                headers={"Content-Type": "application/sdp"}
        ) as response:
            if response.status != 201:
                print(f"WHEP服务器返回错误: {response.status}")
                print(f"响应内容: {await response.text()}")
                raise Exception(f"WHEP服务器返回错误: {response.status}")

            # 获取answer SDP
            answer_sdp = await response.text()

            # 创建RTCSessionDescription对象
            answer = RTCSessionDescription(sdp=answer_sdp, type="answer")

            print(f"收到远程SDP:\n{answer_sdp}")

    # 设置远程描述
    await pc.setRemoteDescription(answer)

    print("连接已建立，开始接收视频流...")

    # 保持连接，直到用户中断
    try:
        while True:
            await asyncio.sleep(1)
            # 检查连接状态
            print(f"当前连接状态: {pc.connectionState}")
    except KeyboardInterrupt:
        print("用户中断，关闭连接...")
    finally:
        await pc.close()


if __name__ == "__main__":
    # 替换为你的WHEP端点URL
    WHEP_URL = "http://192.168.110.25:1985/rtc/v1/whep/?app=live&stream=473b95a47e338301cbd96809ea7ac416"

    # 运行拉流任务
    asyncio.run(whep_pull_video_stream(WHEP_URL))