import os

# ============================================================
# CPU / RUNTIME ENVIRONMENT
# ============================================================

os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["ORT_DISABLE_CUDA"] = "1"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = ""

# ============================================================
# STANDARD LIBRARY
# ============================================================

import sys
import asyncio
import gc
import json
import time
import tempfile
import importlib
import threading
import shutil

from collections import OrderedDict

# ============================================================
# THIRD PARTY
# ============================================================

import numpy as np
import cv2
import gradio as gr
import imageio

from PIL import Image
from huggingface_hub import hf_hub_download

# ============================================================
# CONFIG
# ============================================================

APP_TITLE = "CPU Multi Face Swap Pro"

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

MODEL_DIR = os.path.join(
    BASE_DIR,
    "models"
)

INSIGHTFACE_DIR = os.path.join(
    MODEL_DIR,
    "insightface"
)

HYPERSWAP_DIR = os.path.join(
    MODEL_DIR,
    "hyperswap"
)

FACE_RESTORE_DIR = os.path.join(
    MODEL_DIR,
    "facerestore_models"
)

OUTPUT_DIR = os.path.join(
    BASE_DIR,
    "output"
)

COMFYUI_ROOT = BASE_DIR

REACTOR_PATH = os.path.join(
    COMFYUI_ROOT,
    "custom_nodes",
    "comfyui-reactor-node"
)

if not os.path.isdir(REACTOR_PATH):

    REACTOR_PATH = os.path.join(
        BASE_DIR,
        "custom_nodes",
        "comfyui-reactor-node"
    )

DEFAULT_SWAP_MODEL = (
    "inswapper_128.onnx"
)

DEFAULT_RESTORE_MODEL = (
    "GPEN-BFR-512.onnx"
)

DEFAULT_RESTORE_STRENGTH = 0.70

# ============================================================
# 16 GB RAM / 2 JOB CONFIG
# ============================================================

MAX_CONCURRENT_JOBS = 2
MAX_CONCURRENT_INFERENCE = 2
MAX_PREVIEW_JOBS = 2
GRADIO_QUEUE_SIZE = 8

# ============================================================
# GLOBAL LOCKS
# ============================================================

JOB_SEMAPHORE = threading.BoundedSemaphore(
    MAX_CONCURRENT_JOBS
)

INFERENCE_SEMAPHORE = threading.BoundedSemaphore(
    MAX_CONCURRENT_INFERENCE
)

MODEL_LOCK = threading.RLock()
COMFY_INIT_LOCK = threading.RLock()

# ============================================================
# GLOBAL STATE
# ============================================================

loaded_models = {}

face_cache = OrderedDict()
FACE_CACHE_MAX = 64

source_tensor_cache = OrderedDict()
SOURCE_TENSOR_CACHE_MAX = 16

# ============================================================
# LOG
# ============================================================

def log(*args):

    print(
        *args,
        flush=True
    )

# ============================================================
# STARTUP INFO
# ============================================================

log("=" * 70)

log(
    "[APP] BASE DIR:",
    BASE_DIR
)

log(
    "[APP] MODEL DIR:",
    MODEL_DIR
)

log(
    "[APP] OUTPUT DIR:",
    OUTPUT_DIR
)

log(
    "[APP] MAX JOBS:",
    MAX_CONCURRENT_JOBS
)

log(
    "[APP] MAX INFERENCE:",
    MAX_CONCURRENT_INFERENCE
)

log("=" * 70)

# ============================================================
# DIRECTORIES
# ============================================================

for directory in (
    MODEL_DIR,
    INSIGHTFACE_DIR,
    HYPERSWAP_DIR,
    FACE_RESTORE_DIR,
    OUTPUT_DIR
):

    os.makedirs(
        directory,
        exist_ok=True
    )

# ============================================================
# MODEL DEFINITIONS
# ============================================================

MODEL_REPOS = {

    "inswapper_128.onnx": (
        "ezioruan/inswapper_128.onnx",
        "inswapper_128.onnx",
        INSIGHTFACE_DIR
    ),

    "hyperswap_1a_256.onnx": (
        "facefusion/models-3.3.0",
        "hyperswap_1a_256.onnx",
        HYPERSWAP_DIR
    ),

    "hyperswap_1b_256.onnx": (
        "facefusion/models-3.3.0",
        "hyperswap_1b_256.onnx",
        HYPERSWAP_DIR
    ),

    "hyperswap_1c_256.onnx": (
        "facefusion/models-3.3.0",
        "hyperswap_1c_256.onnx",
        HYPERSWAP_DIR
    ),

    "GPEN-BFR-512.onnx": (
        "martintomov/comfy",
        "facerestore_models/GPEN-BFR-512.onnx",
        FACE_RESTORE_DIR
    )
}

# ============================================================
# MODEL DOWNLOAD
# ============================================================

def download_model(model_name):

    if model_name not in MODEL_REPOS:

        raise RuntimeError(
            f"Unknown model: {model_name}"
        )

    repo, filename, local_dir = (
        MODEL_REPOS[
            model_name
        ]
    )

    basename = os.path.basename(
        filename
    )

    local_path = os.path.join(
        local_dir,
        basename
    )

    if os.path.isfile(
        local_path
    ):

        log(
            "[MODEL] Already exists:",
            local_path
        )

        return local_path

    os.makedirs(
        local_dir,
        exist_ok=True
    )

    # --------------------------------------------------------
    # MIGRATE OLD BROKEN NESTED PATH
    # --------------------------------------------------------

    nested_path = os.path.join(
        local_dir,
        os.path.dirname(filename),
        basename
    )

    if (
        os.path.isfile(nested_path)
        and not os.path.isfile(local_path)
    ):

        log(
            "[MODEL] Migrating old nested model:",
            nested_path
        )

        try:

            shutil.copy2(
                nested_path,
                local_path
            )

            if os.path.isfile(
                local_path
            ):

                log(
                    "[MODEL] Migrated:",
                    local_path
                )

                return local_path

        except Exception as e:

            log(
                "[MODEL] Migration warning:",
                repr(e)
            )

    # --------------------------------------------------------
    # DOWNLOAD
    # --------------------------------------------------------

    log(
        "[MODEL] Downloading:",
        model_name
    )

    download_dir = tempfile.mkdtemp(
        prefix="model_download_"
    )

    try:

        downloaded_path = hf_hub_download(
            repo_id=repo,
            filename=filename,
            local_dir=download_dir
        )

        if not os.path.isfile(
            downloaded_path
        ):

            raise RuntimeError(
                "Model download failed:\n"
                + str(downloaded_path)
            )

        shutil.copy2(
            downloaded_path,
            local_path
        )

    finally:

        try:

            shutil.rmtree(
                download_dir,
                ignore_errors=True
            )

        except Exception:
            pass

    if not os.path.isfile(
        local_path
    ):

        raise RuntimeError(
            "Model was downloaded but "
            "could not be placed at:\n"
            + local_path
        )

    log(
        "[MODEL] Ready:",
        local_path
    )

    return local_path

# ============================================================
# ENSURE MODEL
# ============================================================

def ensure_model(model_name):

    if not model_name:
        return None

    if model_name == "none":
        return None

    with MODEL_LOCK:

        cached = loaded_models.get(
            model_name
        )

        if (
            cached
            and os.path.isfile(cached)
        ):

            return cached

        path = download_model(
            model_name
        )

        loaded_models[
            model_name
        ] = path

        return path

# ============================================================
# COMFYUI PATH
# ============================================================

def setup_comfyui_path():

    if COMFYUI_ROOT in sys.path:

        try:

            sys.path.remove(
                COMFYUI_ROOT
            )

        except ValueError:
            pass

    sys.path.insert(
        0,
        COMFYUI_ROOT
    )

    log(
        "[COMFYUI] Root:",
        COMFYUI_ROOT
    )

    log(
        "[COMFYUI] sys.path[0]:",
        sys.path[0]
    )

setup_comfyui_path()

# ============================================================
# FIX UTILS COLLISION
# ============================================================

def fix_comfy_utils_namespace():

    utils_dir = os.path.join(
        COMFYUI_ROOT,
        "utils"
    )

    if not os.path.isdir(
        utils_dir
    ):

        return

    existing = sys.modules.get(
        "utils"
    )

    if existing is not None:

        existing_path = getattr(
            existing,
            "__path__",
            None
        )

        if existing_path is None:

            try:

                del sys.modules[
                    "utils"
                ]

            except KeyError:
                pass

    if COMFYUI_ROOT in sys.path:

        try:

            sys.path.remove(
                COMFYUI_ROOT
            )

        except ValueError:
            pass

    sys.path.insert(
        0,
        COMFYUI_ROOT
    )

    importlib.invalidate_caches()

    try:

        import utils

        log(
            "[COMFYUI] utils:",
            getattr(
                utils,
                "__file__",
                None
            )
        )

    except Exception as e:

        log(
            "[COMFYUI] utils warning:",
            repr(e)
        )

fix_comfy_utils_namespace()

# ============================================================
# COMFY VERSION
# ============================================================

try:

    import comfyui_version

    comfy_version = getattr(
        comfyui_version,
        "__version__",
        None
    )

    if comfy_version is None:

        comfy_version = getattr(
            comfyui_version,
            "VERSION",
            "unknown"
        )

except Exception:

    comfy_version = "unknown"

log(
    "[COMFYUI] Version:",
    comfy_version
)

# ============================================================
# EXTRA MODEL PATHS
# ============================================================

def add_extra_model_paths():

    config_path = os.path.join(
        COMFYUI_ROOT,
        "extra_model_paths.yaml"
    )

    if not os.path.isfile(
        config_path
    ):

        return

    try:

        from main import (
            load_extra_path_config
        )

        load_extra_path_config(
            config_path
        )

    except Exception as e:

        log(
            "[COMFYUI] Extra paths warning:",
            repr(e)
        )

add_extra_model_paths()

# ============================================================
# TORCH
# ============================================================

import torch

# ============================================================
# COMFY MODEL MANAGEMENT
# ============================================================

import comfy.model_management

from comfy.model_management import (
    CPUState
)

try:

    comfy.model_management.cpu_state = (
        CPUState.CPU
    )

except Exception as e:

    log(
        "[DEVICE] CPU state warning:",
        repr(e)
    )

log(
    "[DEVICE] CPU forced"
)

log(
    "[DEVICE] CUDA available:",
    torch.cuda.is_available()
)

# ============================================================
# COMFY NODES
# ============================================================

def import_custom_nodes():

    with COMFY_INIT_LOCK:

        log(
            "[COMFYUI] Initializing nodes..."
        )

        fix_comfy_utils_namespace()

        import execution

        from nodes import (
            init_extra_nodes
        )

        import server

        loop = asyncio.new_event_loop()

        try:

            asyncio.set_event_loop(
                loop
            )

            server_instance = (
                server.PromptServer(
                    loop
                )
            )

            if not hasattr(
                server_instance,
                "prompt_queue"
            ):

                try:

                    execution.PromptQueue(
                        server_instance
                    )

                except Exception as e:

                    log(
                        "[COMFYUI] PromptQueue warning:",
                        repr(e)
                    )

            result = init_extra_nodes()

            if asyncio.iscoroutine(
                result
            ):

                loop.run_until_complete(
                    result
                )

        finally:

            try:

                asyncio.set_event_loop(
                    None
                )

            except Exception:
                pass

            try:

                loop.close()

            except Exception:
                pass

        log(
            "[COMFYUI] Nodes initialized."
        )

import_custom_nodes()

# ============================================================
# NODE MAPPINGS
# ============================================================

from nodes import (
    NODE_CLASS_MAPPINGS
)

# ============================================================
# VERIFY REACTOR
# ============================================================

if not os.path.isdir(
    REACTOR_PATH
):

    raise RuntimeError(
        "Không tìm thấy ComfyUI-ReActor:\n"
        + REACTOR_PATH
    )

if REACTOR_PATH not in sys.path:

    sys.path.insert(
        0,
        REACTOR_PATH
    )

log(
    "[REACTOR] Path:",
    REACTOR_PATH
)

# ============================================================
# LOAD LOADIMAGE
# ============================================================

try:

    loadimage = (
        NODE_CLASS_MAPPINGS[
            "LoadImage"
        ]()
    )

except Exception as e:

    raise RuntimeError(
        "Không load được LoadImage:\n"
        + repr(e)
    )

# ============================================================
# LOAD REACTOR
# ============================================================

try:

    reactorfaceswap = (
        NODE_CLASS_MAPPINGS[
            "ReActorFaceSwap"
        ]()
    )

except Exception as e:

    raise RuntimeError(
        "Không load được ReActorFaceSwap:\n"
        + repr(e)
    )

# ============================================================
# REACTOR FACE FUNCTIONS
# ============================================================

try:

    from scripts.reactor_swapper import (
        analyze_faces,
        sort_by_order
    )

except Exception as e:

    raise RuntimeError(
        "Không import được reactor_swapper:\n"
        + repr(e)
    )

# ============================================================
# FILE HELPERS
# ============================================================

def get_file_path(file):

    if file is None:
        return None

    if isinstance(
        file,
        str
    ):

        return file

    if hasattr(
        file,
        "path"
    ):

        return file.path

    if hasattr(
        file,
        "name"
    ):

        return file.name

    return str(file)


def get_file_paths(files):

    if not files:
        return []

    if isinstance(
        files,
        (str, bytes)
    ):

        files = [
            files
        ]

    result = []

    for item in files:

        path = get_file_path(
            item
        )

        if (
            path
            and os.path.isfile(
                path
            )
        ):

            result.append(
                path
            )

    return result

# ============================================================
# FACE CACHE
# ============================================================

def clear_face_cache():

    face_cache.clear()

    gc.collect()


def detect_reactor_faces(
    image_path
):

    if not image_path:
        return []

    if not os.path.isfile(
        image_path
    ):

        return []

    try:

        mtime = os.path.getmtime(
            image_path
        )

    except Exception:

        return []

    key = (
        image_path,
        mtime
    )

    if key in face_cache:

        faces = face_cache.pop(
            key
        )

        face_cache[key] = faces

        return faces

    img = cv2.imread(
        image_path
    )

    if img is None:
        return []

    faces = analyze_faces(
        img,
        det_size=(
            640,
            640
        )
    )

    if faces is None:
        faces = []

    faces = sort_by_order(
        faces,
        "large-small"
    )

    face_cache[key] = faces

    while len(face_cache) > FACE_CACHE_MAX:

        face_cache.popitem(
            last=False
        )

    return faces

# ============================================================
# GIF FIRST FRAME
# ============================================================

def get_gif_first_frame(
    path
):

    reader = imageio.get_reader(
        path
    )

    try:

        frame = reader.get_data(
            0
        )

        return np.asarray(
            frame
        ).copy()

    finally:

        reader.close()

# ============================================================
# IMAGE FOR DETECTION
# ============================================================

def get_detection_image(
    path
):

    if path.lower().endswith(
        ".gif"
    ):

        return get_gif_first_frame(
            path
        )

    img = cv2.imread(
        path
    )

    if img is None:
        return None

    return cv2.cvtColor(
        img,
        cv2.COLOR_BGR2RGB
    )

# ============================================================
# FACE CROP
# ============================================================

def crop_face_from_image(
    rgb,
    face,
    padding=0.35
):

    if rgb is None:
        return None

    h, w = rgb.shape[:2]

    x1, y1, x2, y2 = (
        face.bbox
    )

    x1 = int(
        round(x1)
    )

    y1 = int(
        round(y1)
    )

    x2 = int(
        round(x2)
    )

    y2 = int(
        round(y2)
    )

    face_w = max(
        1,
        x2 - x1
    )

    face_h = max(
        1,
        y2 - y1
    )

    pad_x = int(
        face_w * padding
    )

    pad_y = int(
        face_h * padding
    )

    x1 = max(
        0,
        x1 - pad_x
    )

    y1 = max(
        0,
        y1 - pad_y
    )

    x2 = min(
        w,
        x2 + pad_x
    )

    y2 = min(
        h,
        y2 + pad_y
    )

    if x2 <= x1 or y2 <= y1:
        return None

    crop = rgb[
        y1:y2,
        x1:x2
    ]

    if crop.size == 0:
        return None

    return np.asarray(
        crop
    ).copy()

# ============================================================
# MAKE FACE REVIEW GALLERY
# ============================================================

def make_face_review_gallery(
    path,
    prefix
):

    if not path:
        return []

    rgb = get_detection_image(
        path
    )

    if rgb is None:
        return []

    faces = detect_reactor_faces(
        path
    )

    gallery = []

    for face_index, face in enumerate(
        faces
    ):

        crop = crop_face_from_image(
            rgb,
            face,
            padding=0.35
        )

        if crop is None:
            continue

        filename = os.path.basename(
            path
        )

        caption = (
            f"{prefix} FACE {face_index} | "
            f"{filename}"
        )

        gallery.append(
            (
                crop,
                caption
            )
        )

    return gallery

# ============================================================
# FACE REVIEW FOR MULTIPLE FILES
# ============================================================

def make_multi_face_review(
    paths,
    prefix
):

    gallery = []

    for index, path in enumerate(
        paths
    ):

        try:

            items = make_face_review_gallery(
                path,
                f"{prefix} {index}"
            )

            gallery.extend(
                items
            )

        except Exception as e:

            log(
                "[FACE REVIEW ERROR]",
                path,
                repr(e)
            )

    return gallery

# ============================================================
# ANNOTATED FACE PREVIEW
# ============================================================

def make_annotated_preview(
    path,
    prefix
):

    if not path:
        return None, 0

    rgb = get_detection_image(
        path
    )

    if rgb is None:
        return None, 0

    bgr = cv2.cvtColor(
        rgb,
        cv2.COLOR_RGB2BGR
    )

    h, w = bgr.shape[:2]

    faces = detect_reactor_faces(
        path
    )

    for face_index, face in enumerate(
        faces
    ):

        x1, y1, x2, y2 = (
            face.bbox
        )

        x1 = max(
            0,
            min(
                w - 1,
                int(x1)
            )
        )

        y1 = max(
            0,
            min(
                h - 1,
                int(y1)
            )
        )

        x2 = max(
            x1 + 1,
            min(
                w - 1,
                int(x2)
            )
        )

        y2 = max(
            y1 + 1,
            min(
                h - 1,
                int(y2)
            )
        )

        cv2.rectangle(
            bgr,
            (x1, y1),
            (x2, y2),
            (0, 255, 0),
            max(
                2,
                int(
                    min(w, h) / 350
                )
            )
        )

        label = (
            f"{prefix} FACE {face_index}"
        )

        font_scale = max(
            0.55,
            min(
                1.15,
                min(w, h) / 800
            )
        )

        thickness = max(
            1,
            int(
                font_scale * 2
            )
        )

        (
            tw,
            th
        ), baseline = cv2.getTextSize(
            label,
            cv2.FONT_HERSHEY_SIMPLEX,
            font_scale,
            thickness
        )

        label_y1 = max(
            0,
            y1 - th - baseline - 8
        )

        label_y2 = (
            y1
            if y1 > th + baseline + 8
            else y1 + th + baseline + 12
        )

        cv2.rectangle(
            bgr,
            (
                x1,
                label_y1
            ),
            (
                min(
                    w - 1,
                    x1 + tw + 12
                ),
                min(
                    h - 1,
                    label_y2
                )
            ),
            (0, 255, 0),
            -1
        )

        text_y = (
            label_y2 - 6
            if y1 > th + baseline + 8
            else label_y1 + th + 6
        )

        cv2.putText(
            bgr,
            label,
            (
                x1 + 6,
                text_y
            ),
            cv2.FONT_HERSHEY_SIMPLEX,
            font_scale,
            (0, 0, 0),
            thickness,
            cv2.LINE_AA
        )

    preview = cv2.cvtColor(
        bgr,
        cv2.COLOR_BGR2RGB
    )

    return preview, len(faces)

# ============================================================
# MULTI FILE PREVIEW
# ============================================================

def preview_source_target(
    source_files,
    target_files
):

    source_paths = get_file_paths(
        source_files
    )

    target_paths = get_file_paths(
        target_files
    )

    source_gallery = []
    target_gallery = []

    source_choices = []
    target_choices = []

    source_face_choices = []
    target_face_choices = []

    try:

        with JOB_SEMAPHORE:

            # =================================================
            # SOURCE
            # =================================================

            for source_index, path in enumerate(
                source_paths
            ):

                try:

                    faces = detect_reactor_faces(
                        path
                    )

                    count = len(
                        faces
                    )

                    crops = make_face_review_gallery(
                        path,
                        f"SOURCE {source_index}"
                    )

                    source_gallery.extend(
                        crops
                    )

                    source_choices.append(
                        f"SOURCE {source_index}: "
                        f"{os.path.basename(path)}"
                    )

                    for face_index in range(
                        count
                    ):

                        source_face_choices.append(
                            f"SOURCE {source_index} "
                            f"FACE {face_index}"
                        )

                except Exception as e:

                    log(
                        "[PREVIEW SOURCE ERROR]",
                        path,
                        repr(e)
                    )

            # =================================================
            # TARGET
            # =================================================

            for target_index, path in enumerate(
                target_paths
            ):

                try:

                    faces = detect_reactor_faces(
                        path
                    )

                    count = len(
                        faces
                    )

                    crops = make_face_review_gallery(
                        path,
                        f"TARGET {target_index}"
                    )

                    target_gallery.extend(
                        crops
                    )

                    target_choices.append(
                        f"TARGET {target_index}: "
                        f"{os.path.basename(path)}"
                    )

                    for face_index in range(
                        count
                    ):

                        target_face_choices.append(
                            f"TARGET {target_index} "
                            f"FACE {face_index}"
                        )

                except Exception as e:

                    log(
                        "[PREVIEW TARGET ERROR]",
                        path,
                        repr(e)
                    )

    except Exception as e:

        log(
            "[PREVIEW ERROR]",
            repr(e)
        )

    source_value = (
        source_choices[0]
        if source_choices
        else None
    )

    target_value = (
        target_choices[0]
        if target_choices
        else None
    )

    source_face_value = (
        source_face_choices[0]
        if source_face_choices
        else None
    )

    target_face_value = (
        target_face_choices[0]
        if target_face_choices
        else None
    )

    return (

        source_gallery,

        target_gallery,

        gr.update(
            choices=source_choices,
            value=source_value
        ),

        gr.update(
            choices=target_choices,
            value=target_value
        ),

        gr.update(
            choices=source_face_choices,
            value=source_face_value
        ),

        gr.update(
            choices=target_face_choices,
            value=target_face_value
        ),

        (
            f"Preview completed: "
            f"{len(source_paths)} source / "
            f"{len(target_paths)} target"
        )
    )

# ============================================================
# SOURCE FACE SELECTOR
# ============================================================

def source_face_choices_from_source(
    source_files,
    source_selector
):

    paths = get_file_paths(
        source_files
    )

    if not paths:

        return gr.update(
            choices=[],
            value=None
        )

    index = 0

    if source_selector:

        try:

            index = int(
                source_selector.split(
                    ":",
                    1
                )[0].replace(
                    "SOURCE",
                    ""
                ).strip()
            )

        except Exception:

            index = 0

    if index < 0 or index >= len(paths):
        index = 0

    faces = detect_reactor_faces(
        paths[index]
    )

    choices = [

        f"SOURCE {index} FACE {i}"

        for i in range(
            len(faces)
        )
    ]

    return gr.update(
        choices=choices,
        value=(
            choices[0]
            if choices
            else None
        )
    )

# ============================================================
# TARGET FACE SELECTOR
# ============================================================

def target_face_choices_from_target(
    target_files,
    target_selector
):

    paths = get_file_paths(
        target_files
    )

    if not paths:

        return gr.update(
            choices=[],
            value=None
        )

    index = 0

    if target_selector:

        try:

            index = int(
                target_selector.split(
                    ":",
                    1
                )[0].replace(
                    "TARGET",
                    ""
                ).strip()
            )

        except Exception:

            index = 0

    if index < 0 or index >= len(paths):
        index = 0

    faces = detect_reactor_faces(
        paths[index]
    )

    choices = [

        f"TARGET {index} FACE {i}"

        for i in range(
            len(faces)
        )
    ]

    return gr.update(
        choices=choices,
        value=(
            choices[0]
            if choices
            else None
        )
    )

# ============================================================
# PARSE SOURCE FACE
# ============================================================

def parse_source_face(
    value
):

    if not value:

        return (
            -1,
            -1
        )

    text = str(
        value
    ).upper()

    try:

        source_index = int(
            text.split(
                "SOURCE",
                1
            )[1].split(
                "FACE",
                1
            )[0].strip()
        )

        face_index = int(
            text.split(
                "FACE",
                1
            )[1].strip()
        )

        return (
            source_index,
            face_index
        )

    except Exception:

        return (
            -1,
            -1
        )

# ============================================================
# PARSE TARGET FACE
# ============================================================

def parse_target_face(
    value
):

    if not value:

        return (
            -1,
            -1
        )

    text = str(
        value
    ).upper()

    try:

        target_index = int(
            text.split(
                "TARGET",
                1
            )[1].split(
                "FACE",
                1
            )[0].strip()
        )

        face_index = int(
            text.split(
                "FACE",
                1
            )[1].strip()
        )

        return (
            target_index,
            face_index
        )

    except Exception:

        return (
            -1,
            -1
        )

# ============================================================
# MAPPING FORMAT
# ============================================================

def normalize_mapping_line(
    source_face,
    target_face
):

    source_index, source_face_index = (
        parse_source_face(
            source_face
        )
    )

    target_index, target_face_index = (
        parse_target_face(
            target_face
        )
    )

    if (
        source_index < 0
        or source_face_index < 0
        or target_index < 0
        or target_face_index < 0
    ):

        return None

    return (
        f"SOURCE {source_index} "
        f"FACE {source_face_index} "
        f"-> TARGET {target_index} "
        f"FACE {target_face_index}"
    )

# ============================================================
# PARSE MAPPINGS
# ============================================================

def parse_mappings(
    mapping_text
):

    mappings = []

    if not mapping_text:
        return mappings

    for raw in str(
        mapping_text
    ).splitlines():

        line = raw.strip()

        if not line:
            continue

        upper = line.upper()

        if (
            "SOURCE" not in upper
            or "TARGET" not in upper
            or "FACE" not in upper
        ):

            continue

        try:

            left, right = (
                upper.split(
                    "->",
                    1
                )
            )

            s_index = int(
                left.split(
                    "SOURCE",
                    1
                )[1].split(
                    "FACE",
                    1
                )[0].strip()
            )

            s_face = int(
                left.split(
                    "FACE",
                    1
                )[1].strip()
            )

            t_index = int(
                right.split(
                    "TARGET",
                    1
                )[1].split(
                    "FACE",
                    1
                )[0].strip()
            )

            t_face = int(
                right.split(
                    "FACE",
                    1
                )[1].strip()
            )

            mappings.append(
                {
                    "source_index": s_index,
                    "source_face": s_face,
                    "target_index": t_index,
                    "target_face": t_face
                }
            )

        except Exception:

            continue

    return mappings

# ============================================================
# DEFAULT MAPPING
# ============================================================

def default_mapping():

    return (
        "SOURCE 0 FACE 0 "
        "-> TARGET 0 FACE 0"
    )

# ============================================================
# ADD MAPPING
# ============================================================

def add_mapping(
    mapping_text,
    source_face,
    target_face
):

    line = normalize_mapping_line(
        source_face,
        target_face
    )

    if not line:

        raise gr.Error(
            "Mapping không hợp lệ."
        )

    existing = []

    if mapping_text:

        existing = [

            x.strip()

            for x in str(
                mapping_text
            ).splitlines()

            if x.strip()
        ]

    if line not in existing:

        existing.append(
            line
        )

    return "\n".join(
        existing
    )

# ============================================================
# REMOVE LAST
# ============================================================

def remove_last_mapping(
    mapping_text
):

    lines = [

        x.strip()

        for x in str(
            mapping_text or ""
        ).splitlines()

        if x.strip()
    ]

    if lines:
        lines.pop()

    return "\n".join(
        lines
    )

# ============================================================
# CLEAR MAPPING
# ============================================================

def clear_mapping():

    return ""

# ============================================================
# VALIDATE MAPPING
# ============================================================

def validate_mappings(
    mappings,
    source_paths,
    target_paths
):

    if not mappings:

        raise gr.Error(
            "Chưa có Face Mapping."
        )

    for mapping in mappings:

        si = mapping[
            "source_index"
        ]

        sf = mapping[
            "source_face"
        ]

        ti = mapping[
            "target_index"
        ]

        tf = mapping[
            "target_face"
        ]

        if (
            si < 0
            or si >= len(source_paths)
        ):

            raise gr.Error(
                f"Source index không tồn tại: {si}"
            )

        if (
            ti < 0
            or ti >= len(target_paths)
        ):

            raise gr.Error(
                f"Target index không tồn tại: {ti}"
            )

        source_faces = (
            detect_reactor_faces(
                source_paths[si]
            )
        )

        target_faces = (
            detect_reactor_faces(
                target_paths[ti]
            )
        )

        if (
            sf < 0
            or sf >= len(source_faces)
        ):

            raise gr.Error(
                f"SOURCE {si} FACE {sf} "
                "không tồn tại."
            )

        if (
            tf < 0
            or tf >= len(target_faces)
        ):

            raise gr.Error(
                f"TARGET {ti} FACE {tf} "
                "không tồn tại."
            )

# ============================================================
# LOAD IMAGE THROUGH COMFY
# ============================================================

def comfy_load_image(
    path
):

    result = loadimage.load_image(
        image=path
    )

    if result is None:

        raise RuntimeError(
            "ComfyUI LoadImage returned None."
        )

    if isinstance(
        result,
        dict
    ):

        values = list(
            result.values()
        )

        if not values:

            raise RuntimeError(
                "LoadImage returned empty dict."
            )

        image = values[0]

    else:

        try:

            image = result[0]

        except Exception:

            image = result

    if image is None:

        raise RuntimeError(
            "Loaded image is None."
        )

    return image

# ============================================================
# SOURCE CACHE
# ============================================================

def load_source_cached(
    path
):

    try:

        mtime = os.path.getmtime(
            path
        )

    except Exception:

        mtime = 0

    key = (
        path,
        mtime
    )

    if key in source_tensor_cache:

        value = (
            source_tensor_cache.pop(
                key
            )
        )

        source_tensor_cache[
            key
        ] = value

        return value

    image = comfy_load_image(
        path
    )

    source_tensor_cache[
        key
    ] = image

    while (
        len(source_tensor_cache)
        > SOURCE_TENSOR_CACHE_MAX
    ):

        source_tensor_cache.popitem(
            last=False
        )

    return image

# ============================================================
# GET VALUE
# ============================================================

def get_value_at_index(
    obj,
    index
):

    if obj is None:

        raise RuntimeError(
            "ComfyUI node returned None."
        )

    if isinstance(
        obj,
        dict
    ):

        values = list(
            obj.values()
        )

        if index >= len(values):

            raise RuntimeError(
                "ComfyUI output index out of range."
            )

        return values[index]

    return obj[index]

# ============================================================
# RESULT TO PIL
# ============================================================

def result_to_pil(
    result
):

    if result is None:

        raise RuntimeError(
            "ReActor output is None."
        )

    try:

        value = get_value_at_index(
            result,
            0
        )

    except Exception:

        value = result

    if value is None:

        raise RuntimeError(
            "ReActor image output is None."
        )

    if isinstance(
        value,
        (list, tuple)
    ):

        if not value:

            raise RuntimeError(
                "ReActor returned empty image."
            )

        value = value[0]

    if hasattr(
        value,
        "detach"
    ):

        value = (
            value
            .detach()
            .cpu()
            .float()
            .numpy()
        )

    value = np.asarray(
        value
    )

    if value.ndim == 4:

        value = value[0]

    if value.ndim != 3:

        raise RuntimeError(
            "Invalid ReActor output shape: "
            + str(value.shape)
        )

    if value.shape[-1] != 3:

        if value.shape[0] == 3:

            value = np.transpose(
                value,
                (1, 2, 0)
            )

        else:

            raise RuntimeError(
                "Invalid image channels: "
                + str(value.shape)
            )

    if value.max() <= 1.0:

        value = value * 255.0

    value = np.clip(
        value,
        0,
        255
    ).astype(
        np.uint8
    )

    return Image.fromarray(
        value
    ).convert(
        "RGB"
    )

# ============================================================
# PIL -> COMFY IMAGE TENSOR
# ============================================================

def pil_to_comfy_image(
    image
):

    if image is None:

        raise RuntimeError(
            "Cannot convert None image to ComfyUI IMAGE."
        )

    if isinstance(
        image,
        Image.Image
    ):

        pil = image.convert(
            "RGB"
        )

        array = np.asarray(
            pil,
            dtype=np.float32
        )

    else:

        array = np.asarray(
            image
        )

        if array.ndim == 2:

            array = np.stack(
                [
                    array,
                    array,
                    array
                ],
                axis=-1
            )

        if array.ndim == 3:

            if array.shape[-1] == 4:

                array = (
                    array[
                        :, :, :3
                    ]
                )

        elif array.ndim == 4:

            if array.shape[0] == 1:

                array = array[0]

        else:

            raise RuntimeError(
                "Invalid image shape for "
                "ComfyUI conversion: "
                + str(array.shape)
            )

        array = array.astype(
            np.float32,
            copy=False
        )

    if array.max() > 1.0:

        array /= 255.0

    array = np.clip(
        array,
        0.0,
        1.0
    )

    tensor = torch.from_numpy(
        array
    )

    if tensor.ndim == 3:

        tensor = tensor.unsqueeze(
            0
        )

    if tensor.ndim != 4:

        raise RuntimeError(
            "Invalid converted ComfyUI "
            "IMAGE tensor shape: "
            + str(tuple(tensor.shape))
        )

    if tensor.shape[-1] != 3:

        raise RuntimeError(
            "Invalid converted image channels: "
            + str(tuple(tensor.shape))
        )

    return tensor.contiguous()

# ============================================================
# NORMALIZE COMFY IMAGE
# ============================================================

def normalize_target_image(
    image
):

    if image is None:
        return None

    if torch.is_tensor(
        image
    ):

        tensor = image

        if tensor.ndim == 3:

            tensor = tensor.unsqueeze(
                0
            )

        if tensor.ndim != 4:

            raise RuntimeError(
                "Invalid ComfyUI IMAGE "
                "tensor shape: "
                + str(tuple(tensor.shape))
            )

        if tensor.shape[-1] != 3:

            raise RuntimeError(
                "Invalid ComfyUI IMAGE "
                "channels: "
                + str(tuple(tensor.shape))
            )

        if tensor.dtype != torch.float32:

            tensor = tensor.float()

        if tensor.numel() > 0:

            max_value = float(
                tensor.detach()
                .max()
                .cpu()
            )

            if max_value > 1.0:

                tensor = (
                    tensor / 255.0
                )

        return tensor.clamp(
            0.0,
            1.0
        ).contiguous()

    return pil_to_comfy_image(
        image
    )

# ============================================================
# REACTOR SWAP
# ============================================================

def reactor_swap(
    source_image,
    target_image,
    source_face_index,
    target_face_index,
    swap_model,
    restore_model,
    restore_strength
):

    source_image = normalize_target_image(
        source_image
    )

    target_image = normalize_target_image(
        target_image
    )

    if source_image is None:

        raise RuntimeError(
            "ReActor source_image is None."
        )

    if target_image is None:

        raise RuntimeError(
            "ReActor target_image is None."
        )

    swap_model_path = ensure_model(
        swap_model
    )

    if not swap_model_path:

        raise RuntimeError(
            "Swap model unavailable."
        )

    restore_name = (
        restore_model
        if restore_model
        else "none"
    )

    if restore_name != "none":

        restore_path = ensure_model(
            restore_name
        )

        log(
            "[RESTORE] Model path:",
            restore_path
        )

        if not restore_path:

            log(
                "[RESTORE] Model unavailable. "
                "Disabling restore."
            )

            restore_name = "none"

    kwargs = {

        "enabled": True,

        "swap_model": swap_model,

        "facedetection": (
            "retinaface_resnet50"
        ),

        "face_restore_model": (
            restore_name
        ),

        "face_restore_visibility": (
            float(
                restore_strength
            )
        ),

        "codeformer_weight": 0.5,

        "detect_gender_input": "no",

        "detect_gender_source": "no",

        "input_faces_index": str(
            target_face_index
        ),

        "source_faces_index": str(
            source_face_index
        ),

        "console_log_level": 1,

        "input_image": target_image,

        "source_image": source_image
    }

    try:

        result = reactorfaceswap.execute(
            **kwargs
        )

        return result_to_pil(
            result
        )

    except TypeError as e:

        message = str(e)

        if (
            restore_name != "none"
            and (
                "NoneType" in message
                or "Unable to load" in message
                or "No such file" in message
                or "not found" in message.lower()
            )
        ):

            log(
                "[RESTORE FALLBACK]",
                message
            )

            kwargs[
                "face_restore_model"
            ] = "none"

            kwargs[
                "face_restore_visibility"
            ] = 0.0

            result = reactorfaceswap.execute(
                **kwargs
            )

            return result_to_pil(
                result
            )

        raise

    except Exception as e:

        message = str(e)

        if (
            restore_name != "none"
            and (
                "NoneType" in message
                or "Unable to load" in message
                or "No such file" in message
                or "not found" in message.lower()
            )
        ):

            log(
                "[RESTORE FALLBACK]",
                message
            )

            kwargs[
                "face_restore_model"
            ] = "none"

            kwargs[
                "face_restore_visibility"
            ] = 0.0

            result = reactorfaceswap.execute(
                **kwargs
            )

            return result_to_pil(
                result
            )

        raise

# ============================================================
# SWAP ONE IMAGE
# ============================================================

def swap_one_image(
    source_path,
    target_path,
    source_face_index,
    target_face_index,
    swap_model,
    restore_model,
    restore_strength,
    target_image_override=None
):

    source_tensor = load_source_cached(
        source_path
    )

    source_tensor = normalize_target_image(
        source_tensor
    )

    if target_image_override is None:

        target_tensor = comfy_load_image(
            target_path
        )

    else:

        target_tensor = normalize_target_image(
            target_image_override
        )

    if source_tensor is None:

        raise RuntimeError(
            "Source tensor is None."
        )

    if target_tensor is None:

        raise RuntimeError(
            "Target tensor is None."
        )

    log(
        "[SWAP] Source type:",
        type(source_tensor),
        "shape:",
        getattr(
            source_tensor,
            "shape",
            None
        )
    )

    log(
        "[SWAP] Target type:",
        type(target_tensor),
        "shape:",
        getattr(
            target_tensor,
            "shape",
            None
        )
    )

    with INFERENCE_SEMAPHORE:

        result_image = reactor_swap(

            source_tensor,

            target_tensor,

            source_face_index,

            target_face_index,

            swap_model,

            restore_model,

            restore_strength
        )

    if result_image is None:

        raise RuntimeError(
            "Swap returned None."
        )

    return result_image

# ============================================================
# GIF PROCESS
# ============================================================

def process_gif_mapping(
    source_path,
    target_path,
    source_face_index,
    target_face_index,
    swap_model,
    restore_model,
    restore_strength
):

    reader = imageio.get_reader(
        target_path
    )

    frames = []
    durations = []

    try:

        try:

            meta = reader.get_meta_data()

        except Exception:

            meta = {}

        duration = meta.get(
            "duration",
            100
        )

        fps = meta.get(
            "fps",
            None
        )

        if fps and not duration:

            duration = (
                1000.0 /
                float(fps)
            )

        for frame_index, frame in enumerate(
            reader
        ):

            frame_pil = (
                Image.fromarray(
                    frame
                ).convert(
                    "RGB"
                )
            )

            with tempfile.NamedTemporaryFile(
                suffix=".png",
                delete=False
            ) as tmp:

                temp_path = tmp.name

            try:

                frame_pil.save(
                    temp_path
                )

                result = swap_one_image(

                    source_path,

                    temp_path,

                    source_face_index,

                    target_face_index,

                    swap_model,

                    restore_model,

                    restore_strength
                )

                frames.append(
                    result
                )

                durations.append(
                    duration
                )

            finally:

                try:

                    os.remove(
                        temp_path
                    )

                except Exception:
                    pass

            if (
                frame_index % 5
                == 0
            ):

                gc.collect()

    finally:

        reader.close()

    if not frames:

        raise gr.Error(
            "GIF không có frame."
        )

    return (
        frames,
        durations
    )

# ============================================================
# OUTPUT PATH
# ============================================================

def safe_name(
    value
):

    result = []

    for char in str(
        value
    ):

        if (
            char.isalnum()
            or char in (
                "-",
                "_",
                "."
            )
        ):

            result.append(
                char
            )

        else:

            result.append(
                "_"
            )

    return "".join(
        result
    )


def make_output_path(
    source_path,
    target_path,
    batch_index
):

    source_name = safe_name(
        os.path.splitext(
            os.path.basename(
                source_path
            )
        )[0]
    )

    target_name = safe_name(
        os.path.splitext(
            os.path.basename(
                target_path
            )
        )[0]
    )

    filename = (
        f"{source_name}"
        f"_TO_"
        f"{target_name}"
        f"_{batch_index:04d}.webp"
    )

    return os.path.join(
        OUTPUT_DIR,
        filename
    )

# ============================================================
# GENERATE
# ============================================================

def generate_image(
    source_files,
    target_files,
    mapping_text,
    swap_model,
    restore_model,
    restore_strength
):

    source_paths = get_file_paths(
        source_files
    )

    target_paths = get_file_paths(
        target_files
    )

    if not source_paths:

        raise gr.Error(
            "Chưa upload Source."
        )

    if not target_paths:

        raise gr.Error(
            "Chưa upload Target."
        )

    mappings = parse_mappings(
        mapping_text
    )

    if not mappings:

        mappings = [
            {
                "source_index": 0,
                "source_face": 0,
                "target_index": 0,
                "target_face": 0
            }
        ]

    if not swap_model:

        swap_model = (
            DEFAULT_SWAP_MODEL
        )

    if not restore_model:

        restore_model = (
            DEFAULT_RESTORE_MODEL
        )

    restore_strength = float(
        restore_strength
        if restore_strength is not None
        else DEFAULT_RESTORE_STRENGTH
    )

    log("=" * 70)

    log(
        "[GENERATE] START"
    )

    log(
        "[GENERATE] Sources:",
        len(source_paths)
    )

    log(
        "[GENERATE] Targets:",
        len(target_paths)
    )

    log(
        "[GENERATE] Mappings:",
        len(mappings)
    )

    log(
        "[GENERATE] Swap:",
        swap_model
    )

    log(
        "[GENERATE] Restore:",
        restore_model
    )

    log(
        "[GENERATE] Strength:",
        restore_strength
    )

    log("=" * 70)

    with JOB_SEMAPHORE:

        validate_mappings(
            mappings,
            source_paths,
            target_paths
        )

        ensure_model(
            swap_model
        )

        if (
            restore_model
            and restore_model != "none"
        ):

            ensure_model(
                restore_model
            )

        output_paths = []

        # ====================================================
        # GROUP MAPPING BY TARGET
        # ====================================================

        mappings_by_target = {}

        for mapping in mappings:

            target_index = mapping[
                "target_index"
            ]

            mappings_by_target.setdefault(
                target_index,
                []
            ).append(
                mapping
            )

        # ====================================================
        # PROCESS EACH TARGET
        # ====================================================

        for target_index, target_path in enumerate(
            target_paths
        ):

            target_mappings = (
                mappings_by_target.get(
                    target_index,
                    []
                )
            )

            if not target_mappings:
                continue

            log(
                "[TARGET]",
                target_index,
                target_path
            )

            # =================================================
            # GIF
            # =================================================

            if target_path.lower().endswith(
                ".gif"
            ):

                reader = imageio.get_reader(
                    target_path
                )

                frames = []
                durations = []

                try:

                    try:

                        meta = reader.get_meta_data()

                    except Exception:

                        meta = {}

                    duration = meta.get(
                        "duration",
                        100
                    )

                    for frame_index, frame in enumerate(
                        reader
                    ):

                        current = (
                            Image.fromarray(
                                frame
                            ).convert(
                                "RGB"
                            )
                        )

                        for mapping in target_mappings:

                            source_index = mapping[
                                "source_index"
                            ]

                            source_face = mapping[
                                "source_face"
                            ]

                            target_face = mapping[
                                "target_face"
                            ]

                            source_path = (
                                source_paths[
                                    source_index
                                ]
                            )

                            with tempfile.NamedTemporaryFile(
                                suffix=".png",
                                delete=False
                            ) as tmp:

                                temp_path = tmp.name

                            try:

                                current.save(
                                    temp_path
                                )

                                current = (
                                    swap_one_image(

                                        source_path,

                                        temp_path,

                                        source_face,

                                        target_face,

                                        swap_model,

                                        restore_model,

                                        restore_strength,

                                        target_image_override=(
                                            current
                                        )
                                    )
                                )

                            finally:

                                try:

                                    os.remove(
                                        temp_path
                                    )

                                except Exception:
                                    pass

                        frames.append(
                            current
                        )

                        durations.append(
                            duration
                        )

                        if (
                            frame_index % 3
                            == 0
                        ):

                            gc.collect()

                finally:

                    reader.close()

                if not frames:

                    raise gr.Error(
                        "GIF không có frame."
                    )

                output_path = (
                    make_output_path(

                        source_paths[
                            target_mappings[0][
                                "source_index"
                            ]
                        ],

                        target_path,

                        target_index
                    )
                )

                frames[0].save(

                    output_path,

                    save_all=True,

                    append_images=frames[1:],

                    duration=durations,

                    loop=0,

                    format="WEBP",

                    quality=90,

                    method=6
                )

                output_paths.append(
                    output_path
                )

                del frames
                del durations

                gc.collect()

            # =================================================
            # NORMAL IMAGE
            # =================================================

            else:

                current_path = target_path

                current_image = None

                first_source_path = None

                for mapping_index, mapping in enumerate(
                    target_mappings
                ):

                    source_index = mapping[
                        "source_index"
                    ]

                    source_face = mapping[
                        "source_face"
                    ]

                    target_face = mapping[
                        "target_face"
                    ]

                    source_path = (
                        source_paths[
                            source_index
                        ]
                    )

                    if first_source_path is None:

                        first_source_path = (
                            source_path
                        )

                    log(
                        "[MAPPING]",
                        (
                            f"SOURCE {source_index} "
                            f"FACE {source_face} "
                            f"-> TARGET {target_index} "
                            f"FACE {target_face}"
                        )
                    )

                    current_image = (
                        swap_one_image(

                            source_path,

                            current_path,

                            source_face,

                            target_face,

                            swap_model,

                            restore_model,

                            restore_strength,

                            target_image_override=(
                                current_image
                            )
                        )
                    )

                    if current_image is None:

                        raise RuntimeError(
                            "Mapping output is None."
                        )

                    current_path = None

                if current_image is None:

                    raise RuntimeError(
                        "No output generated."
                    )

                output_path = (
                    make_output_path(

                        first_source_path,

                        target_path,

                        target_index
                    )
                )

                current_image.save(

                    output_path,

                    format="WEBP",

                    quality=90,

                    method=6
                )

                output_paths.append(
                    output_path
                )

                del current_image

                gc.collect()

    gc.collect()

    log(
        "[GENERATE] OUTPUT:",
        output_paths
    )

    return output_paths

# ============================================================
# OUTPUT -> TARGET
# ============================================================

def use_output_as_target(
    output_files
):

    paths = get_file_paths(
        output_files
    )

    if not paths:

        raise gr.Error(
            "Chưa có output."
        )

    valid = [

        p

        for p in paths

        if os.path.isfile(p)
    ]

    if not valid:

        raise gr.Error(
            "Không tìm thấy output."
        )

    return valid

# ============================================================
# CLEAR TARGET
# ============================================================

def clear_target():

    clear_face_cache()

    return (

        [],

        [],

        gr.update(
            choices=[],
            value=None
        ),

        gr.update(
            choices=[],
            value=None
        ),

        gr.update(
            choices=[],
            value=None
        ),

        gr.update(
            choices=[],
            value=None
        ),

        "Target cleared."
    )

# ============================================================
# CLEAR SOURCE
# ============================================================

def clear_source():

    clear_face_cache()

    source_tensor_cache.clear()

    return (

        [],

        [],

        gr.update(
            choices=[],
            value=None
        ),

        gr.update(
            choices=[],
            value=None
        ),

        gr.update(
            choices=[],
            value=None
        ),

        gr.update(
            choices=[],
            value=None
        ),

        "Source cleared."
    )

# ============================================================
# SAVE MAPPING JSON
# ============================================================

def save_mapping_json(
    mapping_text
):

    mappings = parse_mappings(
        mapping_text
    )

    if not mappings:

        raise gr.Error(
            "Mapping trống."
        )

    fd, path = tempfile.mkstemp(
        suffix=".json",
        prefix="face_mapping_"
    )

    os.close(fd)

    with open(
        path,
        "w",
        encoding="utf-8"
    ) as f:

        json.dump(

            mappings,

            f,

            ensure_ascii=False,

            indent=2
        )

    return path

# ============================================================
# LOAD MAPPING JSON
# ============================================================

def load_mapping_json(
    mapping_file
):

    path = get_file_path(
        mapping_file
    )

    if not path:

        raise gr.Error(
            "Chưa chọn mapping JSON."
        )

    with open(
        path,
        "r",
        encoding="utf-8"
    ) as f:

        data = json.load(f)

    lines = []

    for item in data:

        try:

            lines.append(

                f"SOURCE "
                f"{int(item['source_index'])} "
                f"FACE "
                f"{int(item['source_face'])} "
                f"-> TARGET "
                f"{int(item['target_index'])} "
                f"FACE "
                f"{int(item['target_face'])}"

            )

        except Exception:

            continue

    return "\n".join(
        lines
    )

# ============================================================
# AUTO DEFAULT MAPPING
# ============================================================

def auto_mapping(source_files, target_files):
    source_paths = get_file_paths(source_files)
    target_paths = get_file_paths(target_files)

    if not source_paths:
        raise gr.Error("Chưa có source.")

    if not target_paths:
        raise gr.Error("Chưa có target.")

    source_index = 0
    source_face_index = 0

    # Kiểm tra source đầu tiên có face
    source_faces = detect_reactor_faces(source_paths[source_index])

    if not source_faces:
        raise gr.Error("Source đầu tiên không có face.")

    if source_face_index >= len(source_faces):
        raise gr.Error(
            f"Source đầu tiên chỉ có {len(source_faces)} face."
        )

    lines = []
    skipped_targets = []

    # Duyệt toàn bộ target
    for target_index, target_path in enumerate(target_paths):
        try:
            target_faces = detect_reactor_faces(target_path)
        except Exception:
            target_faces = []

        # Target không có face thì bỏ qua
        if not target_faces:
            skipped_targets.append(target_index)
            continue

        # Dùng SOURCE 0 FACE 0 cho tất cả face của target
        for target_face_index in range(len(target_faces)):
            lines.append(
                f"SOURCE {source_index} "
                f"FACE {source_face_index} "
                f"-> TARGET {target_index} "
                f"FACE {target_face_index}"
            )

    if not lines:
        raise gr.Error(
            "Không có target nào phát hiện được face."
        )

    # Thông báo các target bị bỏ qua
    if skipped_targets:
        lines.append("")
        lines.append(
            "# BỎ QUA TARGET KHÔNG CÓ FACE: "
            + ", ".join(map(str, skipped_targets))
        )

    return "\n".join(lines)

# ============================================================
# STATUS
# ============================================================

def status_text(
    source_files,
    target_files,
    mapping_text
):

    sources = len(
        get_file_paths(
            source_files
        )
    )

    targets = len(
        get_file_paths(
            target_files
        )
    )

    mappings = len(
        parse_mappings(
            mapping_text
        )
    )

    return (

        f"Sources: {sources} | "

        f"Targets: {targets} | "

        f"Mappings: {mappings} | "

        f"Workers: {MAX_CONCURRENT_JOBS}"
    )

# ============================================================
# OUTPUT PREVIEW VISIBILITY
# ============================================================

def toggle_output_preview(
    is_visible
):

    new_visible = not bool(
        is_visible
    )

    return (

        gr.update(
            visible=new_visible
        ),

        new_visible,

        (
            "Hide Output Preview"
            if new_visible
            else "Show Output Preview"
        )
    )

# ============================================================
# UI
# ============================================================

with gr.Blocks(
    title=APP_TITLE
) as app:

    # ========================================================
    # SOURCE
    # ========================================================

    with gr.Row():

        source_files = gr.File(

            label="Source Images",

            file_count="multiple",

            file_types=[
                ".jpg",
                ".jpeg",
                ".png",
                ".webp"
            ],

            type="filepath"
        )

        target_files = gr.File(

            label="Target Images / GIF",

            file_count="multiple",

            file_types=[
                ".jpg",
                ".jpeg",
                ".png",
                ".webp",
                ".gif"
            ],

            type="filepath"
        )

    # ========================================================
    # PREVIEW CONTROLS
    # ========================================================

    with gr.Row():

        preview_button = gr.Button(
            "Preview All Faces",
            variant="primary"
        )

        clear_source_button = gr.Button(
            "Clear Source"
        )

        clear_target_button = gr.Button(
            "Clear Target"
        )

    # ========================================================
    # SOURCE FACE REVIEW
    # ========================================================

    source_gallery = gr.Gallery(

        label="Source Face Review — Face Crops",

        columns=4,

        rows=2,

        height="auto",

        object_fit="contain",

        preview=True
    )

    # ========================================================
    # TARGET FACE REVIEW
    # ========================================================

    target_gallery = gr.Gallery(

        label="Target Face Review — Face Crops",

        columns=4,

        rows=2,

        height="auto",

        object_fit="contain",

        preview=True
    )

    # ========================================================
    # SELECTORS
    # ========================================================

    with gr.Row():

        source_selector = gr.Dropdown(

            label="Source Image",

            choices=[],

            value=None,

            interactive=True
        )

        target_selector = gr.Dropdown(

            label="Target Image",

            choices=[],

            value=None,

            interactive=True
        )

    with gr.Row():

        source_face_selector = gr.Dropdown(

            label="Source Face",

            choices=[],

            value=None,

            interactive=True
        )

        target_face_selector = gr.Dropdown(

            label="Target Face",

            choices=[],

            value=None,

            interactive=True
        )

    mapping_source_selector = gr.Dropdown(

        label="Mapping Source Face",

        choices=[],

        value=None,

        interactive=True
    )

    mapping_target_selector = gr.Dropdown(

        label="Mapping Target Face",

        choices=[],

        value=None,

        interactive=True
    )

    # ========================================================
    # MAPPING
    # ========================================================

    with gr.Row():

        add_mapping_button = gr.Button(
            "Add Mapping",
            variant="primary"
        )

        remove_mapping_button = gr.Button(
            "Remove Last"
        )

        clear_mapping_button = gr.Button(
            "Clear Mapping"
        )

        default_mapping_button = gr.Button(
            "Default Mapping"
        )

        auto_mapping_button = gr.Button(
            "Auto Mapping"
        )

    mapping_text = gr.Textbox(

        label="Face Mapping",

        value=default_mapping(),

        lines=8,

        max_lines=30,

        interactive=True
    )

    # ========================================================
    # MAPPING JSON
    # ========================================================

    with gr.Row():

        mapping_json_upload = gr.File(

            label="Load Mapping JSON",

            file_count="single",

            file_types=[
                ".json"
            ],

            type="filepath"
        )

        save_mapping_button = gr.Button(
            "Save Mapping JSON"
        )

        load_mapping_button = gr.Button(
            "Load Mapping JSON"
        )

    # ========================================================
    # MODELS
    # ========================================================

    with gr.Row():

        swap_model = gr.Dropdown(

            label="Swap Model",

            choices=[

                "inswapper_128.onnx",

                "hyperswap_1a_256.onnx",

                "hyperswap_1b_256.onnx",

                "hyperswap_1c_256.onnx"
            ],

            value=DEFAULT_SWAP_MODEL,

            interactive=True
        )

        restore_model = gr.Dropdown(

            label="Face Restore",

            choices=[

                "none",

                "GPEN-BFR-512.onnx"
            ],

            value=DEFAULT_RESTORE_MODEL,

            interactive=True
        )

    restore_strength = gr.Slider(

        label="Restore Strength",

        minimum=0.0,

        maximum=1.0,

        value=DEFAULT_RESTORE_STRENGTH,

        step=0.05,

        interactive=True
    )

    # ========================================================
    # GENERATE
    # ========================================================

    generate_button = gr.Button(

        "GENERATE MULTI FACE SWAP",

        variant="primary",

        size="lg"
    )

    # ========================================================
    # STATUS
    # ========================================================

    status = gr.Textbox(

        label="Status",

        value="Ready.",

        interactive=False
    )

    # ========================================================
    # OUTPUT
    # ========================================================

    output_files = gr.File(

        label="Generated Results",

        file_count="multiple",

        interactive=False
    )

    # ========================================================
    # OUTPUT PREVIEW STATE
    #
    # IMPORTANT:
    # DEFAULT = False
    # ========================================================

    output_preview_visible = gr.State(
        False
    )

    # ========================================================
    # OUTPUT PREVIEW TOGGLE BUTTON
    # ========================================================

    toggle_output_preview_button = gr.Button(
        "Show Output Preview"
    )

    # ========================================================
    # OUTPUT PREVIEW
    #
    # IMPORTANT:
    # DEFAULT HIDDEN
    # ========================================================

    output_gallery = gr.Gallery(

        label="Output Preview",

        columns=3,

        rows=2,

        height="auto",

        object_fit="contain",

        preview=True,

        visible=False
    )

    # ========================================================
    # OUTPUT -> TARGET
    # ========================================================

    use_output_button = gr.Button(

        "Use Output As Target + Preview"
    )

    # ========================================================
    # PREVIEW EVENT
    # ========================================================

    preview_button.click(

        fn=preview_source_target,

        inputs=[

            source_files,

            target_files
        ],

        outputs=[

            source_gallery,

            target_gallery,

            source_selector,

            target_selector,

            mapping_source_selector,

            mapping_target_selector,

            status
        ],

        concurrency_limit=2
    )

    # ========================================================
    # SOURCE SELECTOR
    # ========================================================

    source_selector.change(

        fn=source_face_choices_from_source,

        inputs=[

            source_files,

            source_selector
        ],

        outputs=[

            source_face_selector
        ],

        concurrency_limit=2
    )

    # ========================================================
    # TARGET SELECTOR
    # ========================================================

    target_selector.change(

        fn=target_face_choices_from_target,

        inputs=[

            target_files,

            target_selector
        ],

        outputs=[

            target_face_selector
        ],

        concurrency_limit=2
    )

    # ========================================================
    # ADD MAPPING
    # ========================================================

    add_mapping_button.click(

        fn=add_mapping,

        inputs=[

            mapping_text,

            mapping_source_selector,

            mapping_target_selector
        ],

        outputs=[

            mapping_text
        ],

        concurrency_limit=2
    )

    # ========================================================
    # REMOVE
    # ========================================================

    remove_mapping_button.click(

        fn=remove_last_mapping,

        inputs=[

            mapping_text
        ],

        outputs=[

            mapping_text
        ],

        concurrency_limit=2
    )

    # ========================================================
    # CLEAR
    # ========================================================

    clear_mapping_button.click(

        fn=clear_mapping,

        inputs=[],

        outputs=[

            mapping_text
        ],

        concurrency_limit=2
    )

    # ========================================================
    # DEFAULT
    # ========================================================

    default_mapping_button.click(

        fn=default_mapping,

        inputs=[],

        outputs=[

            mapping_text
        ],

        concurrency_limit=2
    )

    # ========================================================
    # AUTO MAPPING
    # ========================================================

    auto_mapping_button.click(

        fn=auto_mapping,

        inputs=[

            source_files,

            target_files
        ],

        outputs=[

            mapping_text
        ],

        concurrency_limit=2
    )

    # ========================================================
    # SAVE MAPPING
    # ========================================================

    save_mapping_button.click(

        fn=save_mapping_json,

        inputs=[

            mapping_text
        ],

        outputs=[

            mapping_json_upload
        ],

        concurrency_limit=2
    )

    # ========================================================
    # LOAD MAPPING
    # ========================================================

    load_mapping_button.click(

        fn=load_mapping_json,

        inputs=[

            mapping_json_upload
        ],

        outputs=[

            mapping_text
        ],

        concurrency_limit=2
    )

    # ========================================================
    # GENERATE
    #
    # Output Preview stays hidden after generation.
    # Generated files are still populated normally.
    # ========================================================

    generate_button.click(

        fn=generate_image,

        inputs=[

            source_files,

            target_files,

            mapping_text,

            swap_model,

            restore_model,

            restore_strength
        ],

        outputs=[

            output_files
        ],

        concurrency_limit=2

    ).then(

        fn=lambda files: files,

        inputs=[

            output_files
        ],

        outputs=[

            output_gallery
        ],

        concurrency_limit=2
    )

    # ========================================================
    # TOGGLE OUTPUT PREVIEW
    # ========================================================

    toggle_output_preview_button.click(

        fn=toggle_output_preview,

        inputs=[

            output_preview_visible
        ],

        outputs=[

            output_gallery,

            output_preview_visible,

            toggle_output_preview_button
        ],

        concurrency_limit=2
    )

    # ========================================================
    # OUTPUT -> TARGET
    # ========================================================

    use_output_button.click(

        fn=use_output_as_target,

        inputs=[

            output_files
        ],

        outputs=[

            target_files
        ],

        concurrency_limit=2

    ).then(

        fn=preview_source_target,

        inputs=[

            source_files,

            target_files
        ],

        outputs=[

            source_gallery,

            target_gallery,

            source_selector,

            target_selector,

            mapping_source_selector,

            mapping_target_selector,

            status
        ],

        concurrency_limit=2
    )

    # ========================================================
    # CLEAR SOURCE
    # ========================================================

    clear_source_button.click(

        fn=clear_source,

        inputs=[],

        outputs=[

            source_gallery,

            target_gallery,

            source_selector,

            target_selector,

            mapping_source_selector,

            mapping_target_selector,

            status
        ],

        concurrency_limit=2
    )

    # ========================================================
    # CLEAR TARGET
    # ========================================================

    clear_target_button.click(

        fn=clear_target,

        inputs=[],

        outputs=[

            source_gallery,

            target_gallery,

            source_selector,

            target_selector,

            mapping_source_selector,

            mapping_target_selector,

            status
        ],

        concurrency_limit=2
    )

# ============================================================
# GRADIO QUEUE
# ============================================================

app.queue(

    max_size=GRADIO_QUEUE_SIZE,

    default_concurrency_limit=MAX_CONCURRENT_JOBS
)

# ============================================================
# START
# ============================================================

log("=" * 70)

log(
    "[APP] Starting:",
    APP_TITLE
)

log(
    "[APP] Device: CPU"
)

log(
    "[APP] Swap:",
    DEFAULT_SWAP_MODEL
)

log(
    "[APP] Restore:",
    DEFAULT_RESTORE_MODEL
)

log(
    "[APP] RAM profile: 16 GB"
)

log(
    "[APP] Concurrent jobs:",
    MAX_CONCURRENT_JOBS
)

log(
    "[APP] Concurrent inference:",
    MAX_CONCURRENT_INFERENCE
)

log(
    "[APP] Multi Source Face: ENABLED"
)

log(
    "[APP] Multi Target Face: ENABLED"
)

log(
    "[APP] Face Mapping: ENABLED"
)

log(
    "[APP] GIF: ENABLED"
)

log(
    "[APP] Face Crop Review: ENABLED"
)

log(
    "[APP] Source Face Crop: ENABLED"
)

log(
    "[APP] Target Face Crop: ENABLED"
)

log(
    "[APP] Preview Batch Mode: ENABLED"
)

log(
    "[APP] Multi-Mapping Tensor Adapter: ENABLED"
)

log(
    "[APP] GPEN Path Fix: ENABLED"
)

log(
    "[APP] Output Preview Toggle: ENABLED"
)

log(
    "[APP] Output Preview Default: HIDDEN"
)

log("=" * 70)

# ============================================================
# LAUNCH
# ============================================================

if __name__ == "__main__":

    app.launch(

        server_name="0.0.0.0",

        server_port=int(
            os.environ.get(
                "PORT",
                "7860"
            )
        ),

        share=True,

        show_error=True
    )