import os

import gradio as gr
from gradio_client import Client

gurl=os.environ.get("GURL")
concurrency_limit =int(os.environ.get("CLIMIT"))
client = Client(gurl)

def gen_prompt(prompt):
    result = client.predict(
            prompt,	
            fn_index=0
    )
    return result

def gen_video(prompt1, prompt2):
    if prompt2 == "":
        prompt2 = prompt1
    print('P1:',prompt1,'P2:',prompt2)
    result = client.predict(
            prompt1,
            prompt2,
            fn_index=1
    )
    return result


def videocrafter_demo():
    with gr.Blocks(analytics_enabled=False) as videocrafter_iface:
        gr.HTML("<div align='center'> <h2> Pixelgen v1</h2></div>")
        gr.Markdown("""
        <b>1. User can enter a short text and then generate a rich prompt by clicking on the Expand Prompt button.<br>2. Two videos will be generated based on the original user input and the rich prompt respectively.<br>3. It will take 2-3 minutes to generate the HD videos.</b> \
        """)
        
        #######t2v#######
        with gr.Tab(label="Text to Video"):
            with gr.Column():
                with gr.Row():
                    with gr.Column():
                        input_text = gr.Text(label='User Input')
                        prompt_btn = gr.Button("Navigate Prompt")
                        output_text = gr.Text(label='Rich Prompt')
                        video_btn = gr.Button("Generate Videos")
                    with gr.Tab(label='Results'):
                        with gr.Row():
                            output_video_1 =  gr.Video(width=512, label='User Input')
                            output_video_2 =  gr.Video(width=512,label='Navigate Prompt')
            prompt_btn.click(
                fn=gen_prompt,
                inputs=[input_text],
                outputs=[output_text,input_text],
                concurrency_limit=1
            )
            video_btn.click(
                fn=gen_video, 
                inputs=[input_text, output_text],
                outputs=[output_video_1, output_video_2],
                concurrency_limit=concurrency_limit
            )
        

    return videocrafter_iface

if __name__ == "__main__":
    videocrafter_iface = videocrafter_demo()
    videocrafter_iface.queue()
    videocrafter_iface.launch()