#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
队伍圆圈识别脚本
识别"王于兴师"客户端窗口右侧的绿色圆圈部队槽位
"""

import json
import cv2
import numpy as np
import pygetwindow as gw
import pyautogui
import os

# 调试开关
DEBUG = True

# 输出目录
OUTPUT_DIR = "debug_output"

# 颜色阈值（HSV）
# 绿色范围
LOWER_GREEN = np.array([35, 50, 50])
UPPER_GREEN = np.array([85, 255, 255])

# 轮廓筛选参数
MIN_AREA = 100  # 最小面积
MAX_AREA = 2000  # 最大面积
MIN_CIRCULARITY = 0.7  # 最小圆度

# 右侧区域比例（相对于窗口宽度）
RIGHT_REGION_RATIO = 0.7  # 只分析窗口右侧30%的区域


def ensure_output_dir():
    """确保输出目录存在"""
    if DEBUG and not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR)


def get_active_window():
    """获取当前激活窗口"""
    try:
        active_window = gw.getActiveWindow()
        if active_window and "王于兴师" in active_window.title:
            return active_window
        else:
            return None
    except Exception as e:
        print(f"获取窗口失败: {e}")
        return None


def capture_window(window):
    """截取窗口图像"""
    try:
        # 获取窗口位置和大小
        left, top, width, height = window.left, window.top, window.width, window.height
        # 截图
        screenshot = pyautogui.screenshot(region=(left, top, width, height))
        # 转换为OpenCV格式
        img = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
        
        if DEBUG:
            cv2.imwrite(os.path.join(OUTPUT_DIR, "window_screenshot.png"), img)
        
        return img
    except Exception as e:
        print(f"截图失败: {e}")
        return None


def process_image(img):
    """处理图像，识别绿色圆圈"""
    try:
        # 转换为HSV
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        
        # 颜色过滤
        mask = cv2.inRange(hsv, LOWER_GREEN, UPPER_GREEN)
        
        if DEBUG:
            cv2.imwrite(os.path.join(OUTPUT_DIR, "green_mask.png"), mask)
        
        # 形态学操作
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        
        # 查找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 筛选轮廓
        circles = []
        height, width = img.shape[:2]
        right_region_left = int(width * RIGHT_REGION_RATIO)
        
        for contour in contours:
            # 计算面积
            area = cv2.contourArea(contour)
            if area < MIN_AREA or area > MAX_AREA:
                continue
            
            # 计算圆度
            perimeter = cv2.arcLength(contour, True)
            if perimeter == 0:
                continue
            circularity = 4 * np.pi * area / (perimeter * perimeter)
            if circularity < MIN_CIRCULARITY:
                continue
            
            # 计算边界矩形
            x, y, w, h = cv2.boundingRect(contour)
            aspect_ratio = w / h
            if aspect_ratio < 0.8 or aspect_ratio > 1.2:
                continue
            
            # 只保留右侧区域的目标
            if x < right_region_left:
                continue
            
            # 计算中心坐标
            center_x = x + w // 2
            center_y = y + h // 2
            circles.append((center_x, center_y))
        
        # 去重（距离小于阈值的视为同一目标）
        circles = remove_duplicates(circles, threshold=20)
        
        # 按y坐标排序
        circles.sort(key=lambda x: x[1])
        
        # 绘制结果
        if DEBUG and circles:
            result_img = img.copy()
            for i, (x, y) in enumerate(circles, 1):
                cv2.circle(result_img, (x, y), 10, (0, 255, 0), 2)
                cv2.putText(result_img, str(i), (x-5, y+5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
            cv2.imwrite(os.path.join(OUTPUT_DIR, "detection_result.png"), result_img)
        
        return circles
    except Exception as e:
        print(f"图像处理失败: {e}")
        return []


def remove_duplicates(points, threshold=20):
    """去除重复的点"""
    if not points:
        return []
    
    unique_points = []
    for point in points:
        is_unique = True
        for unique_point in unique_points:
            distance = np.sqrt((point[0] - unique_point[0])**2 + (point[1] - unique_point[1])**2)
            if distance < threshold:
                is_unique = False
                break
        if is_unique:
            unique_points.append(point)
    return unique_points


def output_result(result):
    """输出结果到控制台和文件"""
    json_string = json.dumps(result, ensure_ascii=False, indent=2)
    
    # 写入文件
    with open("detect_result.json", "w", encoding="utf-8") as f:
        f.write(json_string)
    
    # 输出JSON
    print(json_string)


def main():
    """主函数"""
    ensure_output_dir()
    
    # 获取当前窗口
    window = get_active_window()
    if not window:
        result = {
            "success": False,
            "msg": "当前窗口不是王于兴师客户端"
        }
        output_result(result)
        return
    
    # 截图
    img = capture_window(window)
    if img is None:
        result = {
            "success": False,
            "msg": "截图失败"
        }
        output_result(result)
        return
    
    # 识别圆圈
    circles = process_image(img)
    if not circles:
        result = {
            "success": False,
            "msg": "未识别到队伍圆圈"
        }
        output_result(result)
        return
    
    # 构建结果
    teams = []
    for i, (x, y) in enumerate(circles, 1):
        teams.append({
            "index": i,
            "x": x,
            "y": y
        })
    
    result = {
        "success": True,
        "teams": teams
    }
    
    output_result(result)


if __name__ == "__main__":
    main()
