#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""为测试图自动生成发际线/额头区域的 mask（模拟用户手画遮罩）。
用 RetinaFace + 1k关键点定位额头，生成该区域的白色 mask。
"""
import os
import sys
import cv2
import numpy as np

HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd"
sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1")

import torch
from utils.landmark_processor import get_max_rect


def gen_hairline_mask(img_path, mask_path):
    """用关键点定位额头，生成发际线区域 mask"""
    from models.detector import RetinaFaceDetector
    from utils.MomocvFaceAlignment1K import MomocvFaceAlignment1K

    img = cv2.imread(img_path)
    if img is None:
        print(f"❌ 读取失败: {img_path}")
        return False

    print(f"原图: {img_path} shape={img.shape}")
    gpu_id = 0

    # 1. RetinaFace 检测
    detector = RetinaFaceDetector(gpu_id=gpu_id)
    dets, landms = detector.forward(img)
    if len(dets) == 0:
        print("❌ 未检测到人脸")
        return False
    max_idx = get_max_rect(dets)
    det = dets[max_idx]
    print(f"检测到人脸: {det}")

    # 2. 1k 关键点
    aligner = MomocvFaceAlignment1K(gpu_id=gpu_id)
    landmarks_1k = aligner.stable_forward(img.copy(), [det])
    if landmarks_1k is None or len(landmarks_1k) == 0:
        print("❌ 关键点检测失败")
        return False
    pts = landmarks_1k[0]  # (1000, 2)
    print(f"1k关键点 shape: {pts.shape}")

    # 3. 构造发际线/额头 mask
    # 1k点的前312点是整脸外轮廓（含发际线）
    # 额头区：用眉骨以上的外轮廓点 + 发际线点构造凸包
    h, w = img.shape[:2]
    mask = np.zeros((h, w), dtype=np.uint8)

    # 取脸部上半部分轮廓点（发际线到太阳穴）构造额头区域
    # 索引 0-311 是脸部外轮廓，取上半部分（额头/太阳穴）
    forehead_pts = []
    # 发际线区域：取轮廓的上半段 + 稍微往下扩展
    for i in range(0, 312, 3):  # 隔点采样，减少密度
        pt = pts[i]
        # 只取上半脸（y < 图像中线偏上）
        if pt[1] < h * 0.55:
            forehead_pts.append(pt)
    # 加入额头中心区域一些点，确保覆盖额头
    forehead_pts = np.array(forehead_pts, dtype=np.int32)

    if len(forehead_pts) < 3:
        print("❌ 额头点不足")
        return False

    # 凸包填充
    hull = cv2.convexHull(forehead_pts)
    cv2.fillConvexPoly(mask, hull, 255)

    # 膨胀让区域稍大一点（模拟用户涂抹）
    mask = cv2.dilate(mask, np.ones((5, 5), np.uint8), iterations=2)

    # 统计 mask 区域
    nonzero = cv2.countNonZero(mask)
    print(f"mask 生成完成: 白色区域={nonzero}px ({nonzero/(h*w)*100:.1f}%)")
    cv2.imwrite(mask_path, mask)
    print(f"✅ mask 已保存: {mask_path}")
    return True


if __name__ == "__main__":
    img_path = sys.argv[1] if len(sys.argv) > 1 else \
        "/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
    mask_path = sys.argv[2] if len(sys.argv) > 2 else \
        "/home/xsl/change_hair/project/logs/hairgrow_test_mask.png"
    gen_hairline_mask(img_path, mask_path)