import gradio as gr
import insightface
from insightface.app import FaceAnalysis
from PIL import Image
import numpy as np

# ---------- HTML description ----------
wellcomingMessage = """
    <h1>Face Swap</h1>
    <p>by <a href="https://www.tonyassi.com/" target="_blank" style="color:#97d8be;">Tony Assi</a></p>
    <h2>Try out <a href="https://huggingface.co/spaces/tonyassi/video-face-swap" target="_blank">Video Face Swap</a> ❤️</h2>
"""

# ---------- Check version ----------
assert insightface.__version__ >= '0.7'

# ---------- Initialize models ----------
app = FaceAnalysis(name='buffalo_l')
app.prepare(ctx_id=0, det_size=(640, 640))

# ensure model downloads fully
swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True)

# ---------- Swap logic ----------
def swap_faces(src_img, dest_img):
    # Pre-run nudge
    gr.Warning(
        'Skip the limits — HD, priority queue & no watermark at '
        '<a href="https://www.face-swap.co/?utm_source=hfspace_faceswap&utm_medium=warning" '
        'target="_blank" rel="noopener">face-swap.co</a>'
    )

    if src_img is None or dest_img is None:
        raise gr.Error("Please upload both a source and a target image.")

    src_faces = app.get(src_img)
    dest_faces = app.get(dest_img)

    if len(src_faces) == 0 or len(dest_faces) == 0:
        raise gr.Error("No faces detected in one of the images. Try clearer, front-facing photos.")

    # Just swap first detected face from each image
    source_face = src_faces[0]
    dest_face = dest_faces[0]
    result = swapper.get(dest_img, dest_face, source_face, paste_back=True)

    out_img = Image.fromarray(np.uint8(result)).convert("RGB")

    # Post-success promo
    gr.Info(
        '✨ Like this preview?<br>'
        '<strong>Get HD face swaps</strong> (higher resolution, priority queue & no watermark) — '
        '<a href="https://www.face-swap.co/?utm_source=hfspace_faceswap&utm_medium=info" '
        'target="_blank" rel="noopener">Upgrade on face-swap.co</a>',
        duration=8
    )

    return out_img

def open_side():
    # for 4.x Sidebar API
    return gr.Sidebar(open=True)

# ---------- Custom CSS ----------
CUSTOM_CSS = """
.sticky-cta {
  position: sticky; top: 0; z-index: 1000;
  background: #a5b4fc;
  color: #0f172a;
  padding: 10px 14px;
  text-align: center;
  border-bottom: 1px solid #333;
  display: block;
  text-decoration: none;
  cursor: pointer;
}
.sticky-cta:hover { filter: brightness(0.97); }
.sticky-cta .pill { background:#4f46e5; color:#fff; padding:4px 10px; border-radius:999px; margin-left:10px; }
.sticky-cta .cta-link { font-weight:600; text-decoration: underline; }

/* centered, single-label API CTA */
.api-cta-wrap { text-align:center; margin-top:10px; }
.api-cta-hero {
  display:inline-flex; align-items:center; gap:10px;
  padding:10px 14px; border-radius:14px;
  background: linear-gradient(90deg,#0ea5e9 0%, #a8a9de 100%);
  color:#fff; font-weight:800; letter-spacing:0.1px;
  box-shadow: 0 6px 22px rgba(99,102,241,0.35);
  border: 1px solid rgba(255,255,255,0.22);
  text-decoration:none;
}
.api-cta-hero:hover { filter:brightness(1.05); transform: translateY(-1px); transition: all .15s ease; }

.api-cta-hero .new {
  background:#fff; color:#0ea5e9; font-weight:900;
  padding:2px 8px; border-radius:999px; font-size:12px; line-height:1;
}
.api-cta-hero .txt { font-weight:800; }
.api-cta-hero .chev { opacity:.95; }
@media (max-width: 520px){
  .api-cta-hero { padding:9px 12px; gap:8px; font-size:14px; }
  .api-cta-hero .new { display:none; }
}

/* floating bottom promo */
.bottom-promo {
  position: fixed; left: 50%; transform: translateX(-50%);
  bottom: 16px; z-index: 1001; background:#0b0b0b; color:#fff;
  border: 1px solid #2a2a2a; border-radius: 12px; padding: 10px 14px;
  box-shadow: 0 8px 24px rgba(0,0,0,0.3);
}
.bottom-promo a { color:#4ea1ff; text-decoration:none; font-weight:600; }

/* big CTA button */
.upgrade-btn { width: 100%; font-size: 16px; padding: 10px 14px; }

/* hero markdown centering + larger heading */
#hero-md {
  text-align: center;
}
#hero-md h3,
#hero-md .prose h3 {
  font-size: 2.1rem;
  line-height: 1.2;
  font-weight: 800;
  margin-bottom: 0.25rem;
}
#hero-md p,
#hero-md .prose p {
  font-size: 1.05rem;
}
"""

# ---------- UI (Blocks) ----------
# NOTE: no css= here – we inject CSS via <style> tag instead
with gr.Blocks(title="Image Face Swap") as demo:
    # Inject custom CSS manually
    gr.HTML(f"<style>{CUSTOM_CSS}</style>")

    # Sticky banner
    gr.HTML(
        '<a class="sticky-cta" '
        'href="https://www.face-swap.co/?utm_source=hfspace_faceswap&utm_medium=banner" '
        'target="_blank" rel="noopener" aria-label="Upgrade to Pro on face-swap.co">'
        '⚡ <strong>Upgrade to HD</strong> — priority queue & higher resolution swaps!'
        '<span class="pill">GPU</span>'
        '</a>'
    )

    # Hero / intro
    gr.Markdown(
        f"""
        ### Image Face Swap (Preview)
        [face-swap.co](https://www.face-swap.co/?utm_source=hfspace_faceswap&utm_medium=subtitle)
        
        **Free preview** runs on CPU and may be limited in resolution to keep it fast.  
        Want full-quality **HD AI face swap** with GPU speed? **[Go Pro ↗](https://www.face-swap.co/?utm_source=hfspace_faceswap&utm_medium=go_pro)**
        
        """,
        elem_id="hero-md"
    )

    # API CTA
    gr.HTML(
        """
        <div class="api-cta-wrap">
          <a class="api-cta-hero"
             href="https://www.face-swap.co/api?utm_source=hfspace_faceswap&utm_medium=hero_api_new"
             target="_blank" rel="noopener" aria-label="Face Swap API">
            <span class="new">NEW</span>
            <span class="txt">Face Swap API</span>
            <span class="chev">↗</span>
          </a>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=5):
            src_img = gr.Image(type="numpy", label="Source Image (face to copy)")
            dest_img = gr.Image(type="numpy", label="Target Image (face to replace)")

            go = gr.Button("Swap Face", variant="primary")
            pro = gr.Button("⚡ Upgrade to HD on face-swap.co", elem_classes=["upgrade-btn"])

        with gr.Column(scale=5):
            out_img = gr.Image(label="Result")

    # Examples
    gr.Examples(
        examples=[["./Images/kim.jpg", "./Images/marilyn.jpg"]],
        inputs=[src_img, dest_img],
        outputs=[out_img],
        fn=swap_faces,
        cache_examples=True,
        run_on_click=True,
        label="Try an example"
    )

    # Sidebar CTA
    with gr.Sidebar(open=False) as side:
        gr.Markdown(
            "### Upgrade to HD 1920x1080\n"
            "- Higher resolution face swaps\n"
            "- Priority queue\n"
            "- API access & automation\n"
            "- No watermark"
        )
        pro_pro = gr.Button("Open Pro Checkout", variant="primary")
        pro_api = gr.Button("API Access", variant="primary")

    # Floating bottom promo
    gr.HTML(
        '<div class="bottom-promo">'
        'Want HD & faster processing? '
        '<a href="https://www.face-swap.co/?utm_source=hfspace_faceswap&utm_medium=upgrade" '
        'target="_blank" rel="noopener">Upgrade</a>'
        '</div>'
    )

    # Open sidebar when user starts a swap
    go.click(fn=open_side, inputs=None, outputs=side, queue=False)

    # Main action
    go.click(fn=swap_faces, inputs=[src_img, dest_img], outputs=out_img)

    # JS-only Pro buttons
    pro.click(
        fn=None, inputs=None, outputs=None,
        js="()=>window.open('https://www.face-swap.co/?utm_source=hfspace_faceswap&utm_medium=upgrade_to_hd','_blank')"
    )
    pro_pro.click(
        fn=None, inputs=None, outputs=None,
        js="()=>window.open('https://www.face-swap.co/?utm_source=hfspace_faceswap&utm_medium=sidebar_pro','_blank')"
    )
    pro_api.click(
        fn=None, inputs=None, outputs=None,
        js="()=>window.open('https://www.face-swap.co/api?utm_source=hfspace_faceswap&utm_medium=sidebar_api','_blank')"
    )

    demo.queue()

if __name__ == "__main__":
    # If this line ever complains about theme=, just drop the arg:
    # demo.launch()
    demo.launch(theme=gr.themes.Soft())