# -*- coding: utf-8 -*-
"""区域生发服务（走换发型工作流，端口 8888）

工作流与换发型一致：infer_hairstyle_diy_jy → warpAffine → photo_service+LoRA → webui → 贴回。
mask = origin_matting ∪ new_matting ∪ 手绘mask（不减刘海）。

启动: python app.py
页面: http://<host>:8888
"""
import os
import sys
import time
import base64
import traceback

# 把 hair_service_sd 加入 path，使其模块可被导入
HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
sys.path.insert(0, HAIR_SERVICE_DIR)
# 切换到 hair_service_sd 目录，使 common.logger 能读到 config/configure.ini（相对路径）
os.chdir(HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("CRYPTOGRAPHY_OPENSSL_NO_LEGACY", "1")

import cv2
import numpy as np
from flask import Flask, request, jsonify, send_from_directory, send_file
from gevent import pywsgi

# PyTorch 2.6+ 默认 weights_only=True，无法加载含 numpy 对象的旧权重（yolov5l.pt 等）。
# 统一改回 False（权重均为本机自有可信文件）。
import torch
_orig_torch_load = torch.load
def _torch_load(*args, **kwargs):
    kwargs.setdefault("weights_only", False)
    return _orig_torch_load(*args, **kwargs)
torch.load = _torch_load

PORT = 8899
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

app = Flask(__name__, static_folder="static", static_url_path="/static")

# ===== 模型懒加载（首次请求时初始化，避免启动卡死）=====
_hairstyle_process = None
_landmark_processor = None


def _get_models():
    """懒加载换发型全套模型"""
    global _hairstyle_process, _landmark_processor
    if _hairstyle_process is None:
        print("[init] 加载换发型模型（首次，约30-90秒）...")
        t0 = time.time()
        import torch
        from core.hairstyle_model import HairStyle_Model
        from utils import landmark_processor
        _hairstyle_process = HairStyle_Model(gpu=True, use_enhance=True)
        _landmark_processor = landmark_processor
        print(f"[init] 模型加载完成，耗时 {time.time()-t0:.1f}s")
    return _hairstyle_process, _landmark_processor


@app.route("/")
def index():
    return send_from_directory(os.path.join(BASE_DIR, "static"), "index.html")


@app.route("/swap")
def swap_page():
    """换发型测试页（独立，无遮罩）"""
    return send_from_directory(os.path.join(BASE_DIR, "static"), "swap.html")


@app.route("/test_new")
def test_new_page():
    """新发型效果测试页（聚焦本次训练的30款发型）"""
    return send_from_directory(os.path.join(BASE_DIR, "static"), "test_new.html")


@app.route("/debug")
def debug_page():
    """换发型调试页（全参数可视化）"""
    return send_from_directory(os.path.join(BASE_DIR, "static"), "debug.html")


@app.route("/hairline")
def hairline_page():
    """发际线带重绘实验页"""
    return send_from_directory(os.path.join(BASE_DIR, "static"), "hairline.html")


@app.route("/manual")
def manual_page():
    """手绘mask重绘测试页"""
    return send_from_directory(os.path.join(BASE_DIR, "static"), "manual.html")


@app.route("/api/swap", methods=["POST"])
def api_swap_proxy():
    """代理换发型请求到 8801（避免前端跨域问题）"""
    import requests as req
    try:
        data = request.json
        resp = req.post("http://127.0.0.1:8801/api/swapHair/v1",
                        json=data, timeout=600)
        return jsonify(resp.json())
    except Exception as e:
        return jsonify({"state": -1, "msg": f"换发型代理失败: {e}"}), 500


@app.route("/api/swap_viz", methods=["POST"])
def api_swap_viz():
    """换发型（带可视化中间产物）

    入参 JSON:
        img: 原图 base64
        hair_id: 发型ID
        is_hr: "true"/"false"
    返回:
        {state, result: 最终图b64, steps: [{title, desc, images:[{label, b64}]}]}
    """
    t0 = time.time()
    try:
        data = request.json
        img_b64 = data.get("img", "")
        hair_id = data.get("hair_id", "")
        is_hr = str(data.get("is_hr", "false")).lower() == "true"
        strict_mask = data.get("strict_mask", False)

        if not img_b64 or not hair_id:
            return jsonify({"state": -1, "msg": "img 和 hair_id 不能为空"}), 400

        img = _b64_to_ndarray(img_b64, color=True)
        if img is None:
            return jsonify({"state": -1, "msg": "img 解析失败"}), 400

        print(f"[swap_viz] img={img.shape}, hair_id={hair_id}, is_hr={is_hr}, strict_mask={strict_mask}")
        hairstyle_process, landmark_processor = _get_models()

        from hair_swap_viz import hair_swap_viz
        task_id = f"swapviz_{int(time.time()*1000)}"
        result, steps = hair_swap_viz(
            origin_img=img, hair_id=hair_id,
            hairstyle_process=hairstyle_process,
            landmark_processor=landmark_processor,
            task_id=task_id, is_hr=is_hr, strict_mask=strict_mask)

        _, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
        result_b64 = base64.b64encode(buf).decode("utf-8")
        print(f"[swap_viz] 完成，总耗时 {time.time()-t0:.1f}s，{len(steps)}个步骤")
        return jsonify({"state": 0, "result": result_b64, "steps": steps})

    except Exception as e:
        print(f"[swap_viz] 失败: {e}")
        traceback.print_exc()
        return jsonify({"state": -1, "msg": f"换发型失败: {e}"}), 500


@app.route("/api/swap_debug", methods=["POST"])
def api_swap_debug():
    """换发型调试接口（全参数可调 + 每步可视化）。

    入参 JSON:
        img: 原图 base64
        hair_id: 发型ID
        # 流程开关
        cut_bang: bool, 是否减刘海圆（默认 true）
        strict_mask: bool, 严格按mask贴回（默认 false）
        seamless_blend: bool, 泊松融合消除接缝（默认 true，仅 strict_mask 时生效）
        # 尺寸/对齐
        is_hr: bool, 高清模式（默认 true）
        dilate_kernel: [x,y], mask膨胀核（默认 [6,18]）
        # SD 推理（仅 denoising 可调）
        denoising_strength: float 0~1, 重绘强度（默认 0.6）
        # 贴回/融合
        blend_dilate: [x,y], strict贴回mask膨胀（默认 [5,5]）
        seamless_dilate: [x,y], 泊松融合mask膨胀（默认 [9,9]）
    返回:
        {state, result, steps, params}
    """
    t0 = time.time()
    try:
        d = request.json
        img_b64 = d.get("img", "")
        hair_id = d.get("hair_id", "")
        if not img_b64 or not hair_id:
            return jsonify({"state": -1, "msg": "img 和 hair_id 不能为空"}), 400

        img = _b64_to_ndarray(img_b64, color=True)
        if img is None:
            return jsonify({"state": -1, "msg": "img 解析失败"}), 400

        # 解析参数（带默认值 + 类型转换）
        def g(k, default): return d.get(k, default)
        dk = g("dilate_kernel", [6, 18])
        bd = g("blend_dilate", [5, 5])
        sd = g("seamless_dilate", [9, 9])

        print(f"[swap_debug] hair_id={hair_id}, cut_bang={g('cut_bang',True)}, "
              f"strict_mask={g('strict_mask',False)}, seamless_blend={g('seamless_blend',True)}, "
              f"is_hr={g('is_hr',True)}, denoising={g('denoising_strength',0.6)}")

        hairstyle_process, landmark_processor = _get_models()
        from hair_swap_debug import hair_swap_debug
        task_id = f"swapdbg_{int(time.time()*1000)}"
        result, steps, params = hair_swap_debug(
            origin_img=img, hair_id=hair_id,
            hairstyle_process=hairstyle_process,
            landmark_processor=landmark_processor,
            task_id=task_id,
            cut_bang=bool(g("cut_bang", True)),
            strict_mask=bool(g("strict_mask", False)),
            seamless_blend=bool(g("seamless_blend", True)),
            is_hr=bool(g("is_hr", True)),
            dilate_kernel=(int(dk[0]), int(dk[1])),
            denoising_strength=float(g("denoising_strength", 0.6)),
            blend_dilate=(int(bd[0]), int(bd[1])),
            seamless_dilate=(int(sd[0]), int(sd[1])),
        )

        _, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
        result_b64 = base64.b64encode(buf).decode("utf-8")
        print(f"[swap_debug] 完成，总耗时 {time.time()-t0:.1f}s")
        return jsonify({"state": 0, "result": result_b64, "steps": steps, "params": params})

    except Exception as e:
        print(f"[swap_debug] 失败: {e}")
        traceback.print_exc()
        return jsonify({"state": -1, "msg": f"换发型失败: {e}"}), 500


@app.route("/api/swap_hairline", methods=["POST"])
def api_swap_hairline():
    """发际线带重绘接口（实验性，流程同 swap_debug，仅步骤③改为边界带）。

    比 swap_debug 多一个参数：
        method: str, 发际线mask方案，可选 boundary_band/mediapipe/landmark_1k/deeplab（默认 mediapipe）
        band_width: int, 边界带宽度（形态学核大小，默认15，带宽≈2*band_width）
    其余参数同 swap_debug。
    """
    t0 = time.time()
    try:
        d = request.json
        img_b64 = d.get("img", "")
        hair_id = d.get("hair_id", "")
        if not img_b64 or not hair_id:
            return jsonify({"state": -1, "msg": "img 和 hair_id 不能为空"}), 400

        img = _b64_to_ndarray(img_b64, color=True)
        if img is None:
            return jsonify({"state": -1, "msg": "img 解析失败"}), 400

        def g(k, default): return d.get(k, default)
        dk = g("dilate_kernel", [6, 18])
        bd = g("blend_dilate", [5, 5])
        sd = g("seamless_dilate", [9, 9])

        print(f"[swap_hairline] hair_id={hair_id}, method={g('method','mediapipe')}, "
              f"band_width={g('band_width',15)}, denoising={g('denoising_strength',0.6)}")

        hairstyle_process, landmark_processor = _get_models()
        from hair_swap_hairline import hair_swap_hairline
        task_id = f"hairline_{int(time.time()*1000)}"
        result, steps, params = hair_swap_hairline(
            origin_img=img, hair_id=hair_id,
            hairstyle_process=hairstyle_process,
            landmark_processor=landmark_processor,
            task_id=task_id,
            method=str(g("method", "mediapipe")),
            strict_mask=bool(g("strict_mask", False)),
            seamless_blend=bool(g("seamless_blend", True)),
            is_hr=bool(g("is_hr", True)),
            dilate_kernel=(int(dk[0]), int(dk[1])),
            denoising_strength=float(g("denoising_strength", 0.6)),
            blend_dilate=(int(bd[0]), int(bd[1])),
            seamless_dilate=(int(sd[0]), int(sd[1])),
            band_width=int(g("band_width", 15)),
            preview_only=bool(g("preview_only", False)),
            height_ratio=float(g("height_ratio", 0.432)),
            width_ratio=float(g("width_ratio", 0.144)),
            corner_ratio=float(g("corner_ratio", 0.25)),
            vertical_offset=float(g("vertical_offset", 0.0)),
            refiner_switch_at=float(g("refiner_switch_at", 0.5)),
        )

        _, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
        result_b64 = base64.b64encode(buf).decode("utf-8")
        print(f"[swap_hairline] 完成，总耗时 {time.time()-t0:.1f}s")
        return jsonify({"state": 0, "result": result_b64, "steps": steps, "params": params})

    except Exception as e:
        print(f"[swap_hairline] 失败: {e}")
        traceback.print_exc()
        return jsonify({"state": -1, "msg": f"发际线带重绘失败: {e}"}), 500


@app.route("/api/swap_manual", methods=["POST"])
def api_swap_manual():
    """手绘mask版换发型接口（重绘区域完全由用户手绘mask决定）。

    入参 JSON:
        img: 原图 base64
        mask: 用户手绘mask base64（白色=重绘区，与img同分辨率）
        hair_id: 发型ID
        其余参数同 swap_debug（strict_mask/seamless_blend/is_hr/dilate_kernel/
                  denoising_strength/blend_dilate/seamless_dilate）
    返回: {state, result, steps, params}
    """
    t0 = time.time()
    try:
        d = request.json
        img_b64 = d.get("img", "")
        mask_b64 = d.get("mask", "")
        hair_id = d.get("hair_id", "")
        if not img_b64 or not mask_b64:
            return jsonify({"state": -1, "msg": "img 和 mask 不能为空（请先手绘重绘区域）"}), 400
        if not hair_id:
            return jsonify({"state": -1, "msg": "hair_id 不能为空"}), 400

        img = _b64_to_ndarray(img_b64, color=True)
        hand_mask = _b64_to_ndarray(mask_b64, color=False)
        if img is None:
            return jsonify({"state": -1, "msg": "img 解析失败"}), 400
        if hand_mask is None:
            return jsonify({"state": -1, "msg": "mask 解析失败"}), 400

        def g(k, default): return d.get(k, default)
        dk = g("dilate_kernel", [6, 18])
        bd = g("blend_dilate", [5, 5])
        sd = g("seamless_dilate", [9, 9])

        print(f"[swap_manual] hair_id={hair_id}, img={img.shape}, mask={hand_mask.shape}, "
              f"strict_mask={g('strict_mask',False)}, denoising={g('denoising_strength',0.6)}")

        hairstyle_process, landmark_processor = _get_models()
        from hair_swap_manual import hair_swap_manual
        task_id = f"manual_{int(time.time()*1000)}"
        result, steps, params = hair_swap_manual(
            origin_img=img, hand_mask=hand_mask, hair_id=hair_id,
            hairstyle_process=hairstyle_process,
            landmark_processor=landmark_processor,
            task_id=task_id,
            strict_mask=bool(g("strict_mask", False)),
            seamless_blend=bool(g("seamless_blend", True)),
            is_hr=bool(g("is_hr", True)),
            dilate_kernel=(int(dk[0]), int(dk[1])),
            denoising_strength=float(g("denoising_strength", 0.6)),
            blend_dilate=(int(bd[0]), int(bd[1])),
            seamless_dilate=(int(sd[0]), int(sd[1])),
            feather_px=int(g("feather_px", 0)),
            enhance=bool(g("enhance", False)),
            enhance_denoising=float(g("enhance_denoising", 0.35)),
        )

        _, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
        result_b64 = base64.b64encode(buf).decode("utf-8")
        print(f"[swap_manual] 完成，总耗时 {time.time()-t0:.1f}s")
        return jsonify({"state": 0, "result": result_b64, "steps": steps, "params": params})

    except Exception as e:
        print(f"[swap_manual] 失败: {e}")
        traceback.print_exc()
        return jsonify({"state": -1, "msg": f"手绘mask换发型失败: {e}"}), 500


def _b64_to_ndarray(b64_str, color=True):
    if "," in b64_str and b64_str.startswith("data:"):
        b64_str = b64_str.split(",", 1)[1]
    data = base64.b64decode(b64_str)
    flag = cv2.IMREAD_COLOR if color else cv2.IMREAD_GRAYSCALE
    return cv2.imdecode(np.frombuffer(data, np.uint8), flag)


@app.route("/api/hairstyles")
def api_hairstyles():
    """返回可用发型列表（含性别，用于前端下拉框分组）"""
    try:
        from common.logger import config
        hairstyle_dir = config.get('default', 'hairstyleDir')
        train_dir = config.get('default', 'train_dir')
        upload_dir = config.get('default', 'upload_train_dir')
        from hair_grow_swap import list_hairstyles
        styles = list_hairstyles(hairstyle_dir, train_dir, upload_dir)
        boy = [s for s in styles if s["gender"] == "boy"]
        girl = [s for s in styles if s["gender"] == "girl"]
        print(f"[hairstyles] 共 {len(styles)} 个可用发型 (boy={len(boy)}, girl={len(girl)})")
        return jsonify({"state": 0, "data": styles, "count": len(styles)})
    except Exception as e:
        traceback.print_exc()
        return jsonify({"state": -1, "msg": str(e)}), 500


@app.route("/preview/<hair_id>")
def preview_img(hair_id):
    """直接返回预览图文件（比 base64 API 快，浏览器可缓存）"""
    effect_path = os.path.join(BASE_DIR, "static", "previews", f"{hair_id}.jpg")
    if os.path.exists(effect_path):
        return send_from_directory(os.path.join(BASE_DIR, "static", "previews"), f"{hair_id}.jpg")
    # 回退到 ref_rgb
    try:
        from common.logger import config
        hairstyle_dir = config.get('default', 'hairstyleDir')
        fallback = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
        if os.path.exists(fallback):
            return send_file(fallback)
    except Exception:
        pass
    return ("", 404)


@app.route("/train_src/<hair_id>")
def train_src_img(hair_id):
    """返回发型的训练原图（hair_type_images/<hair_id>.jpg/.png）。
    用于测试页展示发型真实样子，而非套在标准脸上的效果图。
    """
    src_dir = "/home/ubuntu/change_hair/hair_type_images"
    for ext in (".jpg", ".jpeg", ".png"):
        path = os.path.join(src_dir, hair_id + ext)
        if os.path.exists(path):
            return send_from_directory(src_dir, hair_id + ext)
    return ("", 404)


@app.route("/api/hairstyle_preview/<hair_id>")
def api_hairstyle_preview(hair_id):
    """返回某发型的预览图。
    优先用生成的效果图 static/previews/<hair_id>.jpg（发型套在标准脸上的样子），
    没有则回退到 first## 原始上传图。
    """
    try:
        # 1. 优先：生成的效果图（发型套在 boy/girl 标准脸上）
        preview = None
        effect_path = os.path.join(BASE_DIR, "static", "previews", f"{hair_id}.jpg")
        if os.path.exists(effect_path):
            preview = effect_path
        else:
            # 2. 回退：first## 原始上传图
            from common.logger import config
            upload_dir = config.get('default', 'upload_train_dir')
            save_dir = os.path.join(upload_dir, hair_id)
            if os.path.isdir(save_dir):
                for name in os.listdir(save_dir):
                    if "first##" in name:
                        preview = os.path.join(save_dir, name)
                        break
            # 3. 再回退：ref_rgb
            if not preview or not os.path.exists(preview):
                hairstyle_dir = config.get('default', 'hairstyleDir')
                fallback = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
                if os.path.exists(fallback):
                    preview = fallback
        if not preview:
            return jsonify({"state": -1, "msg": "无预览图"}), 404
        img = cv2.imread(preview)
        img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)
        _, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 85])
        b64 = base64.b64encode(buf).decode()
        return jsonify({"state": 0, "preview": b64})
    except Exception as e:
        return jsonify({"state": -1, "msg": str(e)}), 500


@app.route("/api/grow", methods=["POST"])
def api_grow():
    """生发接口（走换发型工作流）

    入参 JSON:
        img: 原图 base64
        mask: 手绘遮罩 base64（白=生发区，与img同分辨率）
        hair_id: 选择的发型ID（必填）
        is_hr: 是否高清 "true"/"false"，默认 "true"
    """
    t0 = time.time()
    try:
        data = request.json
        img_b64 = data.get("img", "")
        mask_b64 = data.get("mask", "")
        hair_id = data.get("hair_id", "")
        is_hr = str(data.get("is_hr", "true")).lower() == "true"

        if not img_b64 or not mask_b64:
            return jsonify({"state": -1, "msg": "img 和 mask 不能为空"}), 400
        if not hair_id:
            return jsonify({"state": -1, "msg": "hair_id 不能为空（请先选择发型）"}), 400

        img = _b64_to_ndarray(img_b64, color=True)
        mask = _b64_to_ndarray(mask_b64, color=False)
        if img is None:
            return jsonify({"state": -1, "msg": "img 解析失败"}), 400
        if mask is None:
            return jsonify({"state": -1, "msg": "mask 解析失败"}), 400

        print(f"[grow] img={img.shape}, mask={mask.shape}, hair_id={hair_id}, is_hr={is_hr}")

        # 加载模型（首次慢）
        hairstyle_process, landmark_processor = _get_models()

        # 调用生发（走换发型工作流）
        from hair_grow_swap import hair_grow_swap
        task_id = f"grow_{int(time.time()*1000)}"
        result = hair_grow_swap(
            origin_img=img, hand_mask=mask, hair_id=hair_id,
            hairstyle_process=hairstyle_process,
            landmark_processor=landmark_processor,
            task_id=task_id, is_hr=is_hr)

        _, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
        result_b64 = base64.b64encode(buf).decode("utf-8")
        print(f"[grow] 完成，总耗时 {time.time()-t0:.1f}s")
        return jsonify({"state": 0, "result": result_b64})

    except Exception as e:
        print(f"[grow] 失败: {e}")
        traceback.print_exc()
        return jsonify({"state": -1, "msg": f"生发失败: {e}"}), 500


if __name__ == "__main__":
    print(f"生发服务启动，端口 {PORT}")
    print(f"测试页面: http://0.0.0.0:{PORT}")
    print(f"注意：首次请求会加载换发型模型（约30-60秒）")
    server = pywsgi.WSGIServer(("0.0.0.0", PORT), app)
    server.serve_forever()