import os, tempfile, uuid, asyncio, shutil, requests
from urllib.parse import urlparse
from fastapi import FastAPI, UploadFile, File, HTTPException, Form, WebSocket
from fastapi.responses import JSONResponse
from fastapi import APIRouter
from extensions import *
from main import *
from tts_api import *
from sadtalker_utils import *
import base64
from stt_api import *
from text_generation import *

router = APIRouter()

@router.post("/sadtalker")
async def create_video(source_image_file: UploadFile = File(...), driven_audio_file: UploadFile = File(...)):
    if not source_image_file: raise HTTPException(status_code=400, detail="Source image file is required")
    if not driven_audio_file: raise HTTPException(status_code=400, detail="Driven audio file is required")

    temp_source_image = tempfile.NamedTemporaryFile(suffix=os.path.splitext(source_image_file.filename)[1], delete=False)
    content_image = await source_image_file.read(); temp_source_image.write(content_image); source_image_path = temp_source_image.name
    temp_driven_audio = tempfile.NamedTemporaryFile(suffix=os.path.splitext(driven_audio_file.filename)[1], delete=False)
    content_audio = await driven_audio_file.read(); temp_driven_audio.write(content_audio); driven_audio_path = temp_driven_audio.name

    try:
        loop = asyncio.get_running_loop()
        output_path = await loop.run_in_executor(None, sadtalker_instance.test, source_image_path, driven_audio_path)
        video_base64 = None
        with open(output_path, 'rb') as video_file: video_bytes = video_file.read(); video_base64 = base64.b64encode(video_bytes).decode('utf-8')
        os.remove(output_path); return {"video_base64": video_base64, "mimetype": "video/mp4"}

    except Exception as e: raise HTTPException(status_code=500, detail=str(e))
    finally: os.remove(temp_source_image.name); os.remove(temp_driven_audio.name)

router = APIRouter()
router.add_api_route("/sadtalker", create_video, methods=["POST"])