app.py

import os
import tempfile
import time

import requests
from fastapi import FastAPI
from pydantic import BaseModel
from starlette.responses import JSONResponse
from supabase import create_client

from src.components.vidgen import VideoCreator

supabase_url = os.getenv('SUPABASE_URL')
supabase_key = os.getenv('SUPABASE_KEY')
supabase = create_client(supabase_url, supabase_key)
app = FastAPI()

RESOURCES_DIR = "/app/resources"

os.makedirs(RESOURCES_DIR, exist_ok=True)
TEMP_VIDEO_DIR = f"{RESOURCES_DIR}/temp_video"

os.environ['MOVIEPY_TEMP_DIR'] = '/tmp/moviepy'


def upload_to_supabase(video_path, bucket_name="JewelmirrorVideoGeneration"):
    try:
        with open(video_path, 'rb') as f:
            file_name = os.path.basename(video_path)
            supabase.storage.from_(bucket_name).upload(file_name, f)

            public_url = supabase.storage.from_(bucket_name).get_public_url(file_name)
            return public_url
    except Exception as e:
        print(f"Error uploading to Supabase: {str(e)}")
        return None


def download_image(url):
    """
    Download an image from the web and save it as a temporary file.
    :param url: URL of the image.
    :return: Path to the temporary file.
    """
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()

        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        with open(temp_file.name, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return temp_file.name
    except Exception as e:
        print(f"Error downloading image from {url}: {str(e)}")
        return None


class VideoGenerator(BaseModel):
    necklace_image: str
    necklace_try_on_output_images: list[str]
    clothing_output_images: list[str]
    makeup_output_images: list[str]
    intro_video_path: str = "JewelMirror_intro.mp4"
    background_audio_path: str = "TraditionalIndianVlogMusic.mp3"
    font_path: str = "PlayfairDisplay-VariableFont_wght.ttf"
    image_display_duration: float = 2.5
    fps: int = 30
    necklace_title: str = "Necklace Try-On"
    clothing_title: str = "Clothing Try-On"
    makeup_title: str = "Makeup Try-On"
    necklace_image_title: str = "Necklace Preview"
    nto_image_title: str = "Necklace Try-On"
    nto_cto_image_title: str = "Clothing Try-On"
    makeup_image_title: str = "Makeup Try-On"


@app.post("/create-video/")
async def create_video(request: VideoGenerator):
    start_time = time.time()
    try:
        temp_files = {
            'necklace': download_image(request.necklace_image),
            'necklace_tryon': [download_image(url) for url in request.necklace_try_on_output_images],
            'clothing': [download_image(url) for url in request.clothing_output_images],
            'makeup': [download_image(url) for url in request.makeup_output_images]
        }
        file_download_time = time.time() - start_time

        if any(file is None for files in temp_files.values() for file in
               (files if isinstance(files, list) else [files])):
            return JSONResponse(content={"status": "error", "message": "Failed to download all required images."},
                                status_code=400)

        intro_path = f"{RESOURCES_DIR}/intro/{request.intro_video_path}"
        output_path = f"{TEMP_VIDEO_DIR}/video_{os.urandom(8).hex()}.mp4"
        font_path = f"{RESOURCES_DIR}/fonts/{request.font_path}"
        audio_path = f"{RESOURCES_DIR}/audio/{request.background_audio_path}"

        video_creator = VideoCreator(
            intro_video_path=intro_path,
            necklace_image=temp_files['necklace'],
            nto_outputs=temp_files['necklace_tryon'],
            nto_cto_outputs=temp_files['clothing'],
            makeup_outputs=temp_files['makeup'],
            font_path=font_path,
            output_path=output_path,
            audio_path=audio_path,
            image_display_duration=request.image_display_duration,
            fps=request.fps,
            necklace_image_title=request.necklace_image_title,
            nto_image_title=request.nto_image_title,
            nto_cto_image_title=request.nto_cto_image_title,
            makeup_image_title=request.makeup_image_title

        )

        # Generate video
        video_creator.create_final_video()
        video_creation_time = time.time() - start_time - file_download_time

        # Upload to Supabase
        url = upload_to_supabase(video_path=output_path)
        supabase_upload_time = time.time() - start_time - file_download_time - video_creation_time

        response = {
            "status": "success",
            "message": "Video created successfully.",
            "video_url": url,
            "timings": {
                "file_download_time": file_download_time,
                "video_creation_time": video_creation_time,
                "supabase_upload_time": supabase_upload_time
            }
        }

        for files in temp_files.values():
            if isinstance(files, list):
                for file in files:
                    if os.path.exists(file):
                        os.unlink(file)
            elif os.path.exists(files):
                os.unlink(files)

        os.remove(output_path)
        return JSONResponse(content=response, status_code=200)

    except Exception as e:
        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)


@app.get("/resources")
async def get_infromation():
    music = os.listdir(RESOURCES_DIR + "/audio")
    fonts = os.listdir(RESOURCES_DIR + "/fonts")
    intro = os.listdir(RESOURCES_DIR + "/intro")

    json = {"music": music, "fonts": fonts, "intro": intro}
    return JSONResponse(content=json, status_code=200)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)