import os
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
import gradio as gr
import onnxruntime as ort
import importlib.util
from huggingface_hub import hf_hub_download
import numpy as np
import cv2
import psutil
import time
import sys

token_ = os.environ.get("HF_TOKEN")
model_repo_id = "heyh97791/NCF" 
model_filename = "ncf.onnx"
file_path = "skin_protection.py"
pretrained_name = "skin.tflite"

protector_instance = None

I18N = {
    "lang_btn": {
        "en": "中文",  
        "zh": "English" 
    },
    "title": {
        "en": "# 🎨 ColorFM-L [CPU] ",
        "zh": "# 🎨 ColorFM-L [CPU] "
    },

    "paper_code": {
        "en": """
        📄 **Paper:** [ColorFM: An Optimization-to-Learning Framework for Color Transfer via Flow Matching (ECCV 2026)](https://github.com/cszn/ColorFM)  |  💻 **Code:** [GitHub](https://github.com/cszn/ColorFM)  |  ⚡ **Model:** Feed-forward model
        """,
        "zh": """
        📄 **论文：** [ColorFM: An Optimization-to-Learning Framework for Color Transfer via Flow Matching (ECCV 2026)](https://github.com/cszn/ColorFM)  |  💻 **代码：** [GitHub](https://github.com/cszn/ColorFM)  |  ⚡ **模型：** 前馈式模型
        """
    },
    "desc_header": {
        "en": "💡 What is Color Transfer?",
        "zh": "💡 追色是什么"
    },
    "desc_content": {
        "en": """
        **1. What is "Color Transfer"?**
        It is a technique that extracts the **color palette and atmosphere** from a reference image (Style) and applies it to your target image (Content), while preserving the original structure.
        
        **2. Use Cases:**
        * 📸 **Photography:** Instantly mimic color grading styles from master photographers.
        * 🎨 **Art & Design:** Unify the color theme of different assets quickly.
        
        **3. Tips for Best Results:**
        * ✅ **Match Content:** Results are best when the content and style images share similar scenes (e.g., Landscape to Landscape).
        * 🎨 **Hue Similarity:** Images with closer hues between content and style may produce better color transfer results.
        * ⚡ **High Resolution:** Supports processing high-resolution images (use the Settings menu to optimize speed).
        
        **🛡️ Privacy Disclaimer:**
        This demo runs entirely on the cloud instance. Your uploaded images are processed in memory and are **NOT saved** or stored permanently on our servers.
        """,
        "zh": """
        **1. 什么是“追色”？**
        追色（Color Transfer）是指从一张参考图（风格图）中提取**色调与氛围**，并将其“迁移”到你的目标图片（内容图）上，同时保留原图的细节纹理。
        
        **2. 主要用途：**
        * 📸 **摄影后期：** 一键复刻摄影大师的调色风格，无需手动调参。
        * 🎨 **设计创作：** 快速统一多张素材的色调风格，提高创作效率。
       
       **3. 如何获得最佳效果？**
        * ✅ **内容匹配：** 当内容图与风格图的场景相似时（例如都是风景），效果通常最好。
        * 🎨 **色相相似：** 内容图像和风格图像之间色相相似，可能会呈现更好的效果
        * ⚡ **高清支持：** 支持高分辨率图片处理（可在“设置”中调整分辨率以获得更快速度）。
        
        **🛡️ 免责声明：**
        本Demo仅供演示体验。您的图片仅在内存中进行临时处理，**不会被保存**、存储或用于任何其他用途，处理结束后即刻销毁。
        """
    },
    "label_content": {
        "en": "Content Image ",
        "zh": "内容图"
    },
    "label_style": {
        "en": "Style Reference",
        "zh": "色彩参考图"
    },
    "settings_header": {
        "en": "⚙️ Settings",
        "zh": "⚙️ 设置"
    },
    "res_label": {
        "en": "Max Resolution Limit",
        "zh": "最大分辨率限制"
    },
    "res_info": {
        "en": "Limits the input image size to speed up inference. 'Original' keeps original size.",
        "zh": "限制输入图像尺寸以加快推理速度。选择“Original”将保持原图尺寸。"
    },
    "prot_label": { 
        "en": "🧪 Enable Skin Protection (Experimental)", 
        "zh": "🧪 肤色保护 (实验性功能)" 
    },
    "btn_run": {
        "en": "🚀 Start Color Transfer",
        "zh": "🚀 开始追色"
    },
    "label_result": {
        "en": "Result Image",
        "zh": "追色结果"
    },
    "example_label": {
        "en": "⚡ Quick Examples (Click to Load)",
        "zh": "⚡ 追色样例 (点击加载示例)"
    }
}

try:
    model_path = hf_hub_download(
        repo_id=model_repo_id,
        filename=model_filename,
        token=token_ 
    )
except Exception as e:
    print(f"{e}")

sess_options = ort.SessionOptions()
sess_options.enable_cpu_mem_arena = False
sess_options.intra_op_num_threads = 1
sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']

try:
    ort_session = ort.InferenceSession(model_path, providers=providers, sess_options=sess_options)
    input_names = [node.name for node in ort_session.get_inputs()] 
except Exception as e:
    print(f"{e}")
    ort_session = None

try:
    script_path = hf_hub_download(
        repo_id=model_repo_id,
        filename=file_path,
        token=token_
    )
    skin_model_path = hf_hub_download(
        repo_id=model_repo_id,
        filename=pretrained_name,
        token=token_
    )
    
    spec = importlib.util.spec_from_file_location("skin_protection", script_path)
    secret_module = importlib.util.module_from_spec(spec)
    sys.modules["skin_protection"] = secret_module
    spec.loader.exec_module(secret_module)

    protector_instance = secret_module.SkinProtector(skin_model_path)
    print("✅ Skin protection module loaded successfully.")

except Exception as e:
    print(f"⚠️ Skin protection not loaded (Running in standard mode): {e}")
    protector_instance = None

def get_memory_mb():
    pid = os.getpid()
    return psutil.Process(pid).memory_info().rss / 1024 / 1024

RESOLUTION_MAP = {
    "Original": 0,
    "4K": 3840 * 2160,     # ~8.3MP
    "2K": 2560 * 1440,     # ~3.7MP
    "1080p": 1920 * 1080,    # ~2.1MP
    "720p": 1280 * 720       # ~0.9MP
}

def preprocess_image(img_rgb, resolution_choice="Original"):
    if img_rgb is None:
        return None

    max_pixel_limit = RESOLUTION_MAP.get(str(resolution_choice).split(" ")[0], 0)

    if max_pixel_limit > 0:
        h, w = img_rgb.shape[:2]
        num_pixels = h * w
        
        if num_pixels > max_pixel_limit:
            scale_factor = np.sqrt(max_pixel_limit / num_pixels)
            new_w = int(w * scale_factor)
            new_h = int(h * scale_factor)

            img_rgb = cv2.resize(img_rgb, (new_w, new_h), interpolation=cv2.INTER_AREA)
            # print(f"Resized image from {w}x{h} to {new_w}x{new_h}")

    img = img_rgb.astype(np.float32) / 255.0
    img = img.transpose(2, 0, 1)
    img = np.expand_dims(img, axis=0)
    return img

def postprocess_tensor(tensor):
    img = np.squeeze(tensor, axis=0)
    img = np.clip(img, 0, 1)
    img = img.transpose(1, 2, 0)
    img = (img * 255.0).astype(np.uint8)
    return img


def run_inference_numpy(content_np, style_np):
    if ort_session is None: return content_np
    c_in = content_np.astype(np.float32) / 255.0
    c_in = c_in.transpose(2, 0, 1)[np.newaxis, :, :, :]
    
    s_in = style_np.astype(np.float32) / 255.0
    s_in = s_in.transpose(2, 0, 1)[np.newaxis, :, :, :]
    
    input_feed = {input_names[0]: c_in, input_names[1]: s_in}
    outputs = ort_session.run(None, input_feed)
    out_tensor = outputs[0]
    out_img = np.squeeze(out_tensor, axis=0)
    out_img = np.clip(out_img, 0, 1)
    out_img = out_img.transpose(1, 2, 0)
    out_img = (out_img * 255.0).astype(np.uint8)
    
    return out_img

def inference(content_img, style_img, resolution_selection, enable_skin_protection):
    if ort_session is None:
        raise gr.Error("model do not exists")
        
    if content_img is None or style_img is None:
        return None

    # mem_start = get_memory_mb()
    # t_start = time.time()

    try:
        content_input = preprocess_image(content_img, resolution_selection)
        style_input = preprocess_image(style_img, resolution_selection)

        input_feed = {
            input_names[0]: content_input, 
            input_names[1]: style_input
        }

        outputs = ort_session.run(None, input_feed)
        final_output_tensor = outputs[0]

        result_img = postprocess_tensor(final_output_tensor)

        if enable_skin_protection and protector_instance is not None:
            # print("skin protection activated")
            result_img = protector_instance.process(
                content_img=content_img, 
                style_img=style_img,        
                global_result_img=result_img, 
                inference_callback=run_inference_numpy 
            )

        return result_img

    except Exception as e:
        print(f"{e}")
        raise gr.Error(f"{str(e)}")

def toggle_language(current_lang):
    target_lang = "zh" if current_lang == "en" else "en"
    
    return (
        target_lang,                                
        gr.update(value=I18N["lang_btn"][target_lang]),      # Button
        gr.update(value=I18N["title"][target_lang]),         # Markdown
        gr.update(label=I18N["desc_header"][target_lang]),   # Accordion 
        gr.update(value=I18N["desc_content"][target_lang]),  # Markdown
        gr.update(label=I18N["label_content"][target_lang]), # Image Label
        gr.update(label=I18N["label_style"][target_lang]),   # Image Label
        gr.update(value=I18N["btn_run"][target_lang]),       # Button
        gr.update(label=I18N["label_result"][target_lang]),  # Image Label
        gr.update(label=I18N["settings_header"][target_lang]), # Accordion Settings
        gr.update(label=I18N["res_label"][target_lang], info=I18N["res_info"][target_lang]), # Radio Label
        gr.update(label=I18N["prot_label"][target_lang]),
    )


custom_css = """
#col-container {
    margin: 0 auto;
    max-width: 1100px;
}
.gallery-container img{
    object-fit: contain;
}
"""

with gr.Blocks(css=custom_css, title="Color-Transfer") as demo:
    lang_state = gr.State("en")
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=5):
                md_title = gr.Markdown(I18N["title"]["en"])
            with gr.Column(scale=0, min_width=80): 
                btn_lang = gr.Button(
                    value=I18N["lang_btn"]["en"], 
                    variant="secondary", 
                    size="sm" 
                )

        md_paper_code = gr.Markdown(I18N["paper_code"]["en"])
        with gr.Accordion(label=I18N["desc_header"]["en"], open=False) as acc_desc:
            md_desc = gr.Markdown(I18N["desc_content"]["en"])
        
        with gr.Row():
            with gr.Column():
                input_content = gr.Image(label=I18N["label_content"]["en"], type="numpy", height=300, interactive=True)
            with gr.Column():
                input_style = gr.Image(label=I18N["label_style"]["en"], type="numpy", height=300, interactive=True)

        with gr.Accordion(label=I18N["settings_header"]["en"], open=False) as acc_settings:
            res_choices = ["Original", "4K", "2K", "1080p", "720p"]
            radio_res = gr.Radio(
                choices=res_choices,
                value="Original",
                label=I18N["res_label"]["en"],
                info=I18N["res_info"]["en"],
                interactive=True
            )
            with gr.Row():
                chk_skin = gr.Checkbox(
                    label=I18N["prot_label"]["en"], 
                    value=False,
                    interactive=True
                )
                
        with gr.Row():
            btn_run = gr.Button(I18N["btn_run"]["en"], variant="primary")
        
        with gr.Row():
            output_result = gr.Image(label=I18N["label_result"]["en"], type="numpy", interactive=False, height=450, format="png")
    
        btn_run.click(
            fn=lambda: None,           
            inputs=None, 
            outputs=output_result
        ).then(                    
            fn=inference,
            inputs=[input_content, input_style, radio_res, chk_skin],
            outputs=output_result
        )

        gr.Examples(
            label=I18N["example_label"]["en"],
            examples=[
                ["./figs/01_c.jpg", "./figs/01_s.jpg", "./figs/01.jpg"], 
                ["./figs/02_c.jpg", "./figs/02_s.jpg", "./figs/02.jpg"], 
                ["./figs/03_c.jpg", "./figs/03_s.jpg", "./figs/03.jpg"], 
                ["./figs/04_c.jpg", "./figs/04_s.jpg", "./figs/04.png"], 
            ],
            
            inputs=[input_content, input_style, output_result], 
    
        )

        btn_lang.click(
            fn=toggle_language,
            inputs=[lang_state],
            outputs=[
                lang_state,     
                btn_lang,    
                md_title,    
                acc_desc,     
                md_desc,       
                input_content,  
                input_style,  
                btn_run,   
                output_result,
                acc_settings, 
                radio_res,
                chk_skin, 
            ]
        )

if __name__ == "__main__":
    demo.launch()