main.py

from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse
import tempfile
import shutil
import os
import logging
import traceback
from moviepy.editor import (
    VideoFileClip,
    concatenate_videoclips,
    AudioFileClip,
    CompositeAudioClip,
    concatenate_audioclips,
    vfx,
)
import io
from PIL import Image

# Configure logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("main")

# Enable MoviePy verbose logging
from moviepy.config import change_settings
change_settings({"VERBOSE": True})

# Safe video clip loader
def safe_video_clip(path, target_width=1280, reference_aspect_ratio=None):
    try:
        clip = VideoFileClip(path)
        if clip.reader is None or clip.duration is None or clip.duration <= 0:
            clip.close()
            raise ValueError(f"Invalid or empty video file: {os.path.basename(path)}")
        logger.debug(f"Loaded video: {path}, duration: {clip.duration}s, resolution: {clip.size}, aspect ratio: {clip.w / clip.h}")

        # Get video aspect ratio
        current_aspect_ratio = clip.w / clip.h

        # Use reference aspect ratio if provided, else use current video's aspect ratio
        target_aspect_ratio = reference_aspect_ratio if reference_aspect_ratio else current_aspect_ratio

        # Check if aspect ratios are compatible (within a small tolerance)
        if reference_aspect_ratio and abs(current_aspect_ratio - target_aspect_ratio) > 0.05:
            logger.warning(f"Video {path} has different aspect ratio ({current_aspect_ratio:.2f}) than reference ({target_aspect_ratio:.2f}). Adjusting to match.")
        
        # Calculate target resolution to preserve aspect ratio
        target_height = int(target_width / target_aspect_ratio)
        # Ensure height is even (required by libx264 codec)
        target_height = target_height if target_height % 2 == 0 else target_height + 1

        try:
            clip = clip.resize((target_width, target_height))
            logger.debug(f"Resized video {path} to {target_width}x{target_height}, aspect ratio: {target_width / target_height}")
        except Exception as e:
            logger.warning(f"Failed to resize video {path}: {e}")
            # Continue without resizing if it fails (may cause concatenation issues)
        
        return clip, target_aspect_ratio
    except Exception as e:
        raise ValueError(f"Failed to load video file '{os.path.basename(path)}': {str(e)}")

# Safe audio clip loader
def safe_audio_clip(path):
    try:
        clip = AudioFileClip(path)
        if clip.reader is None or clip.duration is None or clip.duration <= 0:
            clip.close()
            raise ValueError(f"Invalid or empty audio file: {os.path.basename(path)}")
        logger.debug(f"Loaded audio: {path}, duration: {clip.duration}s")
        return clip
    except Exception as e:
        raise ValueError(f"Failed to load audio file '{os.path.basename(path)}': {str(e)}")

def merge_videos_and_audios(video_paths, audio_paths, output_path, temp_dir, orig_vol=1.0, music_vol=0.5, audio_mode="concatenate"):
    logger.debug(f"Merging {len(video_paths)} videos and {len(audio_paths)} audios to {output_path}")
    
    # Load video clips
    video_clips = []
    reference_aspect_ratio = None
    for vp in video_paths:
        try:
            clip, aspect_ratio = safe_video_clip(vp, reference_aspect_ratio=reference_aspect_ratio)
            if reference_aspect_ratio is None:
                reference_aspect_ratio = aspect_ratio  # Set reference from first valid video
            video_clips.append(clip)
        except Exception as e:
            logger.warning(f"Skipping invalid video file {vp}: {e}")

    if not video_clips:
        raise ValueError("No valid video files provided")

    # Load audio clips
    audio_clips = []
    for ap in audio_paths:
        try:
            clip = safe_audio_clip(ap)
            clip = clip.volumex(music_vol)
            audio_clips.append(clip)
        except Exception as e:
            logger.warning(f"Skipping invalid audio file {ap}: {e}")

    # Concatenate video clips
    final_video = concatenate_videoclips(video_clips, method="compose")
    logger.debug(f"Concatenated video duration: {final_video.duration}s, resolution: {final_video.size}")

    # Handle original video audio
    original_audio = None
    if final_video.audio:
        original_audio = final_video.audio.volumex(orig_vol)
        logger.debug(f"Original video audio preserved, duration: {original_audio.duration}s")

    # Combine audio clips
    final_audio = None
    if audio_clips:
        if audio_mode == "concatenate":
            final_audio = concatenate_audioclips(audio_clips)
            logger.debug(f"Concatenated audio duration: {final_audio.duration}s")
        else:
            final_audio = CompositeAudioClip(audio_clips)
            logger.debug(f"Composite audio duration: {final_audio.duration}s")
        
        # Match audio duration to video
        if final_audio and final_audio.duration < final_video.duration:
            logger.debug(f"Looping audio to match video duration: {final_video.duration}s")
            final_audio = final_audio.fx(vfx.loop, duration=final_video.duration)
        elif final_audio and final_audio.duration > final_video.duration:
            logger.debug(f"Trimming audio to match video duration: {final_video.duration}s")
            final_audio = final_audio.subclip(0, final_video.duration)
        
        if original_audio:
            final_audio = CompositeAudioClip([final_audio, original_audio])
            logger.debug(f"Combined audio (original + uploaded) duration: {final_audio.duration}s")

    # Set the audio for the final video
    if final_audio:
        final_video = final_video.set_audio(final_audio)
    elif original_audio:
        final_video = final_video.set_audio(original_audio)
    else:
        logger.warning("No audio (uploaded or original) available for the final video")

    # Write output
    logger.debug(f"Writing output to {output_path}")
    try:
        final_video.write_videofile(
            output_path,
            codec="libx264",
            audio_codec="aac",
            temp_audiofile=os.path.join(temp_dir, "temp-audio.m4a"),
            remove_temp=True,
            verbose=True,
            logger='bar'
        )
    except Exception as e:
        logger.error(f"MoviePy failed to write output: {str(e)}")
        raise

    # Close clips
    final_video.close()
    for vc in video_clips:
        vc.close()
    for ac in audio_clips:
        ac.close()
    if original_audio:
        original_audio.close()

app = FastAPI()

@app.post("/merge")
async def merge_endpoint(
    files: list[UploadFile] = File(...),
    orig_vol: float = Form(1.0),
    music_vol: float = Form(0.5),
    audio_mode: str = Form("concatenate"),
    background_tasks: BackgroundTasks = BackgroundTasks(),
):
    temp_dir = tempfile.mkdtemp()
    logger.debug(f"Created temporary directory: {temp_dir}")
    
    try:
        saved_files = []
        for uploaded_file in files:
            if not uploaded_file.filename:
                logger.warning("Skipping file with no filename")
                continue
            file_path = os.path.join(temp_dir, uploaded_file.filename)
            with open(file_path, "wb") as out_file:
                content = await uploaded_file.read()
                if not content:
                    logger.warning(f"Empty file uploaded: {uploaded_file.filename}")
                    continue
                out_file.write(content)
            saved_files.append(file_path)
            logger.debug(f"Saved file: {file_path}")

        video_files = [f for f in saved_files if f.lower().endswith(".mp4")]
        audio_files = [f for f in saved_files if f.lower().endswith((".mp3", ".wav", ".aac", ".m4a", ".ogg"))]
        logger.debug(f"Video files: {video_files}")
        logger.debug(f"Audio files: {audio_files}")

        if len(video_files) < 1:
            raise HTTPException(status_code=400, detail="Please upload at least one valid video file")

        output_path = os.path.join(temp_dir, "merged_output.mp4")
        merge_videos_and_audios(video_files, audio_files, output_path, temp_dir, orig_vol, music_vol, audio_mode)

        if not os.path.exists(output_path):
            logger.error(f"Output file not found: {output_path}")
            raise HTTPException(status_code=500, detail="Failed to create merged video file")

        logger.debug(f"Reading output file: {output_path}")
        with open(output_path, "rb") as file:
            content = file.read()
            if not content:
                logger.error(f"Output file is empty: {output_path}")
                raise HTTPException(status_code=500, detail="Output file is empty")

        background_tasks.add_task(shutil.rmtree, temp_dir, ignore_errors=True)
        logger.debug("Sending streaming response")
        return StreamingResponse(
            io.BytesIO(content),
            media_type="video/mp4",
            headers={"Content-Disposition": 'attachment; filename="merged_output.mp4"'}
        )

    except Exception as e:
        error_msg = f"Error: {str(e)}\n\n{traceback.format_exc()}"
        logger.error(error_msg)
        raise HTTPException(status_code=500, detail=error_msg)

    finally:
        logger.debug(f"Cleaning up temporary directory: {temp_dir}")
        shutil.rmtree(temp_dir, ignore_errors=True)

# Log public URL
def log_api_url():
    url = os.getenv("SPACE_PUBLIC_URL")
    if url:
        logger.info(f"API is available at: {url}/merge")
    else:
        logger.info("SPACE_PUBLIC_URL environment variable not found")

log_api_url()