import os
import time
from pathlib import Path

import cv2
import gradio as gr

from optimized_swapper import FaceSwapEngine

ENGINE = FaceSwapEngine()


def _as_index(v, default=1):
    try:
        return max(1, int(v))
    except Exception:
        return default


def _file_to_path(item):
    if item is None:
        return None
    if isinstance(item, (str, Path)):
        return str(item)
    if hasattr(item, "name"):
        return item.name
    if isinstance(item, dict):
        for key in ("path", "name"):
            if key in item and item[key]:
                return str(item[key])
    if isinstance(item, (list, tuple)) and item:
        first = item[0]
        if isinstance(first, (str, Path)):
            return str(first)
        if hasattr(first, "name"):
            return first.name
    return None


def _normalize_video_input(video_value):
    if video_value is None:
        return None
    if isinstance(video_value, (str, Path)):
        return str(video_value)
    if isinstance(video_value, dict):
        if video_value.get("video"):
            return str(video_value["video"])
        if video_value.get("path"):
            return str(video_value["path"])
    if isinstance(video_value, (list, tuple)) and video_value:
        first = video_value[0]
        if isinstance(first, (str, Path)):
            return str(first)
        if isinstance(first, dict):
            return str(first.get("video") or first.get("path") or "") or None
    return None


def _eta_text(start_time, progress_value):
    elapsed = max(0.001, time.time() - start_time)
    if progress_value <= 0:
        return f"Elapsed: {elapsed:.1f}s | ETA: estimating..."
    remaining = elapsed * (1.0 - progress_value) / progress_value
    if remaining >= 60:
        return f"Elapsed: {elapsed:.1f}s | ETA: {remaining / 60:.1f} min"
    return f"Elapsed: {elapsed:.1f}s | ETA: {remaining:.1f}s"


def _status(progress_value, message, start_time):
    pct = int(max(0, min(100, round(progress_value * 100))))
    return f"{message}\nProgress: {pct}%\n{_eta_text(start_time, progress_value)}"


def swap_photo(source, source_idx, target, target_idx, det_size, progress=gr.Progress()):
    start = time.time()
    progress(0.05, desc="Preparing image swap")
    if source is None or target is None:
        raise gr.Error("Upload both source and target images.")
    progress(0.35, desc="Running face detection")
    result = ENGINE.swap_image(
        source,
        _as_index(source_idx),
        target,
        _as_index(target_idx),
        det_size=int(det_size),
    )
    progress(1.0, desc="Done")
    return result, _status(1.0, "Photo swap complete.", start)


def swap_video(source, source_idx, video_path, target_idx, det_size, detection_interval, jpeg_quality, audio, progress=gr.Progress()):
    start = time.time()
    progress(0.02, desc="Validating inputs")
    if source is None:
        raise gr.Error("Upload a source image.")

    video_path = _normalize_video_input(video_path)
    if not video_path:
        raise gr.Error("Upload a target video.")

    progress(0.10, desc="Preparing source face")
    progress(0.20, desc="Starting video processing")
    result = ENGINE.swap_video(
        source=source,
        source_idx=_as_index(source_idx),
        video_path=video_path,
        target_idx=_as_index(target_idx),
        det_size=int(det_size),
        detection_interval=max(1, int(detection_interval)),
        jpeg_quality=int(jpeg_quality),
        preserve_audio=bool(audio),
    )
    progress(1.0, desc="Video complete")
    return result, _status(1.0, "Video swap complete.", start)


def swap_multi_source_single(source_files, target, target_idx, det_size, progress=gr.Progress()):
    start = time.time()
    if not source_files or target is None:
        raise gr.Error("Upload source images and a target image.")

    results = []
    total = len(source_files)
    for i, item in enumerate(source_files, start=1):
        progress(((i - 1) / max(total, 1)) * 0.9 + 0.05, desc=f"Processing source {i}/{total}")
        path = _file_to_path(item)
        if not path:
            continue
        img = cv2.imread(path)
        if img is None:
            results.append(f"Error: could not read source image {path}")
            continue
        try:
            result = ENGINE.swap_image(
                img,
                1,
                target,
                _as_index(target_idx),
                int(det_size),
            )
            results.append(result)
        except Exception as e:
            results.append(f"Error: {e}")

    if not results:
        raise gr.Error("No readable source images were uploaded.")
    progress(1.0, desc="Batch complete")
    return results, _status(1.0, f"Processed {len(results)} output(s).", start)


def swap_multi_source_multi(source_files, target_files, target_indices, det_size, progress=gr.Progress()):
    start = time.time()
    if not source_files or not target_files:
        raise gr.Error("Upload source and target images.")

    indices = [x.strip() for x in str(target_indices).split(",") if x.strip()]
    results = []
    targets = []

    progress(0.05, desc="Loading target images")
    for item in target_files:
        path = _file_to_path(item)
        if not path:
            continue
        img = cv2.imread(path)
        if img is not None:
            targets.append(img)

    if not targets:
        raise gr.Error("No readable target images were uploaded.")

    total_jobs = max(1, len(source_files) * len(targets))
    done = 0
    for src_item in source_files:
        src_path = _file_to_path(src_item)
        if not src_path:
            continue
        src = cv2.imread(src_path)
        if src is None:
            results.append(f"Error: could not read source image {src_path}")
            continue

        try:
            ENGINE.prepare_source(src, 1, int(det_size))
        except Exception as e:
            results.append(f"Error preparing source {src_path}: {e}")
            continue

        for j, dst in enumerate(targets):
            idx = _as_index(indices[j] if j < len(indices) else 1)
            done += 1
            progress((done / total_jobs) * 0.95, desc=f"Processing pair {done}/{total_jobs}")
            try:
                results.append(ENGINE.swap_prepared_source(dst, idx, int(det_size)))
            except Exception as e:
                results.append(f"Error: {e}")

    if not results:
        raise gr.Error("No output images were generated.")
    progress(1.0, desc="Batch complete")
    return results, _status(1.0, f"Processed {len(results)} output(s).", start)


with gr.Blocks(title="Fast Face Swap") as demo:
    gr.Markdown(
        "# Fast Face Swapping Suite\n"
        "CUDA/ONNX Runtime optimized photo and video face swapping with live progress status."
    )

    with gr.Tab("Single Photo"):
        with gr.Row():
            with gr.Column():
                src = gr.Image(type="numpy", label="Source Image")
                src_idx = gr.Number(value=1, precision=0, label="Source Face Index")
            with gr.Column():
                dst = gr.Image(type="numpy", label="Target Image")
                dst_idx = gr.Number(value=1, precision=0, label="Target Face Index")
        det_single = gr.Dropdown([256, 320, 384, 512], value=320, label="Detector Size")
        btn = gr.Button("Swap", variant="primary")
        out = gr.Image(type="numpy", label="Result")
        photo_status = gr.Textbox(label="Status / ETA", lines=3, interactive=False)
        btn.click(
            fn=swap_photo,
            inputs=[src, src_idx, dst, dst_idx, det_single],
            outputs=[out, photo_status],
            show_progress="full",
        )

    with gr.Tab("Fast Video"):
        vsrc = gr.Image(type="numpy", label="Source Image")
        vsrc_idx = gr.Number(value=1, precision=0, visible=False, label="Source Face Index")
        with gr.Row():
            vid = gr.Video(label="Target Video")
            vdst_idx = gr.Number(value=1, precision=0, label="Target Face Index")
        with gr.Row():
            vdet = gr.Dropdown([256, 320, 384, 512], value=320, label="Detector Size")
            interval = gr.Slider(
                1,
                5,
                value=1,
                step=1,
                label="Face detection interval (1 = best tracking accuracy)",
            )
            quality = gr.Slider(75, 98, value=92, step=1, label="JPEG fallback quality")
            audio = gr.Checkbox(value=True, label="Preserve original audio")
        vbtn = gr.Button("Swap Video", variant="primary")
        vout = gr.Video(label="Output Video")
        video_status = gr.Textbox(label="Status / ETA", lines=3, interactive=False)
        vbtn.click(
            fn=swap_video,
            inputs=[vsrc, vsrc_idx, vid, vdst_idx, vdet, interval, quality, audio],
            outputs=[vout, video_status],
            show_progress="full",
        )

    with gr.Tab("Multi Source -> Single Target"):
        ms = gr.File(file_count="multiple", file_types=["image"], type="filepath", label="Source Images")
        md = gr.Image(type="numpy", label="Target Image")
        mi = gr.Number(value=1, precision=0, label="Target Face Index")
        det_multi_single = gr.Dropdown([256, 320, 384, 512], value=320, label="Detector Size")
        mb = gr.Button("Process", variant="primary")
        mo = gr.Gallery(label="Results", columns=3)
        multi_single_status = gr.Textbox(label="Status / ETA", lines=3, interactive=False)
        mb.click(
            fn=swap_multi_source_single,
            inputs=[ms, md, mi, det_multi_single],
            outputs=[mo, multi_single_status],
            show_progress="full",
        )

    with gr.Tab("Multi Source -> Multi Target"):
        mss = gr.File(file_count="multiple", file_types=["image"], type="filepath", label="Source Images")
        mdd = gr.File(file_count="multiple", file_types=["image"], type="filepath", label="Target Images")
        mids = gr.Textbox(value="1", label="Target face indices, comma-separated")
        det_multi_multi = gr.Dropdown([256, 320, 384, 512], value=320, label="Detector Size")
        mdb = gr.Button("Process", variant="primary")
        mdo = gr.Gallery(label="Results", columns=3)
        multi_multi_status = gr.Textbox(label="Status / ETA", lines=3, interactive=False)
        mdb.click(
            fn=swap_multi_source_multi,
            inputs=[mss, mdd, mids, det_multi_multi],
            outputs=[mdo, multi_multi_status],
            show_progress="full",
        )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=8).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        show_error=True,
    )