import os
import sys
import subprocess
import cv2
import gradio as gr
from huggingface_hub import snapshot_download

# 1. Автоматичне завантаження всіх ONNX-ваг для LivePortrait
st.info if "streamlit" in sys.modules else print("Завантаження чекпоінтів FasterLivePortrait...")
checkpoint_dir = "./checkpoints"
if not os.path.exists(checkpoint_dir):
    snapshot_download(repo_id="warmshao/FasterLivePortrait", local_dir=checkpoint_dir)

# Додаємо кастомні ліби у шлях, якщо клонуємо репозиторій локально
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))

# Імпортуємо головний конвеєр FasterLivePortrait
# (Примітка: для роботи цього імпорту у корені мають лежати скрипти з репозиторію warmshao/FasterLivePortrait)
from src.pipelines.live_portrait_pipeline import LivePortraitPipeline 

class LivePortraitAnimator:
    def __init__(self):
        # Налаштування конфігу під ONNX GPU інференс
        self.pipeline = LivePortraitPipeline(
            config_path="src/config/inference.yaml",
            checkpoint_dir="./checkpoints",
            mode="onnx", # Працюємо строго в ONNX-GPU режимі
            device_id=0  # Активуємо твою CUDA карту на PRO акаунті
        )

    def animate(self, source_image_path, driving_video_path, stitching=True, relative_motion=True):
        if not source_image_path or not driving_video_path:
            return None
        
        output_dir = "./animations"
        os.makedirs(output_dir, exist_ok=True)
        output_path = os.path.join(output_dir, "liveportrait_result.mp4")
        
        # Запуск нативного 3D-таргетингу емоцій
        self.pipeline.execute(
            source_image=source_image_path,
            driving_video=driving_video_path,
            output_path=output_path,
            stitching=stitching,            # Безшовне вшивання обличчя назад у кадр
            relative=relative_motion,       # Перенос відносних рухів (тримає геометрію обличчя-джерела)
            flag_crop=True                  # Автоматичне фокусування та кроп обличчя
        )
        return output_path

# Ініціалізація синглтона аніматора
animator = LivePortraitAnimator()

# 4. Побудова UI на Gradio (ідеально підходить для паралельного відображення Video/Image)
def create_ui():
    with gr.Blocks(title="LivePortrait Studio PRO") as demo:
        gr.Markdown("# 🎭 LivePortrait Studio — Емоційний Аніматор Обличчя (PRO GPU)")
        gr.Markdown("Перенесення живої міміки, поглядів та емоцій з будь-якого driving-відео на фото-джерело.")
        
        with gr.Row():
            with gr.Column():
                source_img = gr.Image(label="Завантажте фото обличчя (Source Image)", type="filepath")
                driving_vid = gr.Video(label="Завантажте емоційне відео (Driving Video)", format="mp4")
                
                with gr.Accordion("Просунуті налаштування маскування", open=False):
                    stitching_toggle = gr.Checkbox(label="Безшовне вшивання (Stitching)", value=True)
                    relative_toggle = gr.Checkbox(label="Відносний рух м'язів (Relative Motion)", value=True)
                
                submit_btn = gr.Button("🚀 Запустити 3D Анімацію", variant="primary")
            
            with gr.Column():
                output_vid = gr.Video(label="Результат у кінематографічній якості (MP4)")
        
        submit_btn.click(
            fn=animator.animate,
            inputs=[source_img, driving_vid, stitching_toggle, relative_toggle],
            outputs=output_vid
        )
        
    return demo

if __name__ == "__main__":
    demo = create_ui()
    # Запуск Gradio всередині контейнера Hugging Face Spaces
    demo.launch(server_name="0.0.0.0", server_port=7860)