import os
import cv2
import numpy as np
from PIL import Image
import urllib.request

import gradio as gr
import insightface
from insightface.app import FaceAnalysis

# ---- Загрузка моделей при первом старте ----

def download_file(url: str, filename: str):
    if os.path.exists(filename):
        return
    print(f"Скачиваю {filename} из {url} ...")
    urllib.request.urlretrieve(url, filename)
    print(f"{filename} скачан.")

# inswapper_128.onnx (Hugging Face repo, который ты нашёл)
INSWAPPER_URL = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
INSWAPPER_PATH = "inswapper_128.onnx"

# веса GFPGAN
GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth"
GFPGAN_PATH = "GFPGANv1.3.pth"

download_file(INSWAPPER_URL, INSWAPPER_PATH)
download_file(GFPGAN_URL, GFPGAN_PATH)

assert os.path.exists(INSWAPPER_PATH), "inswapper_128.onnx не найден."
assert os.path.exists(GFPGAN_PATH), "GFPGANv1.3.pth не найден."

# --- SHIM для старого импорта torchvision.transforms.functional_tensor (нужен gfpgan/basicsr) ---
import sys, types
from torchvision.transforms import functional as F
mod = types.ModuleType("torchvision.transforms.functional_tensor")
mod.rgb_to_grayscale = F.rgb_to_grayscale
sys.modules["torchvision.transforms.functional_tensor"] = mod

from gfpgan import GFPGANer


# ---------- ТЕКСТЫ ДЛЯ ЛОКАЛИЗАЦИИ ----------

TEXTS = {
    "ru": {
        "lang_radio_label": "Язык / Language",
        "title_md": (
            "# FaceSwap Pro\n\n"
            "Слева — фото **донора** (может быть несколько лиц), справа — фото, где нужно заменить лица.  \n"
            "Внизу — результат, круговой индикатор прогресса, оценка времени и блок «До / После».\n\n"
            "_Основную архитектуру и код приложения подготовил AI‑ассистент **ChatGPT (OpenAI)** совместно с вами._"
        ),
        "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\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\n"
            "_Core architecture and code prepared by AI assistant **ChatGPT (OpenAI)** together with you._"
        ),
        "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)


# ---------- ИНИЦИАЛИЗАЦИЯ МОДЕЛЕЙ ----------

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

ctx_id = 0 if has_gpu else -1  # 0 = GPU, -1 = CPU
print("GPU доступен:" if has_gpu else "Работаем на CPU")

app_face = FaceAnalysis(name="buffalo_l")
app_face.prepare(ctx_id=ctx_id, det_size=(640, 640))

swapper = insightface.model_zoo.get_model(INSWAPPER_PATH, download=False)

face_enhancer = GFPGANer(
    model_path=GFPGAN_PATH,
    upscale=1,
    arch="clean",
    channel_multiplier=2,
    bg_upsampler=None,
)

MAX_PREVIEWS = 8  # максимум маленьких превью


# ---------- ПРОГНОЗ ВРЕМЕНИ ----------

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_per_face = 6 if use_enhancer else 0
    else:
        base_per_face = 10
        extra_per_face = 25 if use_enhancer else 0

    total = num_faces * (base_per_face + extra_per_face)
    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)


# ---------- ВСПОМОГАТЕЛЬНЫЕ ФУНКЦИИ ДЕТЕКЦИИ ----------

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)


# ---------- 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>
"""


# ---------- ОСНОВНАЯ ФУНКЦИЯ ЗАМЕНЫ (ГЕНЕРАТОР С ПРОГРЕССОМ) ----------

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

    # индекс донорского лица
    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]

    # индексы целевых лиц
    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)
    if total_steps <= 0:
        total_steps = 1

    # Стартовое кольцо
    msg = T(lang, "msg_prep")
    yield cv2.cvtColor(result, cv2.COLOR_BGR2RGB), msg, make_progress_html(0, msg), []

    # Замена лиц
    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)

    # Улучшение качества
    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("Ошибка GFPGAN, продолжаем без улучшения:", 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


# ---------- СОХРАНЕНИЕ РЕЗУЛЬТАТА ----------

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: премиум‑оформление + превью + прогресс‑кольцо ----------

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);
}
.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;
}

/* премиум‑кнопка */
.gr-button-primary {
    background-image: linear-gradient(135deg, #8b5cf6, #ec4899);
    border-radius: 999px;
    border: none;
    font-weight: 600;
    letter-spacing: 0.03em;
    box-shadow: 0 12px 30px rgba(236,72,153,0.45);
}
.gr-button-primary:hover {
    filter: brightness(1.06);
    box-shadow: 0 16px 44px rgba(236,72,153,0.55);
}
"""

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",
)


# ---------- ИНТЕРФЕЙС ----------

with gr.Blocks(
    title="FaceSwap Pro (ChatGPT)",
    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([])

    # --- Блок 1: Донор и целевое фото ---
    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",
            )
            with gr.Row():
                donor_previews = [
                    gr.Image(
                        value=None,
                        visible=False,
                        show_label=False,
                        interactive=False,
                        width=64,
                        height=64,
                        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",
            )
            with gr.Row():
                target_previews = [
                    gr.Image(
                        value=None,
                        visible=False,
                        show_label=False,
                        interactive=False,
                        width=64,
                        height=64,
                        type="numpy",
                        elem_classes=["face-thumb"],
                    )
                    for _ in range(MAX_PREVIEWS)
                ]
            target_choices = gr.CheckboxGroup(
                label=TEXTS["ru"]["step2_target_choices_label"],
                choices=[],
            )

    # --- Блок 2: результат и настройки ---
    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")
            progress_html = gr.HTML("")
            download_btn = gr.DownloadButton(TEXTS["ru"]["download_btn"])

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

    # ---------- ЛОКАЛИЗАЦИЯ UI ----------

    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()