import os
import sys

# Add SadTalker to path
sys.path.append('SadTalker')

import gradio as gr
from SadTalker.src.gradio_demo import SadTalker

def sadtalker_demo():
    """Initialize and return the SadTalker Gradio interface"""
    
    checkpoint_path = 'SadTalker/checkpoints'
    config_path = 'SadTalker/src/config'
    
    sad_talker = SadTalker(checkpoint_path, config_path, lazy_load=True)

    with gr.Blocks(analytics_enabled=False, title="SadTalker - Talking Head Generator") as sadtalker_interface:
        gr.Markdown("""
        <div align='center'> 
            <h2>😭 SadTalker: Audio-Driven Talking Face Animation</h2>
            <p>Upload an image and audio to generate a talking head video</p>
            <a href='https://arxiv.org/abs/2211.12194'>Paper</a> &nbsp;&nbsp;
            <a href='https://sadtalker.github.io'>Homepage</a> &nbsp;&nbsp;
            <a href='https://github.com/OpenTalker/SadTalker'>GitHub</a>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(variant='panel'):
                with gr.Tabs():
                    with gr.TabItem('Upload Image'):
                        source_image = gr.Image(
                            label="Source Image", 
                            source="upload", 
                            type="filepath"
                        )

                with gr.Tabs():
                    with gr.TabItem('Upload Audio'):
                        driven_audio = gr.Audio(
                            label="Input Audio", 
                            source="upload", 
                            type="filepath"
                        )
                            
            with gr.Column(variant='panel'): 
                with gr.Tabs():
                    with gr.TabItem('Settings'):
                        gr.Markdown("Need help? Visit our [best practice page](https://github.com/OpenTalker/SadTalker/blob/main/docs/best_practice.md)")
                        
                        pose_style = gr.Slider(
                            minimum=0, 
                            maximum=46, 
                            step=1, 
                            label="Pose Style", 
                            value=0
                        )
                        size_of_image = gr.Radio(
                            [256, 512], 
                            value=256, 
                            label='Face Model Resolution'
                        )
                        preprocess_type = gr.Radio(
                            ['crop', 'resize', 'full', 'extcrop', 'extfull'], 
                            value='crop', 
                            label='Preprocess Type'
                        )
                        is_still_mode = gr.Checkbox(
                            label="Still Mode (less head motion)"
                        )
                        batch_size = gr.Slider(
                            label="Batch Size", 
                            step=1, 
                            minimum=1,
                            maximum=10, 
                            value=2
                        )
                        enhancer = gr.Checkbox(
                            label="GFPGAN Face Enhancer"
                        )
                        
                        submit = gr.Button('Generate Video', variant='primary')
                        
                with gr.Tabs():
                    with gr.TabItem('Generated Video'):
                        gen_video = gr.Video(label="Generated Video", format="mp4")

        submit.click(
            fn=sad_talker.test, 
            inputs=[
                source_image,
                driven_audio,
                preprocess_type,
                is_still_mode,
                enhancer,
                batch_size,                            
                size_of_image,
                pose_style
            ], 
            outputs=[gen_video]
        )

    return sadtalker_interface

if __name__ == "__main__":
    demo = sadtalker_demo()
    demo.queue()
    demo.launch()