app.py

import os
import shutil
import uuid
from datetime import datetime
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
from moviepy.editor import (
    ImageClip,
    ColorClip,
    TextClip,
    CompositeVideoClip,
    concatenate_videoclips,
    VideoFileClip,
)

# ---------- Simple "video generator" stub ----------
# Replace `generate_video_from_prompt` with your real model call.
# It creates a short MP4 with a first frame (optionally from a previous video),
# then overlays prompt text for a few seconds.

OUT_DIR = "outputs"
USED_DIR = os.path.join(OUT_DIR, "used")
TMP_DIR = "tmp"
os.makedirs(OUT_DIR, exist_ok=True)
os.makedirs(USED_DIR, exist_ok=True)
os.makedirs(TMP_DIR, exist_ok=True)

FPS = 24
W, H = 768, 432  # 16:9 HD-ish to keep files light
DURATION = 3.0   # seconds per generated clip for demo

def _solid_bg(color=(18, 18, 18)):
    return ColorClip(size=(W, H), color=color, duration=DURATION)

def _text_overlay(txt: str):
    # Use TextClip if ImageMagick is available; otherwise fallback to PIL.
    try:
        return TextClip(
            txt,
            fontsize=48,
            color="white",
            font="Arial-Bold",
            size=(W - 80, None),
            method="caption",
        ).set_position(("center", "center")).set_duration(DURATION)
    except Exception:
        # PIL fallback
        img = Image.new("RGBA", (W, H), (0, 0, 0, 0))
        draw = ImageDraw.Draw(img)
        # Try to load a font; fallback to default if not available on HF
        try:
            font = ImageFont.truetype("DejaVuSans-Bold.ttf", 48)
        except Exception:
            font = ImageFont.load_default()
        # simple multiline center
        lines = []
        words = txt.split()
        line = ""
        for w in words:
            test = (line + " " + w).strip()
            if draw.textlength(test, font=font) > (W - 80):
                lines.append(line)
                line = w
            else:
                line = test
        lines.append(line)

        total_h = sum(font.getbbox(l)[3] for l in lines) + (len(lines)-1)*8
        y = (H - total_h)//2
        for l in lines:
            w_px = draw.textlength(l, font=font)
            x = (W - w_px)//2
            draw.text((x, y), l, fill=(255,255,255,255), font=font)
            y += font.getbbox(l)[3] + 8
        pil_path = os.path.join(TMP_DIR, f"txt_{uuid.uuid4().hex}.png")
        img.save(pil_path)
        return ImageClip(pil_path, duration=DURATION).set_position(("center","center"))

def extract_last_frame_as_image(video_path: str) -> str:
    """Save last frame of video to an image file and return its path."""
    with VideoFileClip(video_path) as v:
        frame = v.get_frame(v.duration - 1.0 / max(1, v.fps))
    img = Image.fromarray(frame)
    frame_path = os.path.join(TMP_DIR, f"seed_{uuid.uuid4().hex}.png")
    img.save(frame_path)
    return frame_path

def generate_video_from_prompt(prompt: str, seed_frame_path: str | None) -> str:
    """
    Make a short demo MP4 using:
      - If seed_frame_path: start 0.5s with that still frame
      - Then a solid background + prompt text
    """
    # Clips to concatenate
    clips = []

    if seed_frame_path and os.path.exists(seed_frame_path):
        seed = ImageClip(seed_frame_path, duration=0.5).set_fps(FPS)
        clips.append(seed)

    bg = _solid_bg().set_fps(FPS)
    txt = _text_overlay(prompt)
    comp = CompositeVideoClip([bg, txt]).set_duration(DURATION).set_fps(FPS)
    clips.append(comp)

    final = concatenate_videoclips(clips, method="compose")
    out_name = f"gen_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.mp4"
    out_path = os.path.join(OUT_DIR, out_name)
    final.write_videofile(out_path, fps=FPS, codec="libx264", audio=False, verbose=False, logger=None)
    final.close()
    return out_path

def concat_used_videos(video_paths: list[str]) -> str:
    clips = [VideoFileClip(p) for p in video_paths]
    final = concatenate_videoclips(clips, method="compose")
    out_path = os.path.join(OUT_DIR, f"continuous_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.mp4")
    final.write_videofile(out_path, fps=FPS, codec="libx264", audio=False, verbose=False, logger=None)
    for c in clips:
        c.close()
    return out_path

def zip_used_videos(video_paths: list[str]) -> str:
    # Copy into a temp folder to zip cleanly
    stamp = datetime.utcnow().strftime('%Y%m%d_%H%M%S')
    pack_dir = os.path.join(TMP_DIR, f"used_{stamp}_{uuid.uuid4().hex[:6]}")
    os.makedirs(pack_dir, exist_ok=True)
    for p in video_paths:
        shutil.copy(p, pack_dir)
    zip_base = os.path.join(OUT_DIR, f"used_{stamp}")
    shutil.make_archive(zip_base, "zip", pack_dir)
    shutil.rmtree(pack_dir, ignore_errors=True)
    return f"{zip_base}.zip"

# ---------- Gradio App ----------
with gr.Blocks(css=".grow {flex: 1}") as demo:
    gr.Markdown("# Continuous Video Prompt → Use → Chain → Download")

    # Session state
    state_used_paths = gr.State([])         # list[str]
    state_seed_frame = gr.State(None)       # str | None
    state_current_path = gr.State(None)     # str | None

    with gr.Row():
        prompt = gr.Textbox(
            label="Prompt",
            placeholder="Describe your next shot…",
            lines=2,
            autofocus=True,
        )
    with gr.Row(equal_height=True):
        video_out = gr.Video(label="Video Output", interactive=False).style(height=360)
    with gr.Row():
        btn_generate = gr.Button("Generate", variant="primary")
        btn_use = gr.Button("Use (chain this)", variant="secondary")
        btn_download = gr.Button("Download (A+B+C & ZIP)", variant="secondary")
        btn_reset = gr.Button("Reset Session", variant="stop")

    files_out = gr.Files(label="Downloads (concatenated MP4 + ZIP of used clips)", height=100)

    # ---- Handlers ----
    def do_generate(prompt_text, seed_frame_path):
        if not prompt_text or not prompt_text.strip():
            return None, None
        out_path = generate_video_from_prompt(prompt_text.strip(), seed_frame_path)
        return out_path, out_path  # gr.Video path AND state_current_path

    btn_generate.click(
        do_generate,
        inputs=[prompt, state_seed_frame],
        outputs=[video_out, state_current_path],
    )

    def do_use(current_path, used_paths):
        """
        Save current_path to used list, extract its last frame as the next seed.
        """
        if not current_path or not os.path.exists(current_path):
            # no-op if nothing to use
            return used_paths, gr.update(interactive=True), None

        # Append to used list
        new_used = list(used_paths)
        if current_path not in new_used:
            new_used.append(current_path)

        # Extract last frame for next generation seed
        next_seed = extract_last_frame_as_image(current_path)
        return new_used, gr.update(interactive=True), next_seed

    btn_use.click(
        do_use,
        inputs=[state_current_path, state_used_paths],
        outputs=[state_used_paths, prompt, state_seed_frame],
    )

    def do_download(used_paths):
        """
        Build concatenated video (A+B+C) and a ZIP of used clips.
        Returns list of two files for the Files component.
        """
        if not used_paths:
            return []
        concat_path = concat_used_videos(used_paths)
        zip_path = zip_used_videos(used_paths)
        return [concat_path, zip_path]

    btn_download.click(
        do_download,
        inputs=[state_used_paths],
        outputs=[files_out],
    )

    def do_reset():
        # Clear session state and temp
        try:
            for f in os.listdir(TMP_DIR):
                fp = os.path.join(TMP_DIR, f)
                if os.path.isfile(fp):
                    os.remove(fp)
        except Exception:
            pass
        return None, [], None, None, gr.update(value=None), gr.update(value=[])

    btn_reset.click(
        do_reset,
        inputs=None,
        outputs=[state_seed_frame, state_used_paths, state_current_path, prompt, video_out, files_out],
    )

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