import os
import sys
import cv2
import numpy as np
from PIL import Image

import gradio as gr
import gradio_client.utils as gc_utils
import insightface
from insightface.app import FaceAnalysis

# =========================================================
# Patch for gradio_client schema bugs (schema can be bool)
# =========================================================
_orig_get_type = gc_utils.get_type
_orig_json_schema_to_python_type = gc_utils.json_schema_to_python_type


def _safe_get_type(schema):
    if not isinstance(schema, dict):
        return "Any"
    return _orig_get_type(schema)


def _safe_json_schema_to_python_type(schema):
    if not isinstance(schema, dict):
        return "Any"
    try:
        return _orig_json_schema_to_python_type(schema)
    except Exception:
        return "Any"


gc_utils.get_type = _safe_get_type
gc_utils.json_schema_to_python_type = _safe_json_schema_to_python_type

# =========================================================
# Paths (OFFLINE)
# =========================================================
APP_DIR = os.path.dirname(os.path.abspath(__file__))
MODELS_DIR = os.path.join(APP_DIR, "models")

INSWAPPER_PATH = os.path.join(MODELS_DIR, "inswapper_128.onnx")
GFPGAN_PATH = os.path.join(MODELS_DIR, "GFPGANv1.3.pth")


def require_file(path: str, hint: str = ""):
    if os.path.exists(path):
        print(f"[OK] Found: {path}")
        return
    msg = f"Required file not found: {path}\n"
    if hint:
        msg += hint + "\n"
    raise RuntimeError(msg)


require_file(
    INSWAPPER_PATH,
    hint="Put inswapper_128.onnx into: models/inswapper_128.onnx",
)
require_file(
    GFPGAN_PATH,
    hint="Put GFPGANv1.3.pth into: models/GFPGANv1.3.pth",
)

from gfpgan import GFPGANer  # after files check

# =========================================================
# Localization texts
# =========================================================
TEXTS = {
    "ru": {
        "lang_radio_label": "Язык / Language",
        "title_md": (
            "# FaceSwap Pro (Docker)\n\n"
            "Слева — фото **донора** (может быть несколько лиц), справа — фото, где нужно заменить лица.  \n"
            "Внизу — результат, круговой индикатор прогресса, оценка времени и блок «До / После».\n"
        ),
        "step1_title": "### 1. Фото донора",
        "step1_input_label": "Загрузите фото донора (может быть несколько лиц)",
        "step1_donor_choice_label": "Выберите лицо‑донора",
        "step2_title": "### 2. Фото, которое меняем",
        "step2_input_label": "Загрузите изображение для замены лиц",
        "step2_target_choices_label": "Выберите лица для замены (если ничего не выбрано — меняем все)",
        "step3_title": "### 3. Настройки и экспорт",
        "use_enh_label": "Улучшить качество результата (GFPGAN)",
        "eta_initial": "Оценка времени появится после обнаружения лиц.",
        "fmt_label": "Формат файла для скачивания",
        "run_btn": "3. Запустить замену",
        "download_btn": "Скачать результат",
        "before_after_label": "До / После",
        "eta_fmt_sec": "Оценочное время обработки: ~{sec} сек.",
        "eta_fmt_min": "Оценочное время обработки: ~{min} мин {sec} сек.",
        "msg_need_images": "Загрузите оба изображения, чтобы начать обработку.",
        "msg_no_donor_faces": "На фото донора не найдено ни одного лица.",
        "msg_no_target_faces": "На целевом фото не найдено ни одного лица.",
        "msg_prep": "Подготовка к обработке...",
        "msg_swap_step": "Замена лица {i} из {n}",
        "msg_enh": "Улучшение качества (GFPGAN)...",
        "msg_done": "Готово. Обработано лиц: {n}.",
        "progress_done": "Готово!",
        "donor_option": "Донор {i}",
        "target_option": "Лицо {i}",
    },
    "en": {
        "lang_radio_label": "Language / Язык",
        "title_md": (
            "# FaceSwap Pro (Docker)\n\n"
            "Left — **donor photo** (can contain multiple faces), right — photo where faces will be replaced.  \n"
            "Below — result, circular progress indicator, time estimate and **Before / After** gallery.\n"
        ),
        "step1_title": "### 1. Donor photo",
        "step1_input_label": "Upload donor photo (can contain multiple faces)",
        "step1_donor_choice_label": "Choose donor face",
        "step2_title": "### 2. Target photo",
        "step2_input_label": "Upload image where faces will be replaced",
        "step2_target_choices_label": "Choose faces to replace (if none selected — replace all)",
        "step3_title": "### 3. Settings & Export",
        "use_enh_label": "Enhance result quality (GFPGAN)",
        "eta_initial": "Time estimate will appear after detecting faces.",
        "fmt_label": "Download file format",
        "run_btn": "3. Start swapping",
        "download_btn": "Download result",
        "before_after_label": "Before / After",
        "eta_fmt_sec": "Estimated processing time: ~{sec} sec.",
        "eta_fmt_min": "Estimated processing time: ~{min} min {sec} sec.",
        "msg_need_images": "Please upload both images to start.",
        "msg_no_donor_faces": "No faces detected on donor image.",
        "msg_no_target_faces": "No faces detected on target image.",
        "msg_prep": "Preparing for processing...",
        "msg_swap_step": "Replacing face {i} of {n}",
        "msg_enh": "Enhancing quality (GFPGAN)...",
        "msg_done": "Done. Faces processed: {n}.",
        "progress_done": "Done!",
        "donor_option": "Donor {i}",
        "target_option": "Face {i}",
    },
}


def T(lang: str, key: str, **kwargs) -> str:
    s = TEXTS[lang][key]
    return s.format(**kwargs)


# =========================================================
# Model init
# =========================================================
try:
    import torch
    has_gpu = torch.cuda.is_available()
except Exception:
    has_gpu = False

ctx_id = 0 if has_gpu else -1
print("[INFO] GPU available" if has_gpu else "[INFO] Running on CPU")

# buffalo_l must already exist in /root/.insightface/models/buffalo_l (Dockerfile copies it)
app_face = FaceAnalysis(name="buffalo_l")
app_face.prepare(ctx_id=ctx_id, det_size=(640, 640))

# load inswapper from local path; do not download
swapper = insightface.model_zoo.get_model(INSWAPPER_PATH, download=False)

# GFPGAN from local weights
face_enhancer = GFPGANer(
    model_path=GFPGAN_PATH,
    upscale=1,
    arch="clean",
    channel_multiplier=2,
    bg_upsampler=None,
)

MAX_PREVIEWS = 8


# =========================================================
# ETA
# =========================================================
def estimate_time(lang: str, num_faces: int, use_enhancer: bool) -> str:
    if num_faces <= 0:
        return TEXTS[lang]["eta_initial"]

    if has_gpu:
        base_per_face = 4
        extra = 6 if use_enhancer else 0
    else:
        base_per_face = 10
        extra = 25 if use_enhancer else 0

    total = num_faces * (base_per_face + extra)
    if total < 60:
        return T(lang, "eta_fmt_sec", sec=int(total))
    m = int(total // 60)
    s = int(total % 60)
    return T(lang, "eta_fmt_min", min=m, sec=s)


# =========================================================
# Face detection
# =========================================================
def detect_faces_generic(img):
    if img is None:
        return [], []
    bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    faces = app_face.get(bgr)
    faces = sorted(faces, key=lambda f: f.bbox[0])
    previews = []
    for f in faces:
        x1, y1, x2, y2 = map(int, f.bbox)
        h, w = bgr.shape[:2]
        x1 = max(0, x1)
        y1 = max(0, y1)
        x2 = min(w, x2)
        y2 = min(h, y2)
        crop = bgr[y1:y2, x1:x2]
        previews.append(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
    return previews, faces


def update_donor_faces(img, lang: str):
    previews, faces = detect_faces_generic(img)
    updates = []
    for i in range(MAX_PREVIEWS):
        if i < len(previews):
            updates.append(gr.update(value=previews[i], visible=True))
        else:
            updates.append(gr.update(value=None, visible=False))

    labels = [T(lang, "donor_option", i=i + 1) for i in range(len(previews))]
    default = labels[0] if labels else None
    radio_update = gr.update(choices=labels, value=default)
    return updates + [radio_update, faces]


def update_target_faces(img, use_enhancer: bool, lang: str):
    previews, faces = detect_faces_generic(img)
    updates = []
    for i in range(MAX_PREVIEWS):
        if i < len(previews):
            updates.append(gr.update(value=previews[i], visible=True))
        else:
            updates.append(gr.update(value=None, visible=False))

    labels = [T(lang, "target_option", i=i + 1) for i in range(len(previews))]
    default = labels if labels else []
    checkbox_update = gr.update(choices=labels, value=default)

    eta_text = estimate_time(lang, len(faces), use_enhancer)
    return updates + [checkbox_update, faces, eta_text]


def update_eta_only(faces, use_enhancer: bool, lang: str):
    num = len(faces) if faces else 0
    return estimate_time(lang, num, use_enhancer)


# =========================================================
# Progress HTML
# =========================================================
def make_progress_html(percent: int, text: str) -> str:
    percent = max(0, min(100, int(percent)))
    return f"""
<div class="progress-container">
  <div class="progress-ring" style="--p:{percent};">
    <div class="progress-ring-inner">
      <span class="progress-ring-percent">{percent}%</span>
    </div>
  </div>
  <div class="progress-ring-text">{text}</div>
</div>
"""


# =========================================================
# Main swap generator
# =========================================================
def swap_from_ui(
    donor_img,
    target_img,
    donor_choice,
    target_choices,
    donor_faces,
    target_faces,
    use_enhancer: bool,
    lang: str,
):
    if donor_img is None or target_img is None:
        msg = T(lang, "msg_need_images")
        html = make_progress_html(0, msg)
        yield None, msg, html, []
        return

    if not donor_faces:
        msg = T(lang, "msg_no_donor_faces")
        html = make_progress_html(0, msg)
        yield target_img, msg, html, [target_img]
        return

    if not target_faces:
        msg = T(lang, "msg_no_target_faces")
        html = make_progress_html(0, msg)
        yield target_img, msg, html, [target_img]
        return

    # donor face index
    if donor_choice:
        try:
            donor_idx = int(donor_choice.split()[-1]) - 1
        except Exception:
            donor_idx = 0
    else:
        donor_idx = 0
    donor_idx = max(0, min(donor_idx, len(donor_faces) - 1))
    donor_face = donor_faces[donor_idx]

    # target indices
    if target_choices:
        target_indices = []
        for lbl in target_choices:
            try:
                idx = int(lbl.split()[-1]) - 1
                if 0 <= idx < len(target_faces):
                    target_indices.append(idx)
            except Exception:
                continue
        if not target_indices:
            target_indices = list(range(len(target_faces)))
    else:
        target_indices = list(range(len(target_faces)))

    tar_bgr = cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR)
    result = tar_bgr.copy()

    total_steps = len(target_indices) + (1 if use_enhancer else 0)
    total_steps = max(total_steps, 1)

    msg = T(lang, "msg_prep")
    yield cv2.cvtColor(result, cv2.COLOR_BGR2RGB), msg, make_progress_html(0, msg), []

    # swap loop
    for i, idx in enumerate(target_indices):
        step = i + 1
        msg = T(lang, "msg_swap_step", i=step, n=len(target_indices))
        pct = int(step / total_steps * 100)
        yield cv2.cvtColor(result, cv2.COLOR_BGR2RGB), msg, make_progress_html(pct, msg), []
        face_obj = target_faces[idx]
        result = swapper.get(result, face_obj, donor_face, paste_back=True)

    # GFPGAN
    if use_enhancer:
        msg = T(lang, "msg_enh")
        pct = int((total_steps - 1) / total_steps * 100)
        yield cv2.cvtColor(result, cv2.COLOR_BGR2RGB), msg, make_progress_html(pct, msg), []
        try:
            _, _, enhanced = face_enhancer.enhance(
                result,
                has_aligned=False,
                only_center_face=False,
                paste_back=True,
            )
            result = enhanced
        except Exception as e:
            print("[WARN] GFPGAN failed:", e)

    final_img = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
    msg = T(lang, "msg_done", n=len(target_indices))
    html = make_progress_html(100, T(lang, "progress_done"))
    before_after = [target_img, final_img]
    yield final_img, msg, html, before_after


# =========================================================
# Save
# =========================================================
def save_result(image, fmt: str):
    if image is None:
        return None
    img = Image.fromarray(image)
    fmt = fmt.lower()
    path = f"/tmp/faceswap.{fmt}"
    img.save(path, format=fmt.upper())
    return path


# =========================================================
# CSS + Theme
# =========================================================
custom_css = """
.gradio-container { max-width: 1240px !important; margin: 0 auto !important; }
.step-card {
    background: radial-gradient(circle at top left, #0b1220, #020617 55%);
    border-radius: 18px;
    padding: 18px 20px;
    border: 1px solid #111827;
    box-shadow: 0 22px 60px rgba(15,23,42,0.9);
}
.face-thumb img {
    width: 100% !important;
    height: 100% !important;
    object-fit: cover;
    border-radius: 8px;
}
.progress-container { display:flex; flex-direction:column; align-items:center; justify-content:center; margin-top:14px; }
.progress-ring {
    width:120px; height:120px; border-radius:50%;
    background:conic-gradient(#8b5cf6 calc(var(--p) * 1%), #1f2937 0);
    display:flex; align-items:center; justify-content:center;
    box-shadow: 0 0 25px rgba(139,92,246,0.6);
    transition: background 0.6s ease-out, box-shadow 0.6s ease-out;
}
.progress-ring-inner {
    width:85px; height:85px; border-radius:50%;
    background:#020617; display:flex; align-items:center; justify-content:center;
}
.progress-ring-percent { font-size:22px; font-weight:700; color:#e5e7eb; }
.progress-ring-text { margin-top:8px; font-size:14px; color:#e5e7eb; text-align:center; max-width:260px; }
button#run-btn {
    background-image: linear-gradient(135deg, #a855f7, #ec4899);
    border-radius: 999px; border: none; font-weight: 600;
    letter-spacing: 0.03em;
    box-shadow: 0 12px 30px rgba(236,72,153,0.45);
}
button#run-btn:hover { filter: brightness(1.08); box-shadow: 0 16px 44px rgba(236,72,153,0.6); }
button#download-btn {
    background-image: linear-gradient(135deg, #22c55e, #16a34a);
    border-radius: 999px; border: none; color: #0f172a; font-weight: 600;
    letter-spacing: 0.03em;
    box-shadow: 0 10px 24px rgba(34,197,94,0.45);
}
button#download-btn:hover { filter: brightness(1.05); box-shadow: 0 14px 32px rgba(34,197,94,0.6); }
"""

premium_theme = gr.themes.Default(primary_hue="violet", neutral_hue="slate").set(
    body_text_size="16px",
    body_text_weight="500",
    button_large_text_size="16px",
    button_large_padding="12px 22px",
)

# =========================================================
# UI
# =========================================================
with gr.Blocks(
    title="FaceSwap Pro (Docker)",
    css=custom_css,
    theme=premium_theme,
) as demo:
    lang_state = gr.State("ru")

    lang_radio = gr.Radio(
        choices=["RU", "EN"],
        value="RU",
        label=TEXTS["ru"]["lang_radio_label"],
    )

    title_md = gr.Markdown(TEXTS["ru"]["title_md"])

    donor_faces_state = gr.State([])
    target_faces_state = gr.State([])

    with gr.Row(elem_classes=["step-card"]):
        with gr.Column():
            step1_title_md = gr.Markdown(TEXTS["ru"]["step1_title"])
            donor_img = gr.Image(
                label=TEXTS["ru"]["step1_input_label"],
                type="numpy",
                height=420,
            )

            gr.Markdown("**Найденные лица (донор):**")
            with gr.Row():
                donor_previews = [
                    gr.Image(
                        value=None,
                        visible=False,
                        show_label=False,
                        interactive=False,
                        width=32,
                        height=32,
                        type="numpy",
                        elem_classes=["face-thumb"],
                    )
                    for _ in range(MAX_PREVIEWS)
                ]

            donor_choice = gr.Radio(
                label=TEXTS["ru"]["step1_donor_choice_label"],
                choices=[],
            )

        with gr.Column():
            step2_title_md = gr.Markdown(TEXTS["ru"]["step2_title"])
            target_img = gr.Image(
                label=TEXTS["ru"]["step2_input_label"],
                type="numpy",
                height=420,
            )

            gr.Markdown("**Найденные лица (цель):**")
            with gr.Row():
                target_previews = [
                    gr.Image(
                        value=None,
                        visible=False,
                        show_label=False,
                        interactive=False,
                        width=32,
                        height=32,
                        type="numpy",
                        elem_classes=["face-thumb"],
                    )
                    for _ in range(MAX_PREVIEWS)
                ]

            target_choices = gr.CheckboxGroup(
                label=TEXTS["ru"]["step2_target_choices_label"],
                choices=[],
            )

    with gr.Row(elem_classes=["step-card"]):
        with gr.Column(scale=1):
            step3_title_md = gr.Markdown(TEXTS["ru"]["step3_title"])
            use_enh = gr.Checkbox(label=TEXTS["ru"]["use_enh_label"], value=True)
            eta_text = gr.Markdown(TEXTS["ru"]["eta_initial"])
            fmt = gr.Dropdown(
                label=TEXTS["ru"]["fmt_label"],
                choices=["png", "jpeg", "webp"],
                value="png",
            )
            run_btn = gr.Button(TEXTS["ru"]["run_btn"], variant="primary", elem_id="run-btn")
            progress_html = gr.HTML("")
            download_btn = gr.DownloadButton(TEXTS["ru"]["download_btn"], elem_id="download-btn")

        with gr.Column(scale=2):
            result_img = gr.Image(label="", show_label=False, interactive=False, type="numpy", height=520)
            status_md = gr.Markdown("")
            before_after = gr.Gallery(label=TEXTS["ru"]["before_after_label"], columns=2, height=350)

    def switch_language(choice):
        lang = "ru" if choice == "RU" else "en"
        t = TEXTS[lang]
        return (
            lang,
            gr.update(label=t["lang_radio_label"]),
            gr.update(value=t["title_md"]),
            gr.update(value=t["step1_title"]),
            gr.update(label=t["step1_input_label"]),
            gr.update(label=t["step1_donor_choice_label"]),
            gr.update(value=t["step2_title"]),
            gr.update(label=t["step2_input_label"]),
            gr.update(label=t["step2_target_choices_label"]),
            gr.update(value=t["step3_title"]),
            gr.update(label=t["use_enh_label"]),
            gr.update(value=t["eta_initial"]),
            gr.update(label=t["fmt_label"]),
            gr.update(value=t["run_btn"]),
            gr.update(value=t["download_btn"]),
            gr.update(label=t["before_after_label"]),
        )

    lang_radio.change(
        fn=switch_language,
        inputs=lang_radio,
        outputs=[
            lang_state,
            lang_radio,
            title_md,
            step1_title_md,
            donor_img,
            donor_choice,
            step2_title_md,
            target_img,
            target_choices,
            step3_title_md,
            use_enh,
            eta_text,
            fmt,
            run_btn,
            download_btn,
            before_after,
        ],
    )

    donor_img.change(
        fn=update_donor_faces,
        inputs=[donor_img, lang_state],
        outputs=donor_previews + [donor_choice, donor_faces_state],
    )

    target_img.change(
        fn=update_target_faces,
        inputs=[target_img, use_enh, lang_state],
        outputs=target_previews + [target_choices, target_faces_state, eta_text],
    )

    use_enh.change(
        fn=update_eta_only,
        inputs=[target_faces_state, use_enh, lang_state],
        outputs=eta_text,
    )

    run_btn.click(
        fn=swap_from_ui,
        inputs=[
            donor_img,
            target_img,
            donor_choice,
            target_choices,
            donor_faces_state,
            target_faces_state,
            use_enh,
            lang_state,
        ],
        outputs=[result_img, status_md, progress_html, before_after],
    )

    download_btn.click(
        fn=save_result,
        inputs=[result_img, fmt],
        outputs=download_btn,
    )

demo.launch(
    server_name="0.0.0.0",
    server_port=int(os.getenv("PORT", "7860")),
)