import gradio as gr
import os, uuid, time
from gradio_client import Client, handle_file
from moviepy.editor import VideoFileClip

# ── config ───────────────────────────────────────────────────────────
hf_token   = os.environ.get("TOKEN")
output_dir = "uploads/output"
os.makedirs(output_dir, exist_ok=True)
client = Client("tonyassi/vfs2-cpu", hf_token=hf_token, download_files=output_dir)

UTM = "utm_source=hugging_face_space&utm_medium=banner&utm_campaign=pro_cta"
PRO_URL = f"https://www.face-swap.co/?{UTM}"

# ── helpers ──────────────────────────────────────────────────────────
def preprocess_video(path: str, target_fps: int = 12,
                     target_size: int = 800, target_len: int = 4) -> str:
    clip = VideoFileClip(path)
    if clip.duration > target_len:
        clip = clip.subclip(0, target_len)
    w, h = clip.size
    clip = clip.resize(width=target_size) if w >= h else clip.resize(height=target_size)
    clip = clip.set_fps(target_fps)
    out_path = os.path.join(output_dir, f"pre_{uuid.uuid4().hex}.mp4")
    clip.write_videofile(out_path, codec="libx264", audio_codec="aac",
                         fps=target_fps, verbose=False, logger=None)
    clip.close()
    return out_path

# ── main generate ────────────────────────────────────────────────────
def generate(input_image, input_video, gender):
    # Pre-run nudge (small)
    gr.Warning(
        f'Skip the line — HD, no watermark, priority queue at '
        f'<a href="https://www.face-swap.co/?utm_source=hfspace_deepfakevideo&utm_medium=warning" target="_blank" rel="noopener">face-swap.co</a>'
    )

    if gender == "all":
        gender = None

    try:
        pre_video = preprocess_video(input_video)
        job = client.submit(
            input_image=handle_file(input_image),
            input_video={"video": handle_file(pre_video)},
            device='cpu',
            selector='many',
            gender=gender,
            race=None,
            order=None,
            api_name="/predict"
        )
        while not job.done():
            time.sleep(5)

        if not job.status().success:
            return None

        # Post-success modal (big)
        gr.Info(
            f"✨ Your preview is ready.<br>"
            f"<strong>Get HD</strong> (4× quality, no watermark, priority) — "
            f'<a href="https://www.face-swap.co/?utm_source=hfspace_deepfakevideo&utm_medium=info" target="_blank" rel="noopener">Upgrade on face-swap.co</a>',
            duration=8
        )

        video_path = job.outputs()[0]["video"]
        return video_path

    except Exception as e:
        gr.Error(f"Generation failed: {e}")
        return None

def open_side():  # tiny helper
    return gr.Sidebar(open=True)

# ── UI (Blocks) ──────────────────────────────────────────────────────
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;                 /* full-width clickable */
  text-decoration: none;          /* remove underline */
  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; }



/* 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, /* standard markdown h3 */
#hero-md .prose h3 { /* some gradio themes wrap markdown with .prose */
  font-size: 2.1rem;   /* ~34px */
  line-height: 1.2;
  font-weight: 800;
  margin-bottom: 0.25rem;
}
#hero-md p,
#hero-md .prose p {
  font-size: 1.05rem;  /* slightly larger body text */
}

"""

with gr.Blocks(title="Video Face Swap", theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
    # Sticky banner
    gr.HTML(
        f"""<a class="sticky-cta" href="https://www.face-swap.co/?utm_source=hfspace_deepfakevideo&utm_medium=banner" target="_blank" rel="noopener"
              aria-label="Upgrade to Pro on face-swap.co">
              ⚡ <strong>Upgrade to HD</strong> — priority queue & no duration limit!
              <span class="pill">GPU</span>
            </a>"""
    )


    gr.Markdown(
        f"""
        ### Deep Fake Video (Preview)
        [face-swap.co](https://www.face-swap.co/?utm_source=hfspace_deepfakevideo&utm_medium=subtitle)
        
        **Free preview** is downsampled to 800px • 4s • 12fps to reduce wait time.  
        Want full-length **HD deep fake video** with GPU speed? **[Go Pro ↗](https://www.face-swap.co/?utm_source=hfspace_deepfakevideo&utm_medium=go_pro)**
        """,
        elem_id="hero-md"
    )


    with gr.Row():
        with gr.Column(scale=5):
            
            in_img  = gr.Image(type="filepath", label="Source Image")
            in_vid  = gr.Video(label="Target Video")
            in_gen  = gr.Radio(choices=["all","female","male"], value="all", label="Gender")

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

        with gr.Column(scale=5):
            out_vid = gr.Video(label="Result")

    gr.Examples(
        examples=[["elon.png", "ironman.mp4", "all"], ["bella.jpg", "wizard.mp4", "all"]],
        inputs=[in_img, in_vid, in_gen],
        outputs=[out_vid],
        fn=generate,            # precompute + cache example output
        cache_examples=True,    # store the result on build so it loads instantly
        run_on_click=True,      # clicking the example triggers generate()
        label="Try an example"
    )

    # Sidebar CTA
    with gr.Sidebar(open=False) as side:
        gr.Markdown("### Upgrade to HD 1920x1080\n- 4× quality\n- Priority queue\n- 1-5 minute video duration\n- GPU speed")
        pro_paypergen = gr.Button("Pay per Generation", variant="primary")
        pro_subscription = gr.Button("Monthly Subscription", variant="primary")
        pro_api = gr.Button("Face Swap API", variant="primary")

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

    
    go.click(fn=open_side, inputs=None, outputs=side, queue=False)  # fire instantly

    # Wire events
    go.click(fn=generate, inputs=[in_img, in_vid, in_gen], outputs=out_vid)

    # Open Pro in new tab via JS (no Python call)
    pro.click(fn=None, inputs=None, outputs=None,
              js=f"()=>window.open('https://www.face-swap.co/?utm_source=hfspace_deepfakevideo&utm_medium=upgrade_to_hd','_blank')")
    pro_paypergen.click(fn=None, inputs=None, outputs=None,
               js=f"()=>window.open('https://www.face-swap.co/?utm_source=hfspace_deepfakevideo&utm_medium=sidebar_paypergen','_blank')")
    pro_subscription.click(fn=None, inputs=None, outputs=None,
               js=f"()=>window.open('https://www.face-swap.co/?utm_source=hfspace_deepfakevideo&utm_medium=sidebar_subscription','_blank')")
    pro_api.click(fn=None, inputs=None, outputs=None,
               js=f"()=>window.open('https://www.face-swap.co/?utm_source=hfspace_deepfakevideo&utm_medium=sidebar_api','_blank')")

    # Queue for long jobs + to ensure alerts appear as modals
    demo.queue()

if __name__ == "__main__":
    demo.launch()