"""
Local Dimming 测试模块
功能：
- 生成不同百分比的白色窗口图片
- 通过 UCD 发送图片到显示器
- 自动采集 CA410 亮度数据
- 记录并导出测试结果
"""

import os
import sys
import time
import atexit
import shutil
import numpy as np
from PIL import Image, ImageDraw
import UniTAP


class LocalDimmingController:
    """Local Dimming 控制器 - 用于发送不同百分比窗口 Pattern"""

    def __init__(self, ucd_controller):
        """
        初始化 Local Dimming 控制器

        Args:
            ucd_controller: UCD323 控制器实例
        """
        self.ucd = ucd_controller

        # 兼容打包后的路径
        if getattr(sys, "frozen", False):
            base_dir = os.path.dirname(sys.executable)
        else:
            base_dir = os.getcwd()

        self.temp_dir = os.path.join(base_dir, "temp_local_dimming")

        # 创建临时目录
        if not os.path.exists(self.temp_dir):
            os.makedirs(self.temp_dir)
            print(f"[LD] 创建临时目录: {self.temp_dir}")

        self.cached_images = {}  # 缓存已生成的图片 {(分辨率, 百分比): 文件路径}

        # 注册退出时自动清理
        atexit.register(self.cleanup)

        print("[LD] Local Dimming 控制器已初始化")

    def send_window_pattern_with_resolution(self, percentage, width, height):
        """
        发送指定百分比和分辨率的白色窗口 Pattern

        Args:
            percentage: 窗口面积百分比 (1-100)
            width: 图像宽度
            height: 图像高度

        Returns:
            bool: 是否成功
        """
        try:
            # 检查设备连接状态
            if not self.ucd.status:
                print("[LD 错误] 设备未连接")
                return False

            print(f"\n[LD] 开始发送 {percentage}% 窗口 Pattern")
            print(f"[LD] 使用分辨率: {width}x{height}")

            # 获取 Pattern Generator 和 Audio Generator
            # 兼容 UCDController（仅 HDMI）和 UCD323Controller（多接口）
            if hasattr(self.ucd, 'current_interface'):
                # UCD323Controller（多接口支持）
                interface = self.ucd.current_interface
                if interface == "HDMI":
                    pg = self.ucd.role.hdtx.pg
                    ag = self.ucd.role.hdtx.ag
                elif interface == "Type-C" or interface == "DP":
                    pg = self.ucd.role.dptx.pg
                    ag = self.ucd.role.dptx.ag
                else:
                    print(f"[LD 错误] 不支持的接口类型: {interface}")
                    return False
            else:
                # UCDController（仅 HDMI）
                pg = self.ucd.role.hdtx.pg
                ag = self.ucd.role.hdtx.ag

            # 先停止音频，避免蜂鸣声
            try:
                ag.stop_generate()
                print("[LD] 已停止音频生成")
            except Exception as e:
                print(f"[LD 警告] 停止音频失败: {e}")

            # 检查缓存
            cache_key = (f"{width}x{height}", percentage)

            if cache_key in self.cached_images:
                image_path = self.cached_images[cache_key]
                if os.path.exists(image_path):
                    print(f"[LD] 使用缓存图片: {image_path}")
                else:
                    print(f"[LD] 缓存图片不存在，重新生成...")
                    image_path = self._generate_and_save_image(
                        width, height, percentage, cache_key
                    )
            else:
                print(f"[LD] 正在生成 {percentage}% 窗口图像...")
                image_path = self._generate_and_save_image(
                    width, height, percentage, cache_key
                )

            # 发送图像到设备
            print(f"[LD] 正在发送图像到设备...")

            # 设置 ColorInfo
            color_mode = UniTAP.ColorInfo()
            color_mode.color_format = UniTAP.ColorInfo.ColorFormat.CF_RGB
            color_mode.bpc = 8
            color_mode.colorimetry = UniTAP.ColorInfo.Colorimetry.CM_sRGB

            # 获取当前 timing
            try:
                current_vm = pg.get_vm()
                timing = (
                    current_vm.timing
                    if current_vm and hasattr(current_vm, "timing")
                    else None
                )
            except:
                timing = None

            # 如果有 timing，设置 VideoMode
            if timing:
                video_mode = UniTAP.VideoMode(timing=timing, color_info=color_mode)
                pg.set_vm(vm=video_mode)

            # 设置图片 Pattern
            pg.set_pattern(pattern=image_path)

            # 应用
            pg.apply()

            print(f"[LD] {percentage}% 窗口 Pattern 已发送到设备")
            return True

        except Exception as e:
            print(f"[LD 异常] 发送 {percentage}% 窗口失败: {e}")
            import traceback

            traceback.print_exc()
            return False

    def send_window_pattern(self, percentage):
        """
        发送指定百分比的白色窗口 Pattern（从 GUI 获取分辨率）

        Args:
            percentage: 窗口面积百分比 (1-100)

        Returns:
            bool: 是否成功
        """
        # 从设备当前 timing 获取分辨率
        width, height = self.get_current_resolution()
        return self.send_window_pattern_with_resolution(percentage, width, height)

    def get_current_resolution(self):
        """
        从设备当前 timing 获取显示器分辨率

        Returns:
            tuple: (width, height)
        """
        try:
            # 方式1：从 Pattern Generator 的当前 VideoMode 获取
            if hasattr(self.ucd, 'current_interface'):
                interface = self.ucd.current_interface
                if interface == "HDMI":
                    pg = self.ucd.role.hdtx.pg
                elif interface == "Type-C" or interface == "DP":
                    pg = self.ucd.role.dptx.pg
                else:
                    pg = None
            else:
                pg = self.ucd.role.hdtx.pg

            if pg:
                current_vm = pg.get_vm()
                if current_vm and hasattr(current_vm, "timing") and current_vm.timing:
                    timing = current_vm.timing
                    if hasattr(timing, "h_active") and hasattr(timing, "v_active"):
                        width = timing.h_active
                        height = timing.v_active
                        print(f"[LD] 从当前 timing 获取分辨率: {width}x{height}")
                        return width, height

            # 方式2：从 current_timing 属性获取
            if hasattr(self.ucd, "current_timing") and self.ucd.current_timing:
                timing = self.ucd.current_timing
                if hasattr(timing, "h_active") and hasattr(timing, "v_active"):
                    width = timing.h_active
                    height = timing.v_active
                    print(f"[LD] 从 current_timing 获取分辨率: {width}x{height}")
                    return width, height

        except Exception as e:
            print(f"[LD 警告] 获取分辨率失败: {e}")

        print("[LD 警告] 使用默认分辨率 3840x2160")
        return 3840, 2160

    def _generate_and_save_image(self, width, height, percentage, cache_key):
        """
        生成并保存窗口图像

        Args:
            width: 图像宽度
            height: 图像高度
            percentage: 窗口面积百分比
            cache_key: 缓存键

        Returns:
            str: 图像文件路径
        """
        # 生成图像
        image_array = self._create_window_image(width, height, percentage)

        # 保存到项目目录
        filename = f"window_{width}x{height}_{percentage:03d}percent.png"
        image_path = os.path.join(self.temp_dir, filename)

        image = Image.fromarray(image_array, mode="RGB")
        image.save(image_path, format="PNG")

        # 缓存
        self.cached_images[cache_key] = image_path

        print(f"[LD] 图像已保存: {image_path}")
        return image_path

    def _create_window_image(self, width, height, percentage):
        """
        创建窗口图像
        黑色背景 + 居中白色矩形窗口（保持屏幕比例）

        Args:
            width: 图像宽度
            height: 图像高度
            percentage: 窗口面积百分比 (1-100)

        Returns:
            numpy.ndarray: RGB 图像数组 (height, width, 3)
        """
        # 创建黑色背景
        image = np.zeros((height, width, 3), dtype=np.uint8)

        # 计算窗口尺寸（保持屏幕比例）
        scale_factor = (percentage / 100.0) ** 0.5
        window_width = int(width * scale_factor)
        window_height = int(height * scale_factor)

        # 100% 时强制全屏
        if percentage == 100:
            window_width = width
            window_height = height

        # 计算居中位置
        x1 = (width - window_width) // 2
        y1 = (height - window_height) // 2
        x2 = x1 + window_width
        y2 = y1 + window_height

        # 绘制白色窗口
        image[y1:y2, x1:x2] = [255, 255, 255]

        print(
            f"[LD] 图像生成完成: {width}x{height}, 窗口 {window_width}x{window_height}"
        )
        return image

    def cleanup(self):
        """清理临时文件夹"""
        if os.path.exists(self.temp_dir):
            try:
                shutil.rmtree(self.temp_dir)
                print(f"[LD] 临时文件夹已删除: {self.temp_dir}")
            except Exception as e:
                print(f"[LD 警告] 删除临时文件夹失败: {e}")

    def __del__(self):
        """析构函数：清理临时文件（备用机制）"""
        try:
            self.cleanup()
        except:
            pass


class LocalDimmingTest:
    def __init__(self, ucd_controller, ca_serial, log_callback=None):
        """
        初始化 Local Dimming 测试

        Args:
            ucd_controller: UCD323 控制器实例
            ca_serial: CA410 串口实例
            log_callback: 日志回调函数
        """
        self.ucd = ucd_controller
        self.ca = ca_serial
        self.log = log_callback if log_callback else print

        # 临时图片目录
        self.temp_dir = self._init_temp_dir()

        # 测试结果
        self.test_results = []

        # 测试配置
        self.window_percentages = [1, 2, 5, 10, 18, 25, 50, 75, 100]
        self.wait_time = 2.0  # 每次切换后等待时间（秒）

        # 停止标志
        self.stop_flag = False

        self.log("✓ Local Dimming 测试模块已初始化")

    def _init_temp_dir(self):
        """初始化临时目录"""
        if getattr(sys, "frozen", False):
            base_dir = os.path.dirname(sys.executable)
        else:
            base_dir = os.getcwd()

        temp_dir = os.path.join(base_dir, "temp_local_dimming")
        os.makedirs(temp_dir, exist_ok=True)
        return temp_dir

    def generate_window_image(self, width, height, percentage):
        """
        生成窗口图片（黑色背景 + 居中白色矩形窗口）

        Args:
            width: 图像宽度
            height: 图像高度
            percentage: 窗口面积百分比 (1-100)

        Returns:
            str: 图片文件路径
        """
        # 计算窗口尺寸（保持屏幕比例）
        scale_factor = (percentage / 100.0) ** 0.5
        window_width = int(width * scale_factor)
        window_height = int(height * scale_factor)

        # 100% 时强制全屏
        if percentage == 100:
            window_width = width
            window_height = height

        # 创建黑色背景
        image = np.zeros((height, width, 3), dtype=np.uint8)

        # 计算居中位置
        x1 = (width - window_width) // 2
        y1 = (height - window_height) // 2
        x2 = x1 + window_width
        y2 = y1 + window_height

        # 绘制白色窗口
        image[y1:y2, x1:x2] = [255, 255, 255]

        # 保存图片
        filename = f"window_{width}x{height}_{percentage:03d}percent.png"
        image_path = os.path.join(self.temp_dir, filename)

        pil_image = Image.fromarray(image, mode="RGB")
        pil_image.save(image_path, format="PNG")

        self.log(f"  ✓ 图片已生成: {window_width}×{window_height} px")
        return image_path

    def send_image_to_ucd(self, image_path):
        """
        通过 UCD 发送图片到显示器

        Args:
            image_path: 图片文件路径

        Returns:
            bool: 是否成功
        """
        try:
            # 获取 Pattern Generator 和 Audio Generator
            # 兼容 UCDController（仅 HDMI）和 UCD323Controller（多接口）
            if hasattr(self.ucd, 'current_interface'):
                interface = self.ucd.current_interface
                if interface == "HDMI":
                    pg = self.ucd.role.hdtx.pg
                    ag = self.ucd.role.hdtx.ag
                elif interface == "Type-C" or interface == "DP":
                    pg = self.ucd.role.dptx.pg
                    ag = self.ucd.role.dptx.ag
                else:
                    self.log(f"  ❌ 不支持的接口类型: {interface}")
                    return False
            else:
                # UCDController（仅 HDMI）
                pg = self.ucd.role.hdtx.pg
                ag = self.ucd.role.hdtx.ag

            # 停止音频
            try:
                ag.stop_generate()
            except:
                pass

            # 设置 ColorInfo
            color_mode = UniTAP.ColorInfo()
            color_mode.color_format = UniTAP.ColorInfo.ColorFormat.CF_RGB
            color_mode.bpc = 8
            color_mode.colorimetry = UniTAP.ColorInfo.Colorimetry.CM_sRGB

            # 获取当前 timing
            try:
                current_vm = pg.get_vm()
                timing = (
                    current_vm.timing
                    if current_vm and hasattr(current_vm, "timing")
                    else None
                )
            except:
                timing = None

            # 设置 VideoMode
            if timing:
                video_mode = UniTAP.VideoMode(timing=timing, color_info=color_mode)
                pg.set_vm(vm=video_mode)

            # 设置图片 Pattern
            pg.set_pattern(pattern=image_path)

            # 应用
            pg.apply()

            return True

        except Exception as e:
            self.log(f"  ❌ 发送图片失败: {str(e)}")
            import traceback

            traceback.print_exc()
            return False

    def measure_luminance(self):
        """
        使用 CA410 采集亮度

        Returns:
            tuple: (x, y, lv, X, Y, Z) 或 None
        """
        try:
            if not self.ca:
                self.log("  ❌ CA410 未连接")
                return None

            # 采集数据
            x, y, lv, X, Y, Z = self.ca.readAllDisplay()

            if x is not None and y is not None and lv is not None:
                self.log(f"  ✓ 采集亮度: {lv:.2f} cd/m²")
                return (x, y, lv, X, Y, Z)
            else:
                self.log("  ❌ 采集数据失败")
                return None

        except Exception as e:
            self.log(f"  ❌ 采集亮度异常: {str(e)}")
            return None

    def run_test(self, resolution="3840x2160"):
        """
        执行完整的 Local Dimming 测试

        Args:
            resolution: 分辨率字符串，如 "3840x2160"

        Returns:
            list: 测试结果 [(百分比, x, y, lv, X, Y, Z), ...]
        """
        self.log("=" * 60)
        self.log("开始 Local Dimming 测试")
        self.log("=" * 60)

        # 重置停止标志
        self.stop_flag = False

        # 解析分辨率
        try:
            width, height = map(int, resolution.split("x"))
        except:
            width, height = 3840, 2160
            self.log(f"  ⚠️ 分辨率解析失败，使用默认值: {width}x{height}")

        self.log(f"  分辨率: {width}x{height}")
        self.log(f"  测试窗口: {self.window_percentages}")
        self.log(f"  等待时间: {self.wait_time} 秒")
        self.log("")

        self.test_results = []

        for i, percentage in enumerate(self.window_percentages, start=1):
            # 检查停止标志
            if self.stop_flag:
                self.log("⚠️ 测试已停止")
                break

            self.log(f"[{i}/{len(self.window_percentages)}] 测试 {percentage}% 窗口...")

            # 1. 生成图片
            image_path = self.generate_window_image(width, height, percentage)

            # 2. 发送到 UCD
            if not self.send_image_to_ucd(image_path):
                self.log(f"  ❌ {percentage}% 窗口发送失败，跳过")
                continue

            # 3. 等待稳定
            self.log(f"  ⏳ 等待 {self.wait_time} 秒...")
            time.sleep(self.wait_time)

            # 4. 采集亮度
            result = self.measure_luminance()

            if result:
                x, y, lv, X, Y, Z = result
                self.test_results.append((percentage, x, y, lv, X, Y, Z))
                self.log(f"  ✅ {percentage}% 窗口测试完成")
            else:
                self.log(f"  ❌ {percentage}% 窗口采集失败")

            self.log("")

        self.log("=" * 60)
        self.log("✅ Local Dimming 测试完成")
        self.log(
            f"  成功测试: {len(self.test_results)}/{len(self.window_percentages)} 个窗口"
        )
        self.log("=" * 60)

        return self.test_results

    def stop(self):
        """停止测试"""
        self.stop_flag = True
        self.log("⚠️ 正在停止测试...")

    def get_results_summary(self):
        """获取测试结果摘要"""
        if not self.test_results:
            return None

        luminances = [lv for _, _, _, lv, _, _, _ in self.test_results]

        return {
            "data_points": self.test_results,
            "max_luminance": max(luminances),
            "min_luminance": min(luminances),
            "avg_luminance": sum(luminances) / len(luminances),
        }

    def cleanup(self):
        """清理临时文件"""
        try:
            import shutil

            if os.path.exists(self.temp_dir):
                shutil.rmtree(self.temp_dir)
                self.log(f"✓ 临时文件已清理: {self.temp_dir}")
        except Exception as e:
            self.log(f"⚠️ 清理临时文件失败: {e}")