"""接口10：头部外缘膨胀带遮罩。

先和接口9 一样得到**内缩后的基准遮罩**（含额头的闭合区域外缘朝151内缩 erode_cm、底线不动，默认1.2cm），
在这个基础上：
1. 取基准遮罩的**外轮廓线**（1px），去掉贴着底部分界线的那一段（只留头发/头部外缘弧线）。
2. 把这条外轮廓线膨胀成带子（半径 = dilate_cm/2，即带子**总宽 ≈ dilate_cm**，默认 2cm）。
3. 裁到分界线以上（不越过底线）。
输出这条带子作为 mask。BiSeNet / SegFormer 两套并排对比，分步可视化。

两个可调参数：erode_cm（同接口9 的内缩，默认1.2）+ dilate_cm（带子总宽，默认2）。
复用 `head_mask` 的构件，避免重复实现。
"""
import cv2
import numpy as np

from face_analysis.detector import detector
from face_analysis.calibration import estimate_scale_factor
from face_analysis.head_mask import (
    BASELINE_IDX, ERODE_CM, NoFaceError,
    _baseline_points, _upper_region_mask, _fill_to_baseline, _largest_cc,
    _bisenet_hair_mask, _segformer_hair_mask,
    _erode, _overlay, _draw_baseline, _b64png, _mask_png,
)

DILATE_CM = 2.0          # 膨胀后带子总宽（厘米，默认；半径 = 总宽/2；可由入参覆盖）


def _dilate(mask_bool, r):
    """圆盘核膨胀半径 r（像素）。r<=0 原样返回。"""
    if r <= 0:
        return mask_bool.copy()
    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * r + 1, 2 * r + 1))
    return cv2.dilate(mask_bool.astype(np.uint8), k).astype(bool)


def _baseline_line_mask(baseline_pts, w, h, thickness):
    """底部分割线（含左右水平延长线）画成一条带厚度的线，用于从外轮廓里剔除底边。"""
    m = np.zeros((h, w), np.uint8)
    y_l = baseline_pts[0][1]
    y_r = baseline_pts[-1][1]
    chain = [(0, y_l)] + baseline_pts + [(w - 1, y_r)]
    for a, b in zip(chain[:-1], chain[1:]):
        cv2.line(m, a, b, 1, thickness)
    return m.astype(bool)


def _outer_contour_no_bottom(region, baseline_band):
    """区域外轮廓（1px）去掉贴着底部分界线的那一段。"""
    contour = region & ~_erode(region, 1)
    return contour & ~baseline_band


def _model_result(image_bgr, hair_mask, upper, baseline_pts, baseline_band, r_erode, r_dilate, w):
    """单个分割模型的分步结果（内缩后基准遮罩 / 外轮廓线 / 膨胀带 / 纯遮罩）。"""
    top_fill = _fill_to_baseline(hair_mask, upper)                        # 含额头，延伸到图底
    base = _largest_cc(_erode(top_fill, r_erode) & upper)                # 接口9 内缩后的基准遮罩
    contour = _outer_contour_no_bottom(base, baseline_band)              # 外轮廓，去底线
    band = _largest_cc(_dilate(contour, r_dilate) & upper)              # 膨胀成带、裁到底线以上
    return {
        "base_pixels": int(base.sum()),
        "band_pixels": int(band.sum()),
        "base_mask_base64": _b64png(
            _draw_baseline(_overlay(image_bgr, base, (255, 150, 0)), baseline_pts, w)),
        "contour_base64": _b64png(
            _draw_baseline(_overlay(image_bgr, _dilate(contour, 2), (0, 255, 0)), baseline_pts, w)),
        "band_overlay_base64": _b64png(
            _draw_baseline(_overlay(image_bgr, band, (0, 0, 255)), baseline_pts, w)),
        "mask_base64": _mask_png(band),
    }


def generate_head_band(image_bgr, erode_cm=ERODE_CM, dilate_cm=DILATE_CM):
    """接口10 完整管线。返回可直接进 ok() 的 data dict。

    erode_cm：基准遮罩外缘朝151 内缩距离（厘米，同接口9），页面可调，默认 1.2cm。
    dilate_cm：外轮廓线膨胀后带子总宽（厘米），页面可调，默认 2cm。
    未检出人脸抛 NoFaceError。单个分割模型异常不影响另一个（记为 {"error": ...}）。
    """
    h, w = image_bgr.shape[:2]
    landmarks = detector.detect(image_bgr)
    if landmarks is None:
        raise NoFaceError()

    erode_cm = max(0.0, float(erode_cm))
    dilate_cm = max(0.0, float(dilate_cm))
    px_per_cm = estimate_scale_factor(landmarks, w, h)
    r_erode = int(round(erode_cm * px_per_cm))            # 内缩半径
    r_dilate = int(round((dilate_cm / 2.0) * px_per_cm))  # 膨胀半径 = 总宽/2
    baseline_pts = _baseline_points(landmarks, w, h)
    upper = _upper_region_mask(baseline_pts, w, h)
    # 剔除底边用的分界线带：几像素宽即可，独立于膨胀/内缩半径
    baseline_band = _dilate(_baseline_line_mask(baseline_pts, w, h, 5), 2)

    baseline_viz = _draw_baseline(image_bgr, baseline_pts, w)
    data = {
        "px_per_cm": round(float(px_per_cm), 4),
        "erode_cm": round(erode_cm, 2),
        "erode_px": r_erode,
        "dilate_cm": round(dilate_cm, 2),
        "dilate_radius_px": r_dilate,
        "image_size": {"width": w, "height": h},
        "baseline_landmarks": [
            {"index": idx, "x": p[0], "y": p[1]}
            for idx, p in zip(BASELINE_IDX, baseline_pts)
        ],
        "steps_common": {
            "landmarks_baseline_base64": _b64png(baseline_viz),
            "upper_region_base64": _b64png(_overlay(baseline_viz, upper, (0, 200, 0))),
        },
    }

    seg_fns = {
        "bisenet": lambda: _bisenet_hair_mask(image_bgr, landmarks, w, h),
        "segformer": lambda: _segformer_hair_mask(image_bgr),
    }
    for name, fn in seg_fns.items():
        try:
            hair_mask = fn()
            data[name] = _model_result(image_bgr, hair_mask, upper, baseline_pts,
                                       baseline_band, r_erode, r_dilate, w)
        except Exception as ex:  # noqa: BLE001  单模型失败不影响整体
            data[name] = {"error": f"{type(ex).__name__}: {ex}"}
    return data