import gradio as gr
from gradio_client import Client

PASSWORD = "071295"

def get_speech(text, voice):
    try:
        client = Client("sysf/vixtts-demo")
        result = client.predict(
            text,  # Changed from undefined 'prompt' to 'text'
            language="vi",
            audio_file=voice,  # Changed from undefined 'audio_file_pth' to 'voice'
            normalize_text=True,
            api_name="/predict"
        )
        print(result)
        return result
    except Exception as e:
        raise gr.Error(f"Error in get_speech: {str(e)}")

def get_dreamtalk(image_in, speech):
    try:
        client = Client("https://fffiloni-dreamtalk.hf.space/")
        result = client.predict(
            speech,  # filepath in 'Audio input' Audio component
            image_in,  # filepath in 'Image' Image component
            "M030_front_neutral_level1_001.mat",  # Literal[...] in 'emotional style' Dropdown component
            api_name="/infer"
        )
        print(result)
        return result['video']
    except Exception as e:
        raise gr.Error(f"Error in get_dreamtalk: {str(e)}. Image may not contain any face.")

def pipe(text, voice, image_in):
    try:
        speech = get_speech(text, voice)
        video = get_dreamtalk(image_in, speech)
        return video
    except Exception as e:
        raise gr.Error(f"Pipeline error: {str(e)}")

def authenticate(password):
    if password == PASSWORD:
        return gr.update(visible=True), gr.update(visible=False, value=""), gr.update(visible=False)
    else:
        return gr.update(visible=False), gr.update(visible=True, value="Invalid password"), gr.update(visible=True)

with gr.Blocks() as demo:
    with gr.Column(visible=True) as login_column:
        gr.HTML("""
        <h2 style="text-align: center;">
        Welcome to Chân Dung Biết Nói
        </h2>
        """)
        with gr.Row():
            logo = gr.Image(label="Logo", type="filepath", value="./logo.jpeg", height=200, width=200)
            introduction = gr.Textbox(label="Introduction", value="Đây là ứng dụng giúp bạn tạo một MC ảo từ giọng nói và ảnh bất kỳ.", interactive=False)
        password_input = gr.Textbox(label="Password", type="password")
        login_btn = gr.Button("Login")
        error_message = gr.Textbox(label="", visible=False)

    with gr.Column(visible=False) as main_column:
        gr.HTML("""
        <h2 style="text-align: center;">
        Chân Dung Biết Nói
        </h2>
        <p style="text-align: center;"></p>
        """)
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(label="Portrait IN", type="filepath", value="./maian.PNG")
            with gr.Column():
                voice = gr.Audio(type="filepath", label="Voice")
                text = gr.Textbox(label="Text")
                submit_btn = gr.Button('Submit')
            with gr.Column():
                video_o = gr.Video(label="Video result")

    login_btn.click(
        fn=authenticate,
        inputs=password_input,
        outputs=[main_column, error_message, login_column]
    )

    submit_btn.click(
        fn=pipe,
        inputs=[text, voice, image_in],
        outputs=video_o,  # Removed list brackets as single output expected
        concurrency_limit=3
    )

demo.queue(max_size=10).launch(show_error=True, show_api=False)