import os
import subprocess
from pathlib import Path
from PIL import Image
import streamlit as st

# ---- CONFIG ----
st.set_page_config(
    page_title="Streamlit iCodeIdoia",
    page_icon="images/ilpicon1.png",
    layout="wide",
    initial_sidebar_state="expanded"
)

st.image("images/banner.jpg")

# ---- PATHS ----
FRAME1 = Path("demo/frame1.png")
FRAME2 = Path("demo/frame2.png")
TARGET_DIR = Path("/home/user/app/output/")
PALETTE_PNG = TARGET_DIR / "palette.png"
OUTPUT_GIF = TARGET_DIR / "output.gif"

os.makedirs(TARGET_DIR, exist_ok=True)

# ---- FUNCTION ----
def interpolate_image(img_a_path: str, img_b_path: str) -> str:
    subprocess.run([
        "python3", "inference_img.py",
        "--img", str(img_a_path), str(img_b_path),
        "--exp", "4"
    ], check=True)

    subprocess.run([
        "ffmpeg", "-y", "-r", "14", "-f", "image2",
        "-i", f"{TARGET_DIR}/img%d.png",
        "-vf", "palettegen=stats_mode=single",
        "-frames:v", "1",
        str(PALETTE_PNG)
    ], check=True)

    subprocess.run([
        "ffmpeg", "-y", "-r", "14", "-f", "image2",
        "-i", f"{TARGET_DIR}/img%d.png",
        "-i", str(PALETTE_PNG),
        "-lavfi", "paletteuse",
        str(OUTPUT_GIF)
    ], check=True)

    return str(OUTPUT_GIF)

# ---- TABS ----
tab1, tab2 = st.tabs(["Demo", "Upload your images"])

with tab1:
    st.subheader("Demo: Preloaded images")
    st.image(str(FRAME1), caption="Image A")
    st.image(str(FRAME2), caption="Image B")

    if st.button("Run Interpolation Demo"):
        gif_path = interpolate_image(FRAME1, FRAME2)
        st.image(gif_path, caption="Interpolated GIF")
        st.text(f"Output path: {gif_path}")

with tab2:
    st.subheader("Upload any two images")
    uploaded_a = st.file_uploader("Upload Image A", type=["png", "jpg", "jpeg"])
    uploaded_b = st.file_uploader("Upload Image B", type=["png", "jpg", "jpeg"])

    if uploaded_a and uploaded_b:
        temp_a = TARGET_DIR / "user_a.png"
        temp_b = TARGET_DIR / "user_b.png"
        Image.open(uploaded_a).save(temp_a)
        Image.open(uploaded_b).save(temp_b)

        if st.button("Run Interpolation"):
            gif_path = interpolate_image(temp_a, temp_b)
            st.image(gif_path, caption="Interpolated GIF")
            st.text(f"Output path: {gif_path}")