#!/usr/bin/env python3
"""Hair inpainting service - calls ComfyUI workflow with image + mask."""
import io
import json
import time
import random
import logging
import traceback
import requests
from flask import Flask, request, jsonify, send_file
import numpy as np
from PIL import Image, ImageFilter

# 让 PIL 支持 iPhone 的 HEIC/HEIF 照片（浏览器 accept="image/*" 会允许选中它们）。
try:
    from pillow_heif import register_heif_opener
    register_heif_opener()
    _HEIF_OK = True
except Exception:
    _HEIF_OK = False

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
log = logging.getLogger("hair")

app = Flask(__name__)
COMFYUI_URL = "http://127.0.0.1:8188"

# 允许浏览器跨域直连本服务（如 hair 项目的测试页）。
# 不引入 flask-cors 依赖，直接在响应头 + OPTIONS 预检里处理。
_CORS_HEADERS = {
    "Access-Control-Allow-Origin": "*",
    "Access-Control-Allow-Methods": "POST, OPTIONS",
    "Access-Control-Allow-Headers": "Content-Type",
    "Access-Control-Expose-Headers": "X-Generate-Time",
}


@app.after_request
def _add_cors(resp):
    for k, v in _CORS_HEADERS.items():
        resp.headers[k] = v
    return resp


@app.route("/api/generate", methods=["OPTIONS"])
def _generate_preflight():
    """CORS 预检：浏览器 POST 前会先发 OPTIONS。"""
    return ("", 204)


def build_workflow(image_filename, prompt_text, seed=None):
    """Build ComfyUI API workflow from the 0716add-hair.json structure."""
    if seed is None:
        seed = random.randint(0, 2**53)

    return {
        # Loaders
        "16": {"class_type": "UNETLoader", "inputs": {
            "unet_name": "flux2.0/flux-2-klein-9b-fp8.safetensors",
            "weight_dtype": "fp8_e4m3fn_fast"}},
        "3": {"class_type": "VAELoader", "inputs": {
            "vae_name": "flux2-vae.safetensors"}},
        "61": {"class_type": "CLIPLoader", "inputs": {
            "clip_name": "qwen_3_8b_fp8mixed.safetensors",
            "type": "flux2",
            "device": "default"}},

        # Input image (with mask in alpha channel)
        "26": {"class_type": "LoadImage", "inputs": {
            "image": image_filename}},

        # Prompt
        "60": {"class_type": "JjkText", "inputs": {
            "text": prompt_text}},
        "22": {"class_type": "CLIPTextEncode", "inputs": {
            "clip": ["61", 0],
            "text": ["60", 0]}},

        # Image size
        "31": {"class_type": "easy imageSize", "inputs": {
            "image": ["26", 0]}},

        # Mask processing: fill holes -> convert to image -> scale -> back to mask
        "33": {"class_type": "Mask Fill Holes", "inputs": {
            "masks": ["26", 1]}},
        "36": {"class_type": "Convert Masks to Images", "inputs": {
            "masks": ["33", 0]}},
        "39": {"class_type": "ImageScale", "inputs": {
            "image": ["36", 0],
            "upscale_method": "nearest-exact",
            "width": ["31", 0],
            "height": ["31", 1],
            "crop": "disabled"}},
        "37": {"class_type": "Image To Mask", "inputs": {
            "image": ["39", 0],
            "method": "intensity"}},

        # Scale image+mask by aspect ratio
        "32": {"class_type": "LayerUtility: ImageScaleByAspectRatio V2", "inputs": {
            "image": ["26", 0],
            "mask": ["37", 0],
            "aspect_ratio": "custom",
            "proportional_width": ["31", 0],
            "proportional_height": ["31", 1],
            "fit": "letterbox",
            "method": "lanczos",
            "round_to_multiple": "8",
            "scale_to_side": "None",
            "scale_to_length": 1024,
            "background_color": "#000000"}},

        # Preview (pass_through=true, just passes the image through)
        "44": {"class_type": "ImageAndMaskPreview", "inputs": {
            "image": ["32", 0],
            "mask": ["32", 1],
            "mask_opacity": 1,
            "mask_color": "FFFF00",
            "pass_through": True}},

        # Get size of scaled image
        "14": {"class_type": "GetImageSize+", "inputs": {
            "image": ["44", 0]}},

        # VAE encode the image
        "13": {"class_type": "VAEEncode", "inputs": {
            "pixels": ["44", 0],
            "vae": ["3", 0]}},

        # Flux model setup
        "2": {"class_type": "ModelSamplingFlux", "inputs": {
            "model": ["16", 0],
            "max_shift": 1.15,
            "base_shift": 0.5,
            "width": ["14", 0],
            "height": ["14", 1]}},
        "19": {"class_type": "FluxGuidance", "inputs": {
            "conditioning": ["22", 0],
            "guidance": 1}},
        "5": {"class_type": "ReferenceLatent", "inputs": {
            "conditioning": ["19", 0],
            "latent": ["13", 0]}},

        # Empty latent for sampling
        "7": {"class_type": "EmptySD3LatentImage", "inputs": {
            "width": ["14", 0],
            "height": ["14", 1],
            "batch_size": 1}},

        # Scheduler & guider
        "1": {"class_type": "BasicScheduler", "inputs": {
            "model": ["2", 0],
            "scheduler": "simple",
            "steps": 4,
            "denoise": 1}},
        "20": {"class_type": "BasicGuider", "inputs": {
            "model": ["2", 0],
            "conditioning": ["5", 0]}},

        # Noise & sampler
        "6": {"class_type": "RandomNoise", "inputs": {
            "noise_seed": seed}},
        "8": {"class_type": "KSamplerSelect", "inputs": {
            "sampler_name": "euler"}},
        "9": {"class_type": "SamplerCustomAdvanced", "inputs": {
            "noise": ["6", 0],
            "guider": ["20", 0],
            "sampler": ["8", 0],
            "sigmas": ["1", 0],
            "latent_image": ["7", 0]}},

        # VAE decode
        "10": {"class_type": "VAEDecode", "inputs": {
            "samples": ["9", 0],
            "vae": ["3", 0]}},

        # Color match with original image
        "62": {"class_type": "ColorMatch", "inputs": {
            "image_ref": ["26", 0],
            "image_target": ["10", 0],
            "method": "mkl",
            "strength": 1,
            "multithread": True}},

        # Save result
        "17": {"class_type": "SaveImage", "inputs": {
            "images": ["62", 0],
            "filename_prefix": "hair_inpaint"}},
    }


@app.route("/")
def index():
    return send_file("index.html")


@app.route("/api/generate", methods=["POST"])
def generate():
    t_start = time.time()
    try:
        if "image" not in request.files or "mask" not in request.files:
            msg = f"缺少上传文件: files={list(request.files.keys())}"
            log.warning(msg)
            return jsonify({"error": msg}), 400
        image_file = request.files["image"]
        mask_file = request.files["mask"]
        prompt_text = request.form.get("prompt", "填充遮罩区域的头发，皮肤加一点磨皮，再加一点美颜")
        log.info(
            "收到请求: image=%s mask=%s prompt=%r",
            image_file.filename, mask_file.filename, prompt_text,
        )

        # Load original image as RGB
        try:
            image = Image.open(image_file).convert("RGB")
        except Exception as e:
            log.error("无法解码人物图片 %s: %s", image_file.filename, e)
            hint = "" if _HEIF_OK else "（当前不支持 HEIC）"
            return jsonify({
                "error": f"无法识别人物图片格式{hint}，请改用 JPG/PNG: {e}"
            }), 400

        # Load mask and extract mask data from ALL channels (R, G, B, A)
        # This handles different mask formats:
        #   - Red mask (R=255 where drawn): user-provided PNG
        #   - White mask (R=G=B=255 where drawn): frontend canvas
        #   - Alpha mask (A=255 where drawn): transparent brush
        try:
            mask_img = Image.open(mask_file).convert("RGBA")
        except Exception as e:
            log.error("无法解码遮罩图片 %s: %s", mask_file.filename, e)
            return jsonify({
                "error": f"无法识别遮罩图片格式，请改用 JPG/PNG: {e}"
            }), 400
        mask_arr = np.array(mask_img)
        # Use max of all channels: 255 where any color/alpha is drawn, 0 where empty
        mask_data = np.max(mask_arr, axis=2)  # (H, W) uint8

        # Ensure mask matches image size
        mask_data_img = Image.fromarray(mask_data, mode="L")
        if mask_data_img.size != image.size:
            mask_data_img = mask_data_img.resize(image.size, Image.LANCZOS)

        # Apply slight blur for soft edges (similar to ComfyUI's brush)
        mask_data_img = mask_data_img.filter(ImageFilter.GaussianBlur(radius=4))

        # ComfyUI LoadImage: mask = 1.0 - (alpha/255)
        # So alpha=0 -> mask=1.0 (inpaint), alpha=255 -> mask=0.0 (keep)
        # We want: drawn area (mask_data=255) -> inpaint -> alpha=0
        #          undrawn area (mask_data=0) -> keep -> alpha=255
        comfyui_alpha = Image.eval(mask_data_img, lambda x: 255 - x)

        # Combine into RGBA (split RGB into separate channels first)
        r, g, b = image.split()
        rgba = Image.merge("RGBA", (r, g, b, comfyui_alpha))

        # Upload to ComfyUI
        img_bytes = io.BytesIO()
        rgba.save(img_bytes, format="PNG")
        img_bytes.seek(0)

        upload_resp = requests.post(
            f"{COMFYUI_URL}/upload/image",
            files={"image": ("hair_input.png", img_bytes, "image/png")},
            timeout=30,
        )
        upload_data = upload_resp.json()
        if "name" not in upload_data:
            log.error("ComfyUI 上传图片失败: %s", upload_data)
            return jsonify({"error": f"Upload failed: {upload_data}"}), 500
        image_filename = upload_data["name"]

        # Build and queue workflow
        workflow = build_workflow(image_filename, prompt_text)
        prompt_resp = requests.post(
            f"{COMFYUI_URL}/prompt",
            json={"prompt": workflow},
            timeout=30,
        )
        prompt_data = prompt_resp.json()
        if "error" in prompt_data:
            log.error(
                "ComfyUI /prompt 校验失败: error=%s node_errors=%s",
                prompt_data.get("error"), prompt_data.get("node_errors"),
            )
            return jsonify({"error": json.dumps(prompt_data["error"], ensure_ascii=False)}), 500
        prompt_id = prompt_data["prompt_id"]

        # Poll for completion (5 min timeout, 0.1s interval)
        for _ in range(3000):
            time.sleep(0.1)
            history_resp = requests.get(
                f"{COMFYUI_URL}/history/{prompt_id}", timeout=10
            )
            history_data = history_resp.json()
            if prompt_id in history_data:
                status = history_data[prompt_id].get("status", {})
                if status.get("status_str") == "error":
                    log.error(
                        "ComfyUI 工作流执行失败: %s",
                        json.dumps(status, ensure_ascii=False),
                    )
                    return jsonify({
                        "error": "Workflow execution failed",
                        "detail": status.get("messages", status),
                    }), 500
                outputs = history_data[prompt_id].get("outputs", {})
                if "17" in outputs:  # SaveImage node
                    image_info = outputs["17"]["images"][0]
                    filename = image_info["filename"]
                    subfolder = image_info.get("subfolder", "")
                    img_type = image_info.get("type", "output")
                    view_resp = requests.get(
                        f"{COMFYUI_URL}/view",
                        params={"filename": filename, "subfolder": subfolder, "type": img_type},
                        timeout=30,
                    )
                    elapsed = time.time() - t_start
                    log.info("重绘完成，服务端耗时 %.2fs", elapsed)
                    resp = send_file(
                        io.BytesIO(view_resp.content), mimetype="image/png"
                    )
                    resp.headers["X-Generate-Time"] = f"{elapsed:.2f}"
                    return resp

        return jsonify({"error": "Timeout: workflow did not complete in 5 minutes"}), 500

    except requests.ConnectionError:
        log.error("无法连接 ComfyUI @ %s", COMFYUI_URL)
        return jsonify({"error": "Cannot connect to ComfyUI at " + COMFYUI_URL + ". Is it running?"}), 503
    except Exception as e:
        log.error("生成失败，未捕获异常:\n%s", traceback.format_exc())
        return jsonify({"error": f"{type(e).__name__}: {e}"}), 500


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8899, debug=False)