import streamlit as st
import cv2
import numpy as np
from rembg import remove, new_session
import io
from PIL import Image

# ==========================================
# 核心渲染器：绝对精准颜色匹配
# ==========================================
def render_exact_color(orig_img, mask_3d, target_lab):
    """确保100%零色差，保留褶皱"""
    orig_lab = cv2.cvtColor(orig_img, cv2.COLOR_BGR2LAB).astype(np.float32)
    l_t, a_t, b_t = target_lab

    l_orig = orig_lab[:, :, 0]

    mask_bool = mask_3d[:, :, 0] > 0.5
    if not np.any(mask_bool):
        return orig_img

    l_mean_orig = np.mean(l_orig[mask_bool])

    # 亮度对齐，保留褶皱
    l_new = l_orig - l_mean_orig + l_t
    l_new = np.clip(l_new, 0, 255).astype(np.uint8)

    a_new = np.full_like(l_orig, a_t, dtype=np.uint8)
    b_new = np.full_like(l_orig, b_t, dtype=np.uint8)

    new_lab = cv2.merge([l_new, a_new, b_new])
    new_bgr = cv2.cvtColor(new_lab, cv2.COLOR_LAB2BGR)

    final_out = new_bgr.astype(np.float32) * mask_3d + orig_img.astype(np.float32) * (1.0 - mask_3d)

    return np.clip(final_out, 0, 255).astype(np.uint8)

# ==========================================
# 工具函数
# ==========================================
def get_lab_metrics(img_bgr):
    """获取参考图片中心区域的 LAB 颜色均值"""
    img_lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
    h, w = img_bgr.shape[:2]
    return np.mean(img_lab[int(h * 0.4):int(h * 0.6), int(w * 0.4):int(w * 0.6)], axis=(0, 1))

def generate_ai_mask(orig_bytes, orig_bgr, shape):
    """【终极修复版】人体轮廓 + 肤色剔除 + 暗色剔除 + 背景剔除"""
    session = new_session("u2net")
    output_data = remove(orig_bytes, session=session)
    nparr = np.frombuffer(output_data, np.uint8)
    rgba_img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)

    person_mask = rgba_img[:, :, 3]
    person_mask = cv2.resize(person_mask, (shape[1], shape[0]))
    _, person_mask = cv2.threshold(person_mask, 200, 255, cv2.THRESH_BINARY)

    hsv = cv2.cvtColor(orig_bgr, cv2.COLOR_BGR2HSV)

    lower_skin = np.array([0, 10, 30], dtype=np.uint8)
    upper_skin = np.array([30, 255, 255], dtype=np.uint8)
    skin_mask = cv2.inRange(hsv, lower_skin, upper_skin)
    skin_mask = cv2.dilate(skin_mask, np.ones((5, 5), np.uint8), iterations=3)

    dark_mask = cv2.inRange(hsv, np.array([0, 0, 0]), np.array([180, 255, 50]))
    dark_mask = cv2.dilate(dark_mask, np.ones((7, 7), np.uint8), iterations=2)

    white_mask = cv2.inRange(hsv, np.array([0, 0, 210]), np.array([180, 30, 255]))
    white_mask = cv2.dilate(white_mask, np.ones((5, 5), np.uint8), iterations=2)

    exclude_mask = cv2.bitwise_or(skin_mask, dark_mask)
    exclude_mask = cv2.bitwise_or(exclude_mask, white_mask)

    clothes_mask = cv2.bitwise_and(person_mask, cv2.bitwise_not(exclude_mask))
    clothes_mask = cv2.morphologyEx(clothes_mask, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))
    clothes_mask = cv2.morphologyEx(clothes_mask, cv2.MORPH_CLOSE, np.ones((15, 15), np.uint8))

    contours, _ = cv2.findContours(clothes_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        largest_contour = max(contours, key=cv2.contourArea)
        clean_mask = np.zeros_like(clothes_mask)
        cv2.drawContours(clean_mask, [largest_contour], -1, 255, thickness=cv2.FILLED)
        clothes_mask = clean_mask

    clothes_mask = cv2.erode(clothes_mask, np.ones((3, 3), np.uint8), iterations=1)
    clothes_mask = cv2.GaussianBlur(clothes_mask, (7, 7), 0)

    mask_3d = np.repeat((clothes_mask.astype(np.float32) / 255.0)[:, :, np.newaxis], 3, axis=2)
    return mask_3d

# ==========================================
# Streamlit 网页前端界面
# ==========================================
st.set_page_config(page_title="AI 衣服换色系统", layout="wide")
st.title("🎨 AI 衣服精准换色系统")
st.markdown("上传原图和目标颜色图，AI将自动提取衣物并进行零色差渲染。")

# 侧边栏：上传区域
with st.sidebar:
    st.header("1. 上传图片")
    orig_file = st.file_uploader("上传原图 (需要换色的衣服)", type=['jpg', 'jpeg', 'png'])
    ref_files = st.file_uploader("上传参考图 (目标颜色，可多选)", type=['jpg', 'jpeg', 'png'], accept_multiple_files=True)
    process_btn = st.button("🚀 开始极速换色", use_container_width=True)

# 主显示区域
if orig_file is not None:
    # 将上传的文件转为 OpenCV 格式
    orig_bytes = orig_file.getvalue()
    orig_nparr = np.frombuffer(orig_bytes, np.uint8)
    orig_bgr = cv2.imdecode(orig_nparr, cv2.IMREAD_COLOR)
    
    # 界面预览原图
    st.subheader("原图预览")
    st.image(cv2.cvtColor(orig_bgr, cv2.COLOR_BGR2RGB), width=400)

    if process_btn:
        if not ref_files:
            st.warning("⚠️ 请至少上传一张参考颜色图！")
        else:
            with st.spinner('AI 正在提取衣服轮廓和处理换色，请稍候 (首次运行需下载模型)...'):
                shape = orig_bgr.shape[:2]
                # 1. 提取蒙版
                mask_3d = generate_ai_mask(orig_bytes, orig_bgr, shape)

                st.subheader("🎉 换色结果")
                # 创建多列布局展示结果
                cols = st.columns(len(ref_files))
                
                # 2. 遍历参考图换色
                for idx, ref_file in enumerate(ref_files):
                    ref_nparr = np.frombuffer(ref_file.getvalue(), np.uint8)
                    ref_bgr = cv2.imdecode(ref_nparr, cv2.IMREAD_COLOR)
                    
                    target_lab = get_lab_metrics(ref_bgr)
                    final_img = render_exact_color(orig_bgr, mask_3d, target_lab)
                    
                    # 将 BGR 转为 RGB 供网页显示
                    final_rgb = cv2.cvtColor(final_img, cv2.COLOR_BGR2RGB)
                    
                    with cols[idx]:
                        st.image(cv2.cvtColor(ref_bgr, cv2.COLOR_BGR2RGB), caption=f"参考色: {ref_file.name}", width=150)
                        st.image(final_rgb, caption=f"换色完成: {ref_file.name}", use_column_width=True)
                        
                        # 提供下载按钮
                        result_pil = Image.fromarray(final_rgb)
                        buf = io.BytesIO()
                        result_pil.save(buf, format="JPEG")
                        st.download_button(
                            label=f"⬇️ 下载该结果",
                            data=buf.getvalue(),
                            file_name=f"result_{ref_file.name}",
                            mime="image/jpeg"
                        )