import os
import platform
import uuid
import shutil
from pydub import AudioSegment
import spaces
import torch
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from transformers import pipeline
from huggingface_hub import snapshot_download
from examples.get_examples import get_examples
from src.facerender.pirender_animate import AnimateFromCoeff_PIRender
from src.utils.preprocess import CropAndExtract
from src.test_audio2coeff import Audio2Coeff
from src.facerender.animate import AnimateFromCoeff
from src.generate_batch import get_data
from src.generate_facerender_batch import get_facerender_data
from src.utils.init_path import init_path

checkpoint_path = 'checkpoints'
config_path = 'src/config'
device = "cuda" if torch.cuda.is_available() else "mps" if platform.system() == 'Darwin' else "cpu"

os.environ['TORCH_HOME'] = checkpoint_path
snapshot_download(repo_id='vinthony/SadTalker-V002rc',
                  local_dir=checkpoint_path, local_dir_use_symlinks=True)

app = FastAPI()
app.mount("/results", StaticFiles(directory="results"), name="results")
templates = Jinja2Templates(directory="templates")

def mp3_to_wav(mp3_filename, wav_filename, frame_rate):
    AudioSegment.from_file(file=mp3_filename).set_frame_rate(
        frame_rate).export(wav_filename, format="wav")

def get_pose_style_from_audio(audio_path):
    """Determines pose style based on audio emotion using a pre-trained model."""
    # Load the pre-trained emotion recognition model
    emotion_recognizer = pipeline("sentiment-analysis")

    # Analyze the audio emotion
    results = emotion_recognizer(audio_path)
    emotion = results[0]["label"]

    # Map emotion to pose style (you can adjust these mappings)
    pose_style_mapping = {
        "POSITIVE": 15,  # Happy
        "NEGATIVE": 35,  # Sad
        "NEUTRAL": 0,   # Normal
        # Add more emotion mappings as needed
    }

    return pose_style_mapping.get(emotion, 0)  # Default to neutral pose if unknown

@spaces.GPU(duration=0)
def generate_video(source_image: str, driven_audio: str, preprocess: str = 'crop', still_mode: bool = False,
                   use_enhancer: bool = False, batch_size: int = 1, size: int = 256, 
                   facerender: str = 'facevid2vid', exp_scale: float = 1.0, use_ref_video: bool = False,
                   ref_video: str = None, ref_info: str = None, use_idle_mode: bool = False,
                   length_of_audio: int = 0, use_blink: bool = True, result_dir: str = './results/') -> str:
    # Initialize models and paths
    sadtalker_paths = init_path(
        checkpoint_path, config_path, size, False, preprocess)
    audio_to_coeff = Audio2Coeff(sadtalker_paths, device)
    preprocess_model = CropAndExtract(sadtalker_paths, device)
    animate_from_coeff = AnimateFromCoeff(sadtalker_paths, device) if facerender == 'facevid2vid' and device != 'mps' \
        else AnimateFromCoeff_PIRender(sadtalker_paths, device)

    # Create directories for saving results
    time_tag = str(uuid.uuid4())
    save_dir = os.path.join(result_dir, time_tag)
    os.makedirs(save_dir, exist_ok=True)
    input_dir = os.path.join(save_dir, 'input')
    os.makedirs(input_dir, exist_ok=True)

    # Process source image
    pic_path = os.path.join(input_dir, os.path.basename(source_image))
    shutil.move(source_image, input_dir)

    # Process driven audio
    if driven_audio and os.path.isfile(driven_audio):
        audio_path = os.path.join(input_dir, os.path.basename(driven_audio))
        if '.mp3' in audio_path:
            mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000)
            audio_path = audio_path.replace('.mp3', '.wav')
        else:
            shutil.move(driven_audio, input_dir)
    elif use_idle_mode:
        audio_path = os.path.join(
            input_dir, 'idlemode_'+str(length_of_audio)+'.wav')
        AudioSegment.silent(
            duration=1000*length_of_audio).export(audio_path, format="wav")
    else:
        assert use_ref_video and ref_info == 'all'

    # Process reference video
    if use_ref_video and ref_info == 'all':
        ref_video_videoname = os.path.splitext(os.path.split(ref_video)[-1])[0]
        audio_path = os.path.join(save_dir, ref_video_videoname+'.wav')
        os.system(
            f"ffmpeg -y -hide_banner -loglevel error -i {ref_video} {audio_path}")
        ref_video_frame_dir = os.path.join(save_dir, ref_video_videoname)
        os.makedirs(ref_video_frame_dir, exist_ok=True)
        ref_video_coeff_path, _, _ = preprocess_model.generate(
            ref_video, ref_video_frame_dir, preprocess, source_image_flag=False)
    else:
        ref_video_coeff_path = None

    # Preprocess source image
    first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
    os.makedirs(first_frame_dir, exist_ok=True)
    first_coeff_path, crop_pic_path, crop_info = preprocess_model.generate(
        pic_path, first_frame_dir, preprocess, True, size)
    if first_coeff_path is None:
        raise AttributeError("No face is detected")

    # Determine reference coefficients
    ref_pose_coeff_path, ref_eyeblink_coeff_path = None, None
    if use_ref_video:
        if ref_info == 'pose':
            ref_pose_coeff_path = ref_video_coeff_path
        elif ref_info == 'blink':
            ref_eyeblink_coeff_path = ref_video_coeff_path
        elif ref_info == 'pose+blink':
            ref_pose_coeff_path = ref_eyeblink_coeff_path = ref_video_coeff_path
    else:
        ref_pose_coeff_path = ref_eyeblink_coeff_path = None

    # Generate coefficients from audio or reference video
    if use_ref_video and ref_info == 'all':
        coeff_path = ref_video_coeff_path
    else:
        batch = get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path=ref_eyeblink_coeff_path,
                         still=still_mode, idlemode=use_idle_mode, length_of_audio=length_of_audio, use_blink=use_blink)
        
        # Get pose style from audio
        pose_style = get_pose_style_from_audio(audio_path) 
        
        coeff_path = audio_to_coeff.generate(
            batch, save_dir, pose_style, ref_pose_coeff_path)

    # Generate video from coefficients
    data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode,
                               preprocess=preprocess, size=size, expression_scale=exp_scale, facemodel=facerender)
    return_path = animate_from_coeff.generate(data, save_dir, pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None,
                                              preprocess=preprocess, img_size=size)
    video_name = data['video_name']
    print(f'The generated video is named {video_name} in {save_dir}')

    return return_path

@app.post("/generate")
async def generate_video_api(source_image: UploadFile = File(...), driven_audio: UploadFile = File(None),
                            preprocess: str = Form('crop'), still_mode: bool = Form(False),
                            use_enhancer: bool = Form(False), batch_size: int = Form(1), size: int = Form(256), 
                            facerender: str = Form('facevid2vid'), exp_scale: float = Form(1.0),
                            use_ref_video: bool = Form(False), ref_video: UploadFile = File(None),
                            ref_info: str = Form(None), use_idle_mode: bool = Form(False),
                            length_of_audio: int = Form(0), use_blink: bool = Form(True), result_dir: str = Form('./results/')):
    # Save the uploaded files temporarily
    temp_source_image_path = f"temp/{source_image.filename}"
    os.makedirs("temp", exist_ok=True)
    with open(temp_source_image_path, "wb") as buffer:
        shutil.copyfileobj(source_image.file, buffer)

    if driven_audio:
        temp_driven_audio_path = f"temp/{driven_audio.filename}"
        with open(temp_driven_audio_path, "wb") as buffer:
            shutil.copyfileobj(driven_audio.file, buffer)
    else:
        temp_driven_audio_path = None

    if ref_video:
        temp_ref_video_path = f"temp/{ref_video.filename}"
        with open(temp_ref_video_path, "wb") as buffer:
            shutil.copyfileobj(ref_video.file, buffer)
    else:
        temp_ref_video_path = None

    # Generate the video
    video_path = generate_video(
        source_image=temp_source_image_path,
        driven_audio=temp_driven_audio_path,
        preprocess=preprocess,
        still_mode=still_mode,
        use_enhancer=use_enhancer,
        batch_size=batch_size,
        size=size,
        facerender=facerender,
        exp_scale=exp_scale,
        use_ref_video=use_ref_video,
        ref_video=temp_ref_video_path,
        ref_info=ref_info,
        use_idle_mode=use_idle_mode,
        length_of_audio=length_of_audio,
        use_blink=use_blink,
        result_dir=result_dir
    )

    # Clean up temporary files
    shutil.rmtree("temp")

    # Return the generated video file
    return FileResponse(video_path)


@app.get("/")
async def root(request):
    return templates.TemplateResponse("index.html", {"request": request})

# HTML Template (`templates/index.html`)
html = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>SadTalker API</title>
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/bootstrap@4.6.2/dist/css/bootstrap.min.css">
    <script src="https://cdn.jsdelivr.net/npm/jquery@3.5.1/dist/jquery.slim.min.js"></script>
    <script src="https://cdn.jsdelivr.net/npm/popper.js@1.16.1/dist/umd/popper.min.js"></script>
    <script src="https://cdn.jsdelivr.net/npm/bootstrap@4.6.2/dist/js/bootstrap.min.js"></script>
</head>
<body>
    <div class="container mt-5">
        <h1>SadTalker API</h1>
        <form method="POST" action="/generate" enctype="multipart/form-data">
            <div class="form-group">
                <label for="source_image">Source Image:</label>
                <input type="file" class="form-control-file" id="source_image" name="source_image" required>
            </div>
            <div class="form-group">
                <label for="driven_audio">Driving Audio:</label>
                <input type="file" class="form-control-file" id="driven_audio" name="driven_audio">
            </div>
            <div class="form-group">
                <label for="preprocess">Preprocess:</label>
                <select class="form-control" id="preprocess" name="preprocess">
                    <option value="crop">Crop</option>
                    <option value="resize">Resize</option>
                    <option value="full">Full</option>
                    <option value="extcrop">ExtCrop</option>
                    <option value="extfull">ExtFull</option>
                </select>
            </div>
            <div class="form-check">
                <input type="checkbox" class="form-check-input" id="still_mode" name="still_mode">
                <label class="form-check-label" for="still_mode">Still Mode</label>
            </div>
            <div class="form-check">
                <input type="checkbox" class="form-check-input" id="use_enhancer" name="use_enhancer">
                <label class="form-check-label" for="use_enhancer">Use GFPGAN Enhancer</label>
            </div>
            <div class="form-group">
                <label for="batch_size">Batch Size:</label>
                <input type="number" class="form-control" id="batch_size" name="batch_size" min="1" max="10" value="1">
            </div>
            <div class="form-group">
                <label for="size">Face Model Resolution:</label>
                <select class="form-control" id="size" name="size">
                    <option value="256">256</option>
                    <option value="512">512</option>
                </select>
            </div>
            <div class="form-group">
                <label for="facerender">Face Render:</label>
                <select class="form-control" id="facerender" name="facerender">
                    <option value="facevid2vid">FaceVid2Vid</option>
                    <option value="pirender">PIRender</option>
                </select>
            </div>
            <div class="form-group">
                <label for="exp_scale">Expression Scale:</label>
                <input type="number" class="form-control" id="exp_scale" name="exp_scale" min="0" max="3" step="0.1" value="1.0">
            </div>
            <div class="form-check">
                <input type="checkbox" class="form-check-input" id="use_ref_video" name="use_ref_video">
                <label class="form-check-label" for="use_ref_video">Use Reference Video</label>
            </div>
            <div class="form-group">
                <label for="ref_video">Reference Video:</label>
                <input type="file" class="form-control-file" id="ref_video" name="ref_video">
            </div>
            <div class="form-group">
                <label for="ref_info">Reference Video Information:</label>
                <select class="form-control" id="ref_info" name="ref_info">
                    <option value="pose">Pose</option>
                    <option value="blink">Blink</option>
                    <option value="pose+blink">Pose + Blink</option>
                    <option value="all">All</option>
                </select>
            </div>
            <div class="form-check">
                <input type="checkbox" class="form-check-input" id="use_idle_mode" name="use_idle_mode">
                <label class="form-check-label" for="use_idle_mode">Use Idle Animation</label>
            </div>
            <div class="form-group">
                <label for="length_of_audio">Length of Audio (seconds):</label>
                <input type="number" class="form-control" id="length_of_audio" name="length_of_audio" min="0" value="0">
            </div>
            <div class="form-check">
                <input type="checkbox" class="form-check-input" id="use_blink" name="use_blink" checked>
                <label class="form-check-label" for="use_blink">Use Eye Blink</label>
            </div>
            <button type="submit" class="btn btn-primary">Generate</button>
        </form>
    </div>
</body>
</html>
"""

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)