#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
主界面窗口
包含视频显示区域、控制按钮和车牌号显示区域
"""

import sys
import os
from PyQt5.QtWidgets import (
    QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
    QLabel, QPushButton, QFrame, QTextEdit, QGroupBox,
    QCheckBox, QSpinBox, QSlider, QGridLayout
)
from PyQt5.QtCore import Qt, QTimer, pyqtSignal
from PyQt5.QtGui import QFont, QPixmap, QPalette, QImage

from .video_widget import VideoWidget
from utils.video_capture import VideoCapture
from models.yolo_detector import YOLODetector
from models.plate_recognizer import PlateRecognizerManager

class MainWindow(QMainWindow):
    """主窗口类"""
    
    def __init__(self):
        super().__init__()
        self.video_capture = None
        self.yolo_detector = None
        self.plate_recognizer = PlateRecognizerManager('mock')  # 车牌识别管理器
        self.timer = QTimer()
        self.use_camera = 1  # 1: 摄像头, 0: 视频文件
        self.detected_plates = []  # 存储切割后的车牌图像数组
        self.current_frame = None  # 存储当前帧用于车牌矫正
        
        self.init_ui()
        self.init_detector()
        self.init_video_capture()
        self.connect_signals()
        
    def init_ui(self):
        """初始化用户界面"""
        self.setWindowTitle("车牌检测系统 - YOLO11s")
        self.setGeometry(100, 100, 1200, 800)
        
        # 创建中央widget
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 主布局
        main_layout = QHBoxLayout(central_widget)
        
        # 左侧视频显示区域
        self.create_video_area(main_layout)
        
        # 右侧控制和信息显示区域
        self.create_control_area(main_layout)
        
        # 设置布局比例
        main_layout.setStretch(0, 3)  # 视频区域占3/4
        main_layout.setStretch(1, 1)  # 控制区域占1/4
        
    def create_video_area(self, parent_layout):
        """创建视频显示区域"""
        video_frame = QFrame()
        video_frame.setFrameStyle(QFrame.StyledPanel)
        video_layout = QVBoxLayout(video_frame)
        
        # 视频显示widget
        self.video_widget = VideoWidget()
        video_layout.addWidget(self.video_widget)
        
        parent_layout.addWidget(video_frame)
        
    def create_control_area(self, parent_layout):
        """创建控制和信息显示区域"""
        control_frame = QFrame()
        control_frame.setFrameStyle(QFrame.StyledPanel)
        control_frame.setMaximumWidth(300)
        control_layout = QVBoxLayout(control_frame)
        
        # 控制按钮组
        self.create_control_buttons(control_layout)
        
        # 检测信息显示
        self.create_detection_info(control_layout)
        
        # 车牌号显示区域
        self.create_plate_display(control_layout)
        
        # 系统状态显示
        self.create_status_display(control_layout)
        
        parent_layout.addWidget(control_frame)
        
    def create_control_buttons(self, parent_layout):
        """创建控制按钮"""
        button_group = QGroupBox("控制面板")
        button_layout = QVBoxLayout(button_group)
        
        # 开始/停止按钮
        self.start_btn = QPushButton("开始检测")
        self.start_btn.setMinimumHeight(40)
        self.start_btn.clicked.connect(self.toggle_detection)
        button_layout.addWidget(self.start_btn)
        
        # 视频源切换
        self.camera_checkbox = QCheckBox("使用摄像头")
        self.camera_checkbox.setChecked(True)
        self.camera_checkbox.stateChanged.connect(self.toggle_video_source)
        button_layout.addWidget(self.camera_checkbox)
        
        # 检测开关
        self.detection_checkbox = QCheckBox("启用检测")
        self.detection_checkbox.setChecked(True)
        button_layout.addWidget(self.detection_checkbox)
        
        parent_layout.addWidget(button_group)
        
    def create_detection_info(self, parent_layout):
        """创建检测信息显示"""
        info_group = QGroupBox("检测信息")
        info_layout = QVBoxLayout(info_group)
        
        # FPS显示
        self.fps_label = QLabel("FPS: 0")
        self.fps_label.setFont(QFont("Arial", 12, QFont.Bold))
        info_layout.addWidget(self.fps_label)
        
        # 检测数量
        self.detection_count_label = QLabel("检测到车牌: 0")
        info_layout.addWidget(self.detection_count_label)
        
        # 模型信息
        self.model_info_label = QLabel("模型: YOLO11s (ONNX)")
        info_layout.addWidget(self.model_info_label)
        
        parent_layout.addWidget(info_group)
        
    def create_plate_display(self, parent_layout):
        """创建车牌号显示区域"""
        plate_group = QGroupBox("车牌识别结果")
        plate_layout = QVBoxLayout(plate_group)
        
        # 当前识别的车牌号
        self.current_plate_label = QLabel("当前车牌: 未识别")
        self.current_plate_label.setFont(QFont("Arial", 14, QFont.Bold))
        self.current_plate_label.setStyleSheet("color: blue; padding: 10px; border: 1px solid gray;")
        plate_layout.addWidget(self.current_plate_label)
        
        # 矫正后的车牌图像显示
        self.plate_image_label = QLabel("矫正后车牌图像")
        self.plate_image_label.setAlignment(Qt.AlignCenter)
        self.plate_image_label.setMinimumHeight(100)
        self.plate_image_label.setMaximumHeight(150)
        self.plate_image_label.setStyleSheet("border: 1px solid gray; background-color: #f0f0f0;")
        plate_layout.addWidget(self.plate_image_label)
        
        # 历史车牌记录
        history_label = QLabel("历史记录:")
        plate_layout.addWidget(history_label)
        
        self.plate_history = QTextEdit()
        self.plate_history.setMaximumHeight(150)
        self.plate_history.setReadOnly(True)
        plate_layout.addWidget(self.plate_history)
        
        # 预留接口说明
        interface_label = QLabel("注: 车牌识别接口已预留，可接入OCR模型")
        interface_label.setStyleSheet("color: gray; font-size: 10px;")
        plate_layout.addWidget(interface_label)
        
        parent_layout.addWidget(plate_group)
        
    def create_status_display(self, parent_layout):
        """创建系统状态显示"""
        status_group = QGroupBox("系统状态")
        status_layout = QVBoxLayout(status_group)
        
        self.status_label = QLabel("状态: 就绪")
        status_layout.addWidget(self.status_label)
        
        self.gpu_status_label = QLabel("GPU: 检测中...")
        status_layout.addWidget(self.gpu_status_label)
        
        parent_layout.addWidget(status_group)
        
        # 添加弹性空间
        parent_layout.addStretch()
        
    def init_detector(self):
        """初始化YOLO检测器"""
        try:
            model_path = os.path.join(os.path.dirname(__file__), "..", "yolo11sth50.onnx")
            self.yolo_detector = YOLODetector(model_path)
            self.model_info_label.setText(f"模型: YOLO11s (ONNX) - GPU: {self.yolo_detector.use_gpu}")
            self.gpu_status_label.setText(f"GPU: {'启用' if self.yolo_detector.use_gpu else '禁用'}")
        except Exception as e:
            self.status_label.setText(f"模型加载失败: {str(e)}")
            
    def init_video_capture(self):
        """初始化视频捕获"""
        try:
            self.video_capture = VideoCapture()
            self.status_label.setText("视频捕获初始化成功")
        except Exception as e:
            self.status_label.setText(f"视频捕获初始化失败: {str(e)}")
            
    def connect_signals(self):
        """连接信号和槽"""
        self.timer.timeout.connect(self.update_frame)
        
    def toggle_detection(self):
        """切换检测状态"""
        if self.timer.isActive():
            self.stop_detection()
        else:
            self.start_detection()
            
    def start_detection(self):
        """开始检测"""
        if self.video_capture and self.video_capture.start_capture(self.use_camera):
            # 根据视频源类型设置定时器间隔
            video_fps = self.video_capture.get_video_fps()
            timer_interval = int(1000 / video_fps)  # 转换为毫秒
            self.timer.start(timer_interval)
            
            self.start_btn.setText("停止检测")
            source_type = "摄像头" if self.use_camera else f"视频文件({video_fps:.1f}FPS)"
            self.status_label.setText(f"检测中... - {source_type}")
        else:
            self.status_label.setText("启动失败")
            
    def stop_detection(self):
        """停止检测"""
        self.timer.stop()
        if self.video_capture:
            self.video_capture.stop_capture()
        self.start_btn.setText("开始检测")
        self.status_label.setText("已停止")
        
    def toggle_video_source(self, state):
        """切换视频源"""
        self.use_camera = 1 if state == Qt.Checked else 0
        if self.timer.isActive():
            self.stop_detection()
            self.start_detection()
            
    def update_frame(self):
        """更新帧"""
        if not self.video_capture:
            return
            
        frame, fps = self.video_capture.get_frame()
        if frame is None:
            return
            
        # 保存当前帧用于车牌矫正
        self.current_frame = frame.copy()
        
        # 更新FPS显示
        self.fps_label.setText(f"FPS: {fps:.1f}")
        
        # 进行检测
        if self.detection_checkbox.isChecked() and self.yolo_detector:
            detections = self.yolo_detector.detect(frame)
            frame = self.yolo_detector.draw_detections(frame, detections)
            
            # 切割车牌图像
            if detections:
                self.detected_plates = self.yolo_detector.crop_plates(frame, detections)
                
                # 统计不同类型车牌数量
                blue_count = sum(1 for d in detections if d['class_id'] == 0)
                green_count = sum(1 for d in detections if d['class_id'] == 1)
                total_count = len(detections)
                
                self.detection_count_label.setText(f"检测到车牌: {total_count} (蓝牌:{blue_count}, 绿牌:{green_count})")
                
                # 调用车牌识别接口（预留）
                self.recognize_plates(self.detected_plates, detections)
            else:
                self.detection_count_label.setText("检测到车牌: 0")
        
        # 显示帧
        self.video_widget.update_frame(frame)
        
    def recognize_plates(self, plate_images, detections):
        """车牌识别接口（预留）"""
        # 这里是预留的车牌识别接口
        # 可以接入OCR模型进行车牌号识别
        if plate_images and detections and self.current_frame is not None:
            # 获取最新检测到的车牌信息
            latest_detection = detections[-1]  # 取最后一个检测结果
            plate_type = "Blue Plate" if latest_detection['class_id'] == 0 else "Green Plate"
            confidence = latest_detection['confidence']
            
            # 处理蓝色车牌的矫正
            corrected_image = None
            if latest_detection['class_id'] == 0:  # 蓝色车牌
                try:
                    bbox = latest_detection['bbox']
                    corrected_image = self.plate_recognizer.preprocess_blue_plate(
                        plate_images[-1], self.current_frame, bbox
                    )
                    self._display_plate_image(corrected_image)
                except Exception as e:
                    print(f"蓝色车牌矫正失败: {e}")
                    self.plate_image_label.setText("蓝色车牌矫正失败")
            elif latest_detection['class_id'] == 1:  # 绿色车牌
                # 绿色车牌处理预留
                self.plate_image_label.setText("绿色车牌处理\n(待实现)")
            
            # 模拟识别结果
            plate_text = f"Mock {plate_type}-{len(plate_images)}"
            self.current_plate_label.setText(f"Current Plate: {plate_text} (Confidence: {confidence:.2f})")
            
            # 添加到历史记录
            import datetime
            timestamp = datetime.datetime.now().strftime("%H:%M:%S")
            self.plate_history.append(f"[{timestamp}] {plate_text} (Confidence: {confidence:.2f})")
    
    def _display_plate_image(self, image):
        """在界面上显示车牌图像"""
        try:
            # 将OpenCV图像转换为QPixmap
            if len(image.shape) == 3:
                height, width, channel = image.shape
                bytes_per_line = 3 * width
                q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
            else:
                height, width = image.shape
                bytes_per_line = width
                q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_Grayscale8)
            
            # 缩放图像以适应标签大小
            pixmap = QPixmap.fromImage(q_image)
            scaled_pixmap = pixmap.scaled(self.plate_image_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
            
            self.plate_image_label.setPixmap(scaled_pixmap)
        except Exception as e:
            print(f"显示车牌图像失败: {e}")
            self.plate_image_label.setText(f"图像显示失败: {str(e)}")
            
    def closeEvent(self, event):
        """窗口关闭事件"""
        self.stop_detection()
        event.accept()