app.py

import mimetypes
import os
import tempfile
import time
from typing import List

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

from src.components.each_necklace_video_gen import EachVideoCreator
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 = "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:
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"Video file not found: {video_path}")

        content_type = mimetypes.guess_type(video_path)[0] or 'video/mp4'
        file_name = os.path.basename(video_path)

        options = {
            "content-type": content_type,
            "x-upsert": "true",
            "cache-control": "max-age=3600"
        }

        with open(video_path, 'rb') as f:
            file_data = f.read()

            supabase.storage.from_(bucket_name).upload(
                path=file_name,
                file=file_data,
                file_options=options
            )

            public_url = supabase.storage.from_(bucket_name).get_public_url(file_name)

            response = requests.head(public_url)
            if response.status_code != 200:
                raise Exception(f"Upload verification failed: {response.status_code}")

            return public_url

    except Exception as e:
        print(f"Error uploading to Supabase: {str(e)}")
        return None


def download_image(url):
    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-single-product-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

        )

        video_creator.create_final_video()
        video_creation_time = time.time() - start_time - file_download_time

        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)


class EachNecklaceVideoGeneratorRequest(BaseModel):
    intro_video_path: str = "JewelMirror_intro.mp4"
    font_path: str = "PlayfairDisplay-VariableFont_wght.ttf"
    background_audio_path: str = "TraditionalIndianVlogMusic.mp3"
    image_display_duration: float = 2.5
    fps: int = 10
    necklace_title: List[str]
    nto_image_title: List[str]
    nto_cto_image_title: List[str]
    makeup_image_title: List[str]
    necklace_images: List[str]
    necklace_try_on_output_images: List[List[str]]
    clothing_output_images: List[List[str]]
    makeup_output_images: List[List[str]]


@app.post("/create-multiple-product-video/")
async def create_video(request: EachNecklaceVideoGeneratorRequest):
    start_time = time.time()
    try:
        def process_images(image_list):
            if isinstance(image_list, list):
                return [process_images(item) for item in image_list]
            return download_image(image_list)

        temp_files = {
            'necklaces': process_images(request.necklace_images),
            'necklace_tryon': process_images(request.necklace_try_on_output_images),
            'clothing': process_images(request.clothing_output_images),
            'makeup': process_images(request.makeup_output_images)
        }
        file_download_time = time.time() - start_time

        def verify_files(files):
            if isinstance(files, list):
                return all(verify_files(item) for item in files)  # Recurse for nested lists
            return files is not None  # Check single file

        if not all(verify_files(files) for files in temp_files.values()):
            return JSONResponse(
                content={"status": "error", "message": "Failed to download all required images."},
                status_code=400
            )

        # Prepare paths
        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 = EachVideoCreator(
            intro_video_path=intro_path,
            necklace_image=temp_files['necklaces'],
            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_title=request.necklace_title,
            nto_title=request.nto_image_title,

            cto_title=request.nto_cto_image_title,
            makeup_title=request.makeup_image_title

        )

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

        # Upload video 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
            }
        }

        def cleanup_files(files):
            if isinstance(files, list):
                for item in files:
                    cleanup_files(item)
            elif os.path.exists(files):
                os.unlink(files)

        if os.path.exists(output_path):
            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)