# coding: utf-8

"""
The entrance of the gradio
"""

import tyro
import gradio as gr
import os.path as osp
from src.utils.helper import load_description
from src.gradio_pipeline_stf import GradioPipeline
from src.config.crop_config import CropConfig
from src.config.argument_config import ArgumentConfig
from src.config.inference_config import InferenceConfig
import spaces
import cv2

# import gdown
# folder_url = f"https://drive.google.com/drive/folders/1UtKgzKjFAOmZkhNK-OYT0caJ_w2XAnib"
# gdown.download_folder(url=folder_url, output="pretrained_weights", quiet=False)

import os
import sys
from pathlib import Path
from src.utils.video import extract_audio
#from elevenlabs_utils import ElevenLabsPipeline
from setup_environment import initialize_environment

initialize_environment()



sys.path.append('/home/user/.local/lib/python3.10/site-packages')
sys.path.append('/home/user/.local/lib/python3.10/site-packages/stf_alternative/src/stf_alternative')
sys.path.append('/home/user/.local/lib/python3.10/site-packages/stf_tools/src/stf_tools')
sys.path.append('/home/user/app/')
sys.path.append('/home/user/app/stf/')
sys.path.append('/home/user/app/stf/stf_alternative/')
sys.path.append('/home/user/app/stf/stf_alternative/src/stf_alternative')
sys.path.append('/home/user/app/stf/stf_tools')
sys.path.append('/home/user/app/stf/stf_tools/src/stf_tools')



# CUDA 경로를 환경 변수로 설정
os.environ['PATH'] = '/usr/local/cuda/bin:' + os.environ.get('PATH', '')
os.environ['LD_LIBRARY_PATH'] = '/usr/local/cuda/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')
# 확인용 출력
print("PATH:", os.environ['PATH'])
print("LD_LIBRARY_PATH:", os.environ['LD_LIBRARY_PATH'])

from stf_utils import STFPipeline

stf_pipeline = STFPipeline()
audio_path="assets/examples/driving/test_aud.mp3"
#audio_path_component = gr.Textbox(label="Input", value="assets/examples/driving/test_aud.mp3")


@spaces.GPU(duration=120)
def gpu_wrapped_stf_pipeline_execute(audio_path):
    return stf_pipeline.execute(audio_path)

    
###### 테스트중 ######
    

stf_pipeline = STFPipeline()
driving_video_path=gr.Video()

# set tyro theme
tyro.extras.set_accent_color("bright_cyan")
args = tyro.cli(ArgumentConfig)

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    with gr.Row():
        audio_path_component = gr.Textbox(label="Input", value="assets/examples/driving/test_aud.mp3")
        stf_button = gr.Button("stf test", variant="primary")
        stf_button.click(
                    fn=gpu_wrapped_stf_pipeline_execute,
                    inputs=[
                        audio_path_component
                    ],
                    outputs=[driving_video_path]
                )
    with gr.Row():
        driving_video_path.render()

###### 테스트중 ######



# def partial_fields(target_class, kwargs):
#     return target_class(**{k: v for k, v in kwargs.items() if hasattr(target_class, k)})

# # set tyro theme
# tyro.extras.set_accent_color("bright_cyan")
# args = tyro.cli(ArgumentConfig)

# # specify configs for inference
# inference_cfg = partial_fields(InferenceConfig, args.__dict__)  # use attribute of args to initial InferenceConfig
# crop_cfg = partial_fields(CropConfig, args.__dict__)  # use attribute of args to initial CropConfig

# # gradio_pipeline = GradioPipeline(
# #     inference_cfg=inference_cfg,
# #     crop_cfg=crop_cfg,
# #     args=args
# # )

# @spaces.GPU(duration=120)
# def gpu_wrapped_execute_video(*args, **kwargs):
#     return gradio_pipeline.execute_video(*args, **kwargs)

# @spaces.GPU(duration=120)
# def gpu_wrapped_execute_image(*args, **kwargs):
#     return gradio_pipeline.execute_image(*args, **kwargs)

# @spaces.GPU(duration=120)
# def gpu_wrapped_stf_pipeline_execute(audio_path):
#     return stf_pipeline.execute(audio_path)

# @spaces.GPU(duration=120)
# def gpu_wrapped_elevenlabs_pipeline_generate_voice(text, voice):
#     return elevenlabs_pipeline.generate_voice(text, voice)

    
# def is_square_video(video_path):
#     video = cv2.VideoCapture(video_path)

#     width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
#     height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))

#     video.release()
#     if width != height:
#         raise gr.Error("Error: the video does not have a square aspect ratio. We currently only support square videos")

#     return gr.update(visible=True)
    

# # assets
# title_md = "assets/gradio_title.md"
# example_portrait_dir = "assets/examples/source"
# example_video_dir = "assets/examples/driving"
# data_examples = [
#     [osp.join(example_portrait_dir, "s9.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],
#     [osp.join(example_portrait_dir, "s6.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],
#     [osp.join(example_portrait_dir, "s10.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],
#     [osp.join(example_portrait_dir, "s5.jpg"), osp.join(example_video_dir, "d18.mp4"), True, True, True, True],
#     [osp.join(example_portrait_dir, "s7.jpg"), osp.join(example_video_dir, "d19.mp4"), True, True, True, True],
#     [osp.join(example_portrait_dir, "s22.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],
# ]
# #################### interface logic ####################



# # Define components first
# eye_retargeting_slider = gr.Slider(minimum=0, maximum=0.8, step=0.01, label="target eyes-open ratio")
# lip_retargeting_slider = gr.Slider(minimum=0, maximum=0.8, step=0.01, label="target lip-open ratio")
# retargeting_input_image = gr.Image(type="filepath")
# output_image = gr.Image(type="numpy")
# output_image_paste_back = gr.Image(type="numpy")
# output_video = gr.Video()
# output_video_concat = gr.Video()

        


# def run_end_to_end(image_path, text, voice, input_video, flag_relative, flag_do_crop, flag_remap, flag_crop_driving_video, male): #, animal):
    
#     # # animal 체크 여부에 따라 다른 pipeline 사용
#     # if animal:
#     #     gradio_pipeline = GradioPipelineAnimal(
#     #         inference_cfg=inference_cfg,
#     #         crop_cfg=crop_cfg,
#     #         args=args
#     #     )
#     # else:
#     #     gradio_pipeline = GradioPipeline(
#     #         inference_cfg=inference_cfg,
#     #         crop_cfg=crop_cfg,
#     #         args=args
#     #     )


#     if input_video is None:
    
#         if not male:
#             stf_pipeline = STFPipeline()
#         else:
#             stf_pipeline = STFPipeline(template_video_path="/home/user/app/stf/TEMP/Cam2_2309071202_0012_Natural_Looped.mp4",
#                config_path="/home/user/app/stf/TEMP/front_config_v3.json",
#                checkpoint_path="/home/user/app/stf/TEMP/0157.pth",
#                )

#     if input_video is None:
#         #audio_path = elevenlabs_pipeline.generate_voice(text, voice)
#         audio_path = gpu_wrapped_elevenlabs_pipeline_generate_voice(text, voice)
#         #driving_video_path = stf_pipeline.execute(audio_path)
#         driving_video_path = gpu_wrapped_stf_pipeline_execute(audio_path)
#     else:
#         driving_video_path = input_video
#         os.makedirs("animations",exist_ok=True)
#         audio_path = osp.join("animations", Path(input_video).stem+".wav")
#         extract_audio(driving_video_path, audio_path)


#     #output_path, crop_output_path = gradio_pipeline.execute_video(
#     output_path, crop_output_path = gpu_wrapped_execute_video(
#             input_image_path=image_path,
#             input_video_path=driving_video_path,
#             # input_driving_video_pickle_path=None,
#             flag_do_crop_input=flag_do_crop,
#             flag_remap_input=flag_remap,
#             flag_relative_input=flag_relative,
#             # driving_multiplier=1.0,
#             # flag_stitching=False,
#             # flag_crop_driving_video_input=flag_crop_driving_video,
#             # scale=2.3,
#             # vx_ratio=0.0,
#             # vy_ratio=-0.125,
#             # scale_crop_driving_video=2.2,
#             # vx_ratio_crop_driving_video=0.0,
#             # vy_ratio_crop_driving_video=-0.1,
#             # tab_selection=None,
#             audio_path=audio_path
#             )
   
#     return output_path, crop_output_path




# ###### 테스트중 ######
    

# stf_pipeline = STFPipeline()

# with gr.Blocks(theme=gr.themes.Soft()) as demo:
#     stf_button = gr.Button("stf test", variant="primary")
#     stf_button.click(
#                 fn=gpu_wrapped_stf_pipeline_execute,
#                 inputs=[
#                     audio_path
#                 ],
#                 outputs=[driving_video_path]
#             )

# ###### 테스트중 ######


# with gr.Blocks(theme=gr.themes.Soft()) as demo:
#     with gr.Tabs():
#         # 첫 번째 탭: Text to LipSync
#         with gr.Tab("Text to LipSync"):
#             gr.Markdown("# Text to LipSync")
#             with gr.Row():
#                 script_txt = gr.Text()
#                 voice = gr.Audio(label="사용자 음성", type="filepath")
#                 input_video = gr.Video()

#             with gr.Row():
#                 image_input = gr.Image(type="filepath")  # 여기서 image_input을 정의합니다.
#                 output_video.render()
#                 #crop_output_video.render()
#                 output_video_concat.render()

#             with gr.Row():
#                 flag_relative_input = gr.Checkbox(value=True, label="relative motion")
#                 flag_do_crop_input = gr.Checkbox(value=True, label="do crop")
#                 flag_remap_input = gr.Checkbox(value=True, label="paste-back")
#                 flag_crop_driving_video_input = gr.Checkbox(value=False, label="do crop (driving video)")
#                 male = gr.Checkbox(value=False, label="male")
#                 #animal = gr.Checkbox(value=False, label="animal")  # animal 체크박스 추가

#             with gr.Row():
#                 generate_speech = gr.Button("🚀 Generate Speech", variant="primary")

#             generate_speech.click(
#                 fn=run_end_to_end,
#                 inputs=[
#                     image_input,
#                     script_txt,
#                     voice,
#                     input_video,
#                     flag_relative_input,
#                     flag_do_crop_input,
#                     flag_remap_input,
#                     flag_crop_driving_video_input,
#                     male,
#                     #animal  # 추가된 animal 입력
#                 ],
#                 outputs=[output_video, output_video_concat]
#             )

#         # # 두 번째 탭: FLUX 이미지 생성
#         # with gr.Tab("FLUX 이미지 생성"):
#         #     flux_tab(image_input)  # FLUX 이미지 생성을 위한 별도의 탭

#         # # 세 번째 탭: Flux 개발용 탭
#         # with gr.Tab("FLUX Dev"):
#         #     flux_demo = create_flux_tab()  # Flux 개발용 탭 생성
#         #     #flux_demo.render()  # 해당 UI를 별도의 탭에만 렌더링







    
# with gr.Blocks(theme=gr.themes.Soft()) as demo:
#     gr.HTML(load_description(title_md))
#     gr.Markdown(load_description("assets/gradio_description_upload.md"))
#     with gr.Row():
#         with gr.Accordion(open=True, label="Source Portrait"):
#             image_input = gr.Image(type="filepath")
#             gr.Examples(
#                 examples=[
#                     [osp.join(example_portrait_dir, "s9.jpg")],
#                     [osp.join(example_portrait_dir, "s6.jpg")],
#                     [osp.join(example_portrait_dir, "s10.jpg")],
#                     [osp.join(example_portrait_dir, "s5.jpg")],
#                     [osp.join(example_portrait_dir, "s7.jpg")],
#                     [osp.join(example_portrait_dir, "s12.jpg")],
#                     [osp.join(example_portrait_dir, "s22.jpg")],
#                 ],
#                 inputs=[image_input],
#                 cache_examples=False,
#             )
#         with gr.Accordion(open=True, label="Driving Video"):
#             video_input = gr.Video()
#             gr.Examples(
#                 examples=[
#                     [osp.join(example_video_dir, "d0.mp4")],
#                     [osp.join(example_video_dir, "d18.mp4")],
#                     [osp.join(example_video_dir, "d19.mp4")],
#                     [osp.join(example_video_dir, "d14_trim.mp4")],
#                     [osp.join(example_video_dir, "d6_trim.mp4")],
#                 ],
#                 inputs=[video_input],
#                 cache_examples=False,
#             )
#     with gr.Row():
#         with gr.Accordion(open=False, label="Animation Instructions and Options"):
#             gr.Markdown(load_description("assets/gradio_description_animation.md"))
#             with gr.Row():
#                 flag_relative_input = gr.Checkbox(value=True, label="relative motion")
#                 flag_do_crop_input = gr.Checkbox(value=True, label="do crop")
#                 flag_remap_input = gr.Checkbox(value=True, label="paste-back")
#     gr.Markdown(load_description("assets/gradio_description_animate_clear.md"))
#     with gr.Row():
#         with gr.Column():
#             process_button_animation = gr.Button("🚀 Animate", variant="primary")
#         with gr.Column():
#             process_button_reset = gr.ClearButton([image_input, video_input, output_video, output_video_concat], value="🧹 Clear")
#     with gr.Row():
#         with gr.Column():
#             with gr.Accordion(open=True, label="The animated video in the original image space"):
#                 output_video.render()
#         with gr.Column():
#             with gr.Accordion(open=True, label="The animated video"):
#                 output_video_concat.render()
#     with gr.Row():
#         # Examples
#         gr.Markdown("## You could also choose the examples below by one click ⬇️")
#     with gr.Row():
#         gr.Examples(
#             examples=data_examples,
#             fn=gpu_wrapped_execute_video,
#             inputs=[
#                 image_input,
#                 video_input,
#                 flag_relative_input,
#                 flag_do_crop_input,
#                 flag_remap_input
#             ],
#             outputs=[output_image, output_image_paste_back],
#             examples_per_page=6,
#             cache_examples=False,
#         )
#     gr.Markdown(load_description("assets/gradio_description_retargeting.md"), visible=True)
#     with gr.Row(visible=True):
#         eye_retargeting_slider.render()
#         lip_retargeting_slider.render()
#     with gr.Row(visible=True):
#         process_button_retargeting = gr.Button("🚗 Retargeting", variant="primary")
#         process_button_reset_retargeting = gr.ClearButton(
#             [
#                 eye_retargeting_slider,
#                 lip_retargeting_slider,
#                 retargeting_input_image,
#                 output_image,
#                 output_image_paste_back
#             ],
#             value="🧹 Clear"
#         )
#     with gr.Row(visible=True):
#         with gr.Column():
#             with gr.Accordion(open=True, label="Retargeting Input"):
#                 retargeting_input_image.render()
#                 gr.Examples(
#                     examples=[
#                         [osp.join(example_portrait_dir, "s9.jpg")],
#                         [osp.join(example_portrait_dir, "s6.jpg")],
#                         [osp.join(example_portrait_dir, "s10.jpg")],
#                         [osp.join(example_portrait_dir, "s5.jpg")],
#                         [osp.join(example_portrait_dir, "s7.jpg")],
#                         [osp.join(example_portrait_dir, "s12.jpg")],
#                         [osp.join(example_portrait_dir, "s22.jpg")],
#                     ],
#                     inputs=[retargeting_input_image],
#                     cache_examples=False,
#                 )
#         with gr.Column():
#             with gr.Accordion(open=True, label="Retargeting Result"):
#                 output_image.render()
#         with gr.Column():
#             with gr.Accordion(open=True, label="Paste-back Result"):
#                 output_image_paste_back.render()
#     # binding functions for buttons
#     process_button_retargeting.click(
#         # fn=gradio_pipeline.execute_image,
#         fn=gpu_wrapped_execute_image,
#         inputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image, flag_do_crop_input],
#         outputs=[output_image, output_image_paste_back],
#         show_progress=True
#     )
#     process_button_animation.click(
#         fn=gpu_wrapped_execute_video,
#         inputs=[
#             image_input,
#             video_input,
#             flag_relative_input,
#             flag_do_crop_input,
#             flag_remap_input
#         ],
#         outputs=[output_video, output_video_concat],
#         show_progress=True
#     )
#     # image_input.change(
#     #     fn=gradio_pipeline.prepare_retargeting,
#     #     inputs=image_input,
#     #     outputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image]
#     # )
#     video_input.upload(
#         fn=is_square_video,
#         inputs=video_input,
#         outputs=video_input
#     )

demo.launch(
    server_port=args.server_port,
    share=args.share,
    server_name=args.server_name
)