endpoints/rumble.py

from fastapi import Body, HTTPException
from fastapi.responses import StreamingResponse
from typing import Optional
from fastapi import APIRouter
from generate_rumble_clips import generate_rumble_clips
from endpoints.utils import remove_content_from_dir
from endpoints.utils import rumble_link_scraper
import io
import zipfile
import pydantic
import os

router = APIRouter()

class ClipData(pydantic.BaseModel):
	start: str
	end: str
	name: str
	clip_link: Optional[str] = None

class RumbleData(pydantic.BaseModel):
	video_url: str
	color: str
	size: int
	font: str
	credit: str
	credit_size: int
	clips: list[ClipData]


@router.post("/rumble")
async def create_rumble_video(request_data: RumbleData = Body(...)):

	# Extract values from the request_data dictionary
	video_url = request_data.video_url
	color = request_data.color
	size = request_data.size
	font = request_data.font
	credit = request_data.credit
	credit_size = request_data.credit_size
	clips = request_data.clips

	options = {
		"video_url": video_url,
		"clips": clips,
		"font_color": color,
		"font_size": size,
		"font_family": font,
		"border_size": 1,
		"credit": credit,
		"credit_size": credit_size,
	}

	try:
		# scrape the web for the rumble link
		options["video_url"] = None

		print(options["video_url"])

		# Create the directory if it doesn't exist
		if not os.path.exists("./data/rumble/clips"):
			os.makedirs("./data/rumble/clips")

		# Generate the clips
		clips = generate_rumble_clips(options)
		
		# Create a zip archive in memory
		zip_buffer = io.BytesIO()
		with zipfile.ZipFile(zip_buffer, "w") as zip_file:
			for clip in clips:
				zip_file.write(clip.clip_link, arcname=f"clips/{clip.name}.mp4")

		# Reset the buffer position to the beginning
		zip_buffer.seek(0)		

		# remove the clips directory content
		remove_content_from_dir("./data/rumble/clips")

		# Return the zip archive as a StreamingResponse
		return StreamingResponse(io.BytesIO(zip_buffer.read()), media_type="application/zip", headers={
			'Content-Disposition': 'attachment; filename="clips.zip"'
		})
	except Exception as e:
		# Handle exceptions appropriately
		raise HTTPException(status_code=500, detail=str(e))