# app.py
import sys
import os
import tempfile
from pathlib import Path
import streamlit as st
from PIL import Image

# Force app root and pipeline to be importable (for HF docker)
sys.path.append(str(Path(__file__).parent.resolve()))
sys.path.append(str(Path(__file__).parent.resolve() / "pipeline"))

# === CUDA DIAGNOSTICS (LOGS) ===
try:
    import torch
    print("=== CUDA Info ===")
    print("CUDA available:", torch.cuda.is_available())
    print("CUDA device count:", torch.cuda.device_count())
    print("torch version:", torch.__version__)
    print("=================")
except Exception as e:
    print(f"[WARN] Could not print CUDA info: {e}")

# === Import the two-stage pipeline ===
try:
    from pipeline.two_stage_pipeline import process_two_stage
except ImportError as e:
    st.error(f"Pipeline not found or misconfigured: {e}")
    st.stop()

st.set_page_config(
    page_title="BackgroundFX Pro",
    layout="wide",
    page_icon="🎬"
)

st.title("🎬 BackgroundFX Pro — AI Video Background Replacer")
st.write(
    "Upload a video and a background image (optional). The system will remove the original background "
    "using **SAM2** & **MatAnyone**, and composite the subject over your new background image."
)

def progress_callback(msg, prog):
    """Streamlit-friendly progress reporting (persistent)."""
    if "progress_bar" not in st.session_state:
        st.session_state.progress_bar = st.progress(0)
        st.session_state.progress_text = st.empty()
    st.session_state.progress_bar.progress(min(max(prog,0.0),1.0))
    st.session_state.progress_text.text(msg)

# --- Uploaders ---
video_file = st.file_uploader(
    "1️⃣ Upload video file (mp4, mov, avi, mkv)", 
    type=["mp4", "mov", "avi", "mkv"]
)
bg_file = st.file_uploader(
    "2️⃣ Upload new background (jpg/png, optional)", 
    type=["jpg", "jpeg", "png"]
)

video_path, bg_path = None, None

if video_file:
    # Save video to temp file
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as vf:
        vf.write(video_file.read())
        video_path = vf.name
    st.video(video_path)

if bg_file:
    # Save background to temp file
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as bf:
        bf.write(bg_file.read())
        bg_path = bf.name
    background_img = Image.open(bg_path).convert("RGB")
    st.image(background_img, caption="Background Image")
elif video_file:
    # Fallback: green background sized to video or 1280x720
    try:
        import cv2
        cap = cv2.VideoCapture(video_path)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) or 1280
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) or 720
        cap.release()
    except Exception:
        width, height = 1280, 720
    background_img = Image.new("RGB", (width, height), (0,255,0))
    st.info("No background image uploaded. Using default green background.")

# --- Processing ---
if video_path and background_img:
    if st.button("🚀 Run Background Replacement"):
        st.session_state.progress_bar = st.progress(0)
        st.session_state.progress_text = st.empty()
        with st.spinner("Processing video... (May take several minutes for long videos)"):
            try:
                out_path = process_two_stage(
                    video_path,
                    background_img,
                    progress=progress_callback
                )
                st.success("✅ Done! Here is your composited video:")
                st.video(str(out_path))
                with open(out_path, "rb") as f:
                    st.download_button(
                        label="💾 Download Result",
                        data=f,
                        file_name="backgroundfx_result.mp4",
                        mime="video/mp4"
                    )
            except Exception as e:
                st.error(f"❌ Processing failed: {e}")

        # Clean up temp files (best effort)
        try:
            if video_path and os.path.exists(video_path):
                os.unlink(video_path)
        except Exception: pass
        try:
            if bg_path and os.path.exists(bg_path):
                os.unlink(bg_path)
        except Exception: pass

# HF Spaces tip
st.caption("Built for Hugging Face T4 Spaces • CUDA & GPU acceleration required • v2025.10")