# -*- coding: utf-8 -*-
"""mask 区域重绘（inpainting）服务

两种模式（底层都是 webui img2img inpainting，严格只在mask区重绘）：
  1. enhance: 用 enhance_hair.webui_img2img（默认提示词=发丝增强，denoising默认0.35）
  2. pure_inpaint: 纯SD inpainting（默认提示词=自然头发，denoising默认0.5）

严格只重绘 mask 区域：inpainting_fill=1（mask外保留原图）
"""
import cv2
import numpy as np
from utils import enhance_hair


def inpaint_mask(img, mask, prompt="", denoising_strength=0.5, mode="enhance"):
    """在 mask 区域做 SD inpainting 重绘。

    参数:
        img: 输入图 BGR
        mask: 灰度mask，白色(255)=重绘区，与img同尺寸
        prompt: 提示词（空则用模式默认值）
        denoising_strength: 重绘强度 0~1
        mode: "enhance" 或 "pure_inpaint"

    返回:
        result: 重绘后的图（mask区已重绘，其余保留原图）
        info: 本次参数说明
    """
    h, w = img.shape[:2]

    # 归一化 mask（保证 0/255，尺寸一致）
    if mask.shape[:2] != (h, w):
        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
    if mask.ndim == 3:
        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    _, mask_bin = cv2.threshold(mask, 30, 255, cv2.THRESH_BINARY)

    # 模式默认值
    if not prompt:
        if mode == "enhance":
            prompt = "high quality, detailed natural hair strands, realistic hair texture, sharp focus, photorealistic"
        else:
            prompt = "natural realistic hair, soft hair strands, photorealistic, high quality, detailed"

    # 统计 mask 区域占比
    mask_ratio = float((mask_bin > 0).sum()) / mask_bin.size * 100

    # 调用 enhance_hair.webui_img2img（底层是 inpainting，inpainting_fill=1 严格只在mask区重绘）
    result = enhance_hair.webui_img2img(
        img, mask_bin, prompt=prompt,
        denoising_strength=denoising_strength
    )

    # webui 可能调整尺寸（调整到8的倍数），resize 回原图尺寸
    if result.shape[:2] != (h, w):
        result = cv2.resize(result, (w, h), interpolation=cv2.INTER_LANCZOS4)

    # 双重保险：强制 mask 外区域用原图（虽然 inpainting_fill=1 已保证，再加一层确保严格）
    mask_3c = (mask_bin > 0)[..., None]
    result = np.where(mask_3c, result, img).astype(np.uint8)

    info = (f"模式={mode}，提示词=\"{prompt[:50]}...\"，denoising={denoising_strength}，"
            f"mask区域={mask_ratio:.1f}%。严格只重绘mask区，其余保留原图。")
    return result, mask_bin, info