# coding=utf-8
import os
import platform
import tkinter
from tkinter import *
from HCNetSDK import *
from PlayCtrl import *
from datetime import datetime
import json
from yolo_detector import YOLODetector
import cv2
import numpy as np
from PIL import Image, ImageTk
import redis
from PIL import ImageFont, ImageDraw, Image
import base64
from threading import Thread
from queue import Queue
import queue
import time
import threading
from yolo_processor import process_frame_with_yolo


class DISPLAY_INFO(Structure):
    _fields_ = [
        ("pBuf", POINTER(c_ubyte)),
        ("nSize", c_ulong),
        ("nWidth", c_long),
        ("nHeight", c_long),
        ("nStamp", c_ulong),
        ("nType", c_ulong),
        ("nReserved", c_ulong),
    ]


# 读取配置文件
def load_nvr_config(config_path='roi_config.json'):
    with open(config_path, 'r') as f:
        config = json.load(f)
        return config['nvr'], config['cameraCfgs']


# 获取NVR配置和摄像头配置
nvr_config, camera_configs = load_nvr_config()

# 登录的设备信息
DEV_IP = create_string_buffer(nvr_config['ip'].encode())
DEV_PORT = nvr_config['port']
DEV_USER_NAME = create_string_buffer(nvr_config['username'].encode())
DEV_PASSWORD = create_string_buffer(nvr_config['password'].encode())
YOLO_FPS = 2  # 每秒处理的帧数

WINDOWS_FLAG = True
win = None  # 预览窗口
funcRealDataCallBack_V30 = None  # 实时预览回调函数，需要定义为全局的

PlayCtrl_Port = c_long(-1)  # 播放句柄
Playctrldll = None  # 播放库
FuncDecCB = None  # 播放库解码回调函数，需要定义为全局的
max_channels = 9  # 最大支持9个通道
windows = [None] * max_channels  # 预览窗口数组
canvases = [None] * max_channels  # 画布数组
PlayCtrl_Ports = [c_long(-1)] * max_channels  # 播放句柄数组
funcRealDataCallBack_V30_array = [None] * max_channels  # 预览回调函数数组
funcDecCB_array = [None] * max_channels  # 预览回调函数数组
frame_queues = [Queue(maxsize=2) for _ in range(max_channels)]

# 线程控制参数
WORKER_THREADS = 4  # 根据CPU核心数调整（建议物理核心数×2）
EXIT_SIGNAL = "EXIT"
consumers = []  # 消费者线程池
# 初始化队列
frame_queue = queue.Queue(maxsize=200)  # 原始数据队列
result_queue = queue.Queue()            # GUI更新队列

processing_threads = [None] * max_channels
enable_yolo = True
last_process_times = [datetime.now()] * max_channels  # 每个通道上次处理的时间
redis_client = redis.Redis(host='192.168.110.229', port=36379, db=0, password="yunda123")


# 获取当前系统环境
def GetPlatform():
    sysstr = platform.system()
    print('' + sysstr)
    if sysstr != "Windows":
        global WINDOWS_FLAG
        WINDOWS_FLAG = False


# 设置SDK初始化依赖库路径
def SetSDKInitCfg():
    # 设置HCNetSDKCom组件库和SSL库加载路径
    # print(os.getcwd())
    if WINDOWS_FLAG:
        strPath = os.getcwd().encode('gbk')
        sdk_ComPath = NET_DVR_LOCAL_SDK_PATH()
        sdk_ComPath.sPath = strPath
        Objdll.NET_DVR_SetSDKInitCfg(2, byref(sdk_ComPath))
        Objdll.NET_DVR_SetSDKInitCfg(3, create_string_buffer(strPath + b'\\libcrypto-1_1-x64.dll'))
        Objdll.NET_DVR_SetSDKInitCfg(4, create_string_buffer(strPath + b'\\libssl-1_1-x64.dll'))
    else:
        strPath = os.getcwd().encode('utf-8')
        sdk_ComPath = NET_DVR_LOCAL_SDK_PATH()
        sdk_ComPath.sPath = strPath
        Objdll.NET_DVR_SetSDKInitCfg(2, byref(sdk_ComPath))
        Objdll.NET_DVR_SetSDKInitCfg(3, create_string_buffer(strPath + b'/libcrypto.so.1.1'))
        Objdll.NET_DVR_SetSDKInitCfg(4, create_string_buffer(strPath + b'/libssl.so.1.1'))


def LoginDev(Objdll):
    # 登录注册设备
    device_info = NET_DVR_DEVICEINFO_V30()
    lUserId = Objdll.NET_DVR_Login_V30(DEV_IP, DEV_PORT, DEV_USER_NAME, DEV_PASSWORD, byref(device_info))
    return (lUserId, device_info)


def start_processing_thread(channel_index):
    def run():
        while True:
            frame_data = frame_queues[channel_index].get()
            print("out")
            if frame_data is None:  # 终止信号
                break
            # 在此处执行YOLO处理和Redis发布
            # process_frame_with_yolo(frame_data, channel_index)
            process_frame_with_yolo(frame_data, channel_index, camera_configs, yolo_detector, redis_client)

    thread = Thread(target=run, daemon=True)
    thread.start()
    return thread


def on_closing():
    for i in range(num_channels):
        frame_queues[i].put(None)  # 发送终止信号
        processing_threads[i].join()
    root.destroy()


def DecCBFun(nPort, pBuf, nSize, pFrameInfo, nUser, nReserved2):
    # 解码回调函数
    channel_index = nUser - 1
    # print(f"回调函数被调用 - 通道: {channel_index + 1}, 数据类型: {dwDataType}, 数据大小: {dwBufSize}")

    if channel_index < 0 or channel_index >= max_channels:
        print(f"无效的通道索引: {channel_index}")
        return

    port = PlayCtrl_Ports[channel_index]
    canvas = canvases[channel_index]
    if pFrameInfo.contents.nType == 3:
        # 获取当前时间
        now = datetime.now()
        # 格式化为 HH:mm:ss:fff 格式
        formatted_time = now.strftime("%H:%M:%S:%f")[:-3]

        sFileName = ('../../pic/test_stamp[%s].jpg' % pFrameInfo.contents.nStamp)
        nWidth = pFrameInfo.contents.nWidth
        nHeight = pFrameInfo.contents.nHeight
        nType = pFrameInfo.contents.nType
        dwFrameNum = pFrameInfo.contents.dwFrameNum
        nStamp = pFrameInfo.contents.nStamp
        #print(nWidth, nHeight, nType, dwFrameNum, nStamp, sFileName)

        if yolo_detector is not None:
            try:
                if nWidth > 0 and nHeight > 0:
                    try:
                        # 提取YUV数据并转换为RGB
                        yuv_data = np.frombuffer(pBuf[:nSize], dtype=np.uint8)
                        height = pFrameInfo.contents.nHeight
                        width = pFrameInfo.contents.nWidth
                        # YUV420转RGB
                        yuv_frame = yuv_data.reshape((height * 3 // 2, width))
                        # rgb_frame = cv2.cvtColor(yuv_frame, cv2.COLOR_YUV2RGB_I420)
                        try:
                            # 仅传递原始数据到处理队列
                            if not frame_queue.full():
                                frame_queue.put((channel_index, yuv_frame), block=False)
                        except Exception as e:
                            print(f"回调异常: {str(e)}")
                        # cv2.imshow("frame", frame)
                        # cv2.waitKey(0)
                        # cv2.destroyAllWindows()
                    except Exception as e:
                        print(f"YUV转RGB失败: {str(e)}")

            except Exception as e:
                print(f"YOLO处理异常: {str(e)}")
                import traceback
                traceback.print_exc()


def save_detections(detections, frame, output_dir="."):
    os.makedirs(output_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_image = os.path.join(output_dir, f"detection_{timestamp}.jpg")
    output_json = os.path.join(output_dir, f"detection_{timestamp}.json")
    # 保存可视化图片
    cv2.imwrite(output_image, frame)

    # 保存结构化数据
    with open(output_json, 'w') as f:
        json.dump(detections, f, indent=2)
# 在全局变量部分修改
# 删除原有的win1, win2, cv1, cv2等变量
# 替换为通道数组
# 修改RealDataCallBack_V30函数，支持多通道
def RealDataCallBack_V30(lPlayHandle, dwDataType, pBuffer, dwBufSize, pUser):
    try:
        channel_index = pUser - 1
        # print(f"回调函数被调用 - 通道: {channel_index + 1}, 数据类型: {dwDataType}, 数据大小: {dwBufSize}")

        if channel_index < 0 or channel_index >= max_channels:
            print(f"无效的通道索引: {channel_index}")
            return

        port = PlayCtrl_Ports[channel_index]
        canvas = canvases[channel_index]

        if dwDataType == NET_DVR_SYSHEAD:
            print(f"通道 {channel_index + 1} 收到系统头数据")
            Playctrldll.PlayM4_SetStreamOpenMode(port, 0)
            if Playctrldll.PlayM4_OpenStream(port, pBuffer, dwBufSize, 1024 * 1024):
                canvas_id = canvas.winfo_id()
                print(f"通道 {pUser} 尝试使用画布ID: {canvas_id}")

                # 确保画布仍然有效
                if not canvas.winfo_exists():
                    print(f"通道 {channel_index + 1} 的画布已经不存在")
                    return
                # 设置解码回调，可以返回解码后YUV视频数据
                funcDecCB_array[channel_index] = DECCBFUNWIN(DecCBFun)
                Playctrldll.PlayM4_SetDecCallBackExMend(port, funcDecCB_array[channel_index], None, 0, pUser)

                Playctrldll.PlayM4_Stop(port)

                if Playctrldll.PlayM4_Play(port, canvas_id):
                    print(f'通道 {pUser} 播放库播放成功，使用画布ID: {canvas_id}')
                    Playctrldll.PlayM4_SetDisplayBuf(port, 15)
                    Playctrldll.PlayM4_SetOverlayMode(port, 0, 0)
                    # 禁用自动显示
                    Playctrldll.PlayM4_SetDisplayCallBack(port, None)
                else:
                    err = Playctrldll.PlayM4_GetLastError(port)
                    print(f'通道 {pUser} 播放库播放失败，错误码: {err}')
            else:
                err = Playctrldll.PlayM4_GetLastError(port)
                print(f'通道 {pUser} 播放库打开流失败，错误码: {err}')

        elif dwDataType == NET_DVR_STREAMDATA:
            try:
                # 检查端口是否有效
                if port.value < 0:
                    print(f"通道 {channel_index + 1} 的播放端口无效")
                    return

                # 直接输入数据，但不显示
                if not Playctrldll.PlayM4_InputData(port, pBuffer, dwBufSize):
                    err = Playctrldll.PlayM4_GetLastError(port)
                    # print(f"通道 {channel_index + 1} 输入数据失败，错误码: {err}")
                    # return

            except Exception as e:
                print(f"通道 {channel_index + 1} 数据处理异常: {str(e)}")
                import traceback
                traceback.print_exc()

    except Exception as e:
        print(f"回调函数发生未知异常: {str(e)}")
        import traceback
        traceback.print_exc()


# 添加一个计算布局的函数
def calculate_layout(num_channels):
    """根据通道数计算最佳布局"""
    if num_channels <= 1:
        return 1, 1
    elif num_channels <= 4:
        return 2, 2
    elif num_channels <= 6:
        return 2, 3
    else:  # 7-9通道
        return 3, 3


def OpenPreview(Objdll, lUserId, callbackFun, channel_index):
    '''
    打开预览
    '''
    preview_info = NET_DVR_PREVIEWINFO()
    preview_info.hPlayWnd = 0
    preview_info.lChannel = camera_configs[channel_index]['channel']  # 根据索引获取对应通道号
    preview_info.dwStreamType = 0
    preview_info.dwLinkMode = 0
    preview_info.bBlocked = 1

    lRealPlayHandle = Objdll.NET_DVR_RealPlay_V40(lUserId, byref(preview_info), callbackFun, channel_index + 1)
    return lRealPlayHandle


def InputData(fileMp4, Playctrldll):
    while True:
        pFileData = fileMp4.read(4096)
        if pFileData is None:
            break

        if not Playctrldll.PlayM4_InputData(PlayCtrl_Port, pFileData, len(pFileData)):
            break


# ------------------ 线程管理函数 ------------------
def start_consumers():
    """启动消费者线程池"""
    global consumers
    # 清理残留线程
    consumers = [t for t in consumers if t.is_alive()]

    # 创建新线程
    for _ in range(WORKER_THREADS):
        t = threading.Thread(
            target=consumer_worker,
            daemon=True  # 主退出时自动结束
        )
        t.start()
        consumers.append(t)
    print(f"已启动 {len(consumers)} 个消费者线程")


def stop_consumers():
    """安全停止所有消费者线程"""
    for _ in range(WORKER_THREADS):
        frame_queue.put((EXIT_SIGNAL, None))  # 发送退出信号

    # 等待线程结束
    for t in consumers:
        t.join(timeout=5)
        if t.is_alive():
            print(f"警告: 线程 {t.ident} 未正常退出")
    print("所有消费者线程已停止")


# 修改消费者工作线程
def consumer_worker():
    while True:
        try:
            channel, yuv_array = frame_queue.get(timeout=5)
            if channel is EXIT_SIGNAL:
                break
            # YUV转RGB（实际分辨率从配置读取）
            config = camera_configs[channel]
            rgb_frame = cv2.cvtColor(yuv_array, cv2.COLOR_YUV2RGB_I420)
            # # 执行推理
            # process_frame_with_yolo(rgb_frame, channel)
            process_frame_with_yolo(rgb_frame, channel, camera_configs, yolo_detector, redis_client)
            # 准备GUI更新数据
            resized_frame = cv2.resize(rgb_frame, (320, 240))
            img = cv2.cvtColor(resized_frame, cv2.COLOR_RGB2BGR)
        except queue.Empty:
            continue


# 增强型GUI更新函数
def update_gui():
    while not result_queue.empty():
        try:
            channel, img_tk, status = result_queue.get_nowait()

            # 安全更新GUI组件
            canvases[channel].configure(image=img_tk)
            canvases[channel].image = img_tk  # 保持引用
            canvases[channel].itemconfigure("status", text=status)

        except Exception as e:
            print(f"GUI更新异常: {str(e)}")

    root.after(50, update_gui)


# 在界面初始化时预加载画布ID
def create_canvas_grid():
    global canvas_ids
    canvas_ids = {}

    for i in range(num_channels):
        canvases[i].update()
        canvas_ids[i] = canvases[i].winfo_id()  # 预获取窗口ID
        print(f"通道 {i} 画布ID: {canvas_ids[i]}")


# 在设备初始化时设置播放参数
def set_play_params(port, channel):
    # 使用预存储的画布ID
    Playctrldll.PlayM4_SetDisplayWindow(
        port,
        canvas_ids[channel],  # 使用预获取的ID
        0, 0,
        single_width, single_height
    )


# 修改SDK初始化部分
def Initcustom():
    # 预加载所有画布ID
    create_canvas_grid()

    # 启动消费者线程
    start_consumers()

    # 启动GUI更新循环
    root.after(0, update_gui)

    # 启动主循环
    root.mainloop()
# 程序退出时清理
    stop_consumers()


if __name__ == '__main__':
    try:
        print("程序启动...")

        # YOLO初始化
        print("正在初始化YOLO检测器...")
        try:
            yolo_detector = YOLODetector() if enable_yolo else None
            print("YOLO检测器初始化成功")
        except Exception as e:
            print(f"YOLO检测器初始化失败: {str(e)}")
            raise

        # 创建窗口
        print("正在创建主窗口...")
        root = tkinter.Tk()
        root.title("多通道预览")

        # 加载配置
        print("正在加载配置文件...")
        num_channels = min(len(camera_configs), max_channels)
        print(f"检测到 {len(camera_configs)} 个摄像头配置，将显示 {num_channels} 个通道")

        # 计算布局
        rows, cols = calculate_layout(num_channels)

        # 设置单个视频窗口的大小
        single_width, single_height = 320, 240

        # 设置主窗口大小
        window_width = cols * single_width
        window_height = rows * single_height
        root.geometry(f"{window_width}x{window_height}+100+100")

        # 创建帧来容纳所有画布
        frame = tkinter.Frame(root, width=window_width, height=window_height)
        frame.pack()

        # 创建画布网格
        for i in range(num_channels):
            row = i // cols
            col = i % cols

            # 创建每个通道的画布，并确保它们有唯一的标识
            # 使用Frame作为容器，可能有助于解决显示问题
            frame_container = tkinter.Frame(frame, width=single_width, height=single_height)
            frame_container.grid(row=row, column=col, padx=2, pady=2)
            frame_container.grid_propagate(False)  # 防止frame自动调整大小

            canvases[i] = tkinter.Canvas(frame_container, bg='black',
                                         width=single_width,
                                         height=single_height,
                                         name=f"canvas_{i}")
            canvases[i].pack(fill=tkinter.BOTH, expand=True)

            # 添加通道标签
            canvases[i].create_text(10, 10, text=f"通道 {i + 1}", fill="white", anchor="nw")

            # 确保画布已经更新并有有效的窗口ID
            canvases[i].update()

        # 获取系统平台
        print("\n系统环境检测:")
        GetPlatform()
        print(f"当前系统平台标志: {'Windows' if WINDOWS_FLAG else 'Linux'}")

        # 加载库
        print("\n正在加载SDK库...")
        try:
            if WINDOWS_FLAG:
                print("当前工作目录:", os.getcwd())
                print("尝试切换到SDK目录: ./lib/win")
                os.chdir(r'./lib/win')
                print("加载 HCNetSDK.dll...")
                Objdll = ctypes.CDLL(r'./HCNetSDK.dll')
                print("加载 PlayCtrl.dll...")
                Playctrldll = ctypes.CDLL(r'./PlayCtrl.dll')
            else:
                # Linux加载代码保持不变...
                pass
            print("SDK库加载成功")
        except Exception as e:
            print(f"SDK库加载失败: {str(e)}")
            raise

        print("\n正在初始化SDK...")
        SetSDKInitCfg()

        # 初始化DLL
        if Objdll.NET_DVR_Init() == 0:
            print(f"SDK初始化失败，错误码: {Objdll.NET_DVR_GetLastError()}")
            raise Exception("SDK初始化失败")
        print("SDK初始化成功")

        # 登录设备
        print(f"\n正在登录设备 {DEV_IP.value.decode()}:{DEV_PORT}...")
        (lUserId, device_info) = LoginDev(Objdll)
        if lUserId < 0:
            err = Objdll.NET_DVR_GetLastError()
            print(f'设备登录失败，错误码: {err}')
            Objdll.NET_DVR_Cleanup()
            raise Exception("设备登录失败")
        print(f"设备登录成功，用户ID: {lUserId}")

        # 预览通道初始化
        print("\n开始初始化预览通道...")
        lRealPlayHandles = []

        for i in range(num_channels):
            PlayCtrl_Ports[i] = c_long(-1)
            if not Playctrldll.PlayM4_GetPort(byref(PlayCtrl_Ports[i])):
                error_code = Playctrldll.PlayM4_GetLastError(PlayCtrl_Ports[i])
                print(f'通道 {i + 1} 获取播放库句柄失败，错误码: {error_code}')
                PlayCtrl_Ports[i] = c_long(-1)  # 明确标记为无效
                continue
            print(f'通道 {i + 1} 获取播放库端口成功: {PlayCtrl_Ports[i].value}')

            # 创建回调函数
            print(f'正在为通道 {i + 1} 创建回调函数...')
            funcRealDataCallBack_V30_array[i] = REALDATACALLBACK(RealDataCallBack_V30)

            # 开启预览
            print(f'正在开启通道 {i + 1} 的预览...')
            lRealPlayHandle = OpenPreview(Objdll, lUserId, funcRealDataCallBack_V30_array[i], i)
            lRealPlayHandles.append(lRealPlayHandle)

            if lRealPlayHandle < 0:
                print(f'通道 {i + 1} 预览失败，错误码: {Objdll.NET_DVR_GetLastError()}')
            else:
                print(f'通道 {i + 1} 预览成功，句柄: {lRealPlayHandle}')

        Initcustom()
        root.mainloop()
    except Exception as e:
        print(f"\n程序发生异常: {str(e)}")
        print("错误详细信息:")
        import traceback

        traceback.print_exc()

    finally:
        print("\n开始清理资源...")
        # 关闭预览
        if 'lRealPlayHandles' in locals():
            for handle in lRealPlayHandles:
                if handle >= 0:
                    print(f"正在关闭预览句柄: {handle}")
                    Objdll.NET_DVR_StopRealPlay(handle)

        # 释放播放库资源
        if 'PlayCtrl_Ports' in locals():
            for i in range(num_channels):
                if PlayCtrl_Ports[i].value > -1:
                    print(f"正在释放通道 {i + 1} 的播放库资源")
                    Playctrldll.PlayM4_Stop(PlayCtrl_Ports[i])
                    Playctrldll.PlayM4_CloseStream(PlayCtrl_Ports[i])
                    Playctrldll.PlayM4_FreePort(PlayCtrl_Ports[i])

        # 登出设备
        if 'lUserId' in locals() and lUserId >= 0:
            print("正在登出设备...")
            Objdll.NET_DVR_Logout(lUserId)

        # 清理SDK资源
        if 'Objdll' in locals():
            print("正在清理SDK资源...")
            Objdll.NET_DVR_Cleanup()

        print("程序退出")