import os
import cv2
import numpy as np
from datetime import datetime
from ultralytics import YOLO


DEFAULT_YOLOV8S_SEG_MODEL = './models/yolov8s-seg.pt'

def get_center_distance(box, image_center):
    x1, y1, x2, y2 = box
    box_center = ((x1 + x2) / 2, (y1 + y2) / 2)
    return np.linalg.norm(np.array(box_center) - np.array(image_center))

def generate_mask_from_first_frame(video_path, filter_type:str = 'dist',mask_type:str="array",mask_file_path:str=None):    
    """
    获取视频第一帧的 人像 MASK:
    filter_type：
    dist  距离 画面中心最近的人;  default
    box : 画面中面积最大的人;
    mask_type :
    array: 直接返回 np.array   ; default
    file：保存成mask_file_path ，返回mask_file_path
    """
    if not os.path.exists(video_path):
        raise FileNotFoundError(f"视频文件不存在：{video_path}")
    
    cap = cv2.VideoCapture(video_path)
    ret, frame = cap.read()
    cap.release()

    if not ret:
        raise RuntimeError("无法读取视频第一帧")

    height, width, _ = frame.shape
    image_center = (width / 2, height / 2)


    print(os.path.exists(DEFAULT_YOLOV8S_SEG_MODEL))

    model = YOLO(DEFAULT_YOLOV8S_SEG_MODEL)  # 或 yolov8s-seg.pt，体积更大但更准

    # 3. 进行人形检测+分割
    results = model(frame)[0]

    # 筛选出所有类别为“person”的目标
    person_masks = []
    for i, cls in enumerate(results.boxes.cls):
        if int(cls) == 0:  # class 0 是 "person"
            box = results.boxes.xyxy[i].cpu().numpy()  # [x1, y1, x2, y2]
            mask = results.masks.data[i].cpu().numpy()

            dist = get_center_distance(box, image_center)
            box_area = (box[2] - box[0]) * (box[3] - box[1])

            person_masks.append((dist, box_area, mask))

    if not person_masks:
        return None

    # 4. 取距离中心最近的人形
    if filter_type=="box":
        person_masks.sort(key=lambda x: x[1], reverse=True)
    else:
        person_masks.sort(key=lambda x: x[0])
    
    best_mask = person_masks[0][2]  
    best_mask_resized = cv2.resize(best_mask, (width, height), interpolation=cv2.INTER_NEAREST)

    if mask_type == "file":
        if not mask_file_path:
            now_str = datetime.now().strftime("%Y%m%d_%H%M%S")
            mask_file_path =f"./temp/{now_str}.png"
        mask_img = (best_mask_resized * 255).astype(np.uint8)
        cv2.imwrite(mask_file_path, mask_img)
        return mask_file_path
    else:
        mask = (best_mask_resized * 255).astype(np.uint8)
        return mask