import os
import torch
import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download
import tempfile
import subprocess
from pathlib import Path

# Получаем токен из переменных окружения
HF_HUB_TOKEN = os.environ.get("HF_HUB_TOKEN")

# Инициализация модели будет происходить лениво
model = None
processor = None

def load_model():
    """Ленивая загрузка модели"""
    global model, processor
    if model is None:
        try:
            # Скачиваем файлы модели
            model_files = [
                "diffusion_pytorch_model.safetensors.index.json",
                "multitalk.safetensors"
            ]
            
            local_dir = "./multitalk_weights"
            os.makedirs(local_dir, exist_ok=True)
            
            for file in model_files:
                if not os.path.exists(os.path.join(local_dir, file)):
                    print(f"Скачиваем {file}...")
                    hf_hub_download(
                        repo_id="MeiGen-AI/MeiGen-MultiTalk",
                        filename=file,
                        local_dir=local_dir,
                        token=HF_HUB_TOKEN
                    )
            
            print("Модель загружена!")
            return True
            
        except Exception as e:
            print(f"Ошибка загрузки модели: {e}")
            return False
    return True

def generate_talking_video(audio_file, reference_image, prompt_text="A person is talking"):
    """
    Генерирует видео с говорящим персонажем
    """
    try:
        if not load_model():
            return None, "Ошибка загрузки модели"
        
        if audio_file is None or reference_image is None:
            return None, "Пожалуйста, загрузите аудио файл и референсное изображение"
        
        # Создаем временные файлы
        with tempfile.TemporaryDirectory() as temp_dir:
            # Сохраняем входные файлы
            audio_path = os.path.join(temp_dir, "input_audio.wav")
            image_path = os.path.join(temp_dir, "reference_image.jpg")
            output_path = os.path.join(temp_dir, "output_video.mp4")
            
            # Копируем файлы
            import shutil
            shutil.copy2(audio_file, audio_path)
            
            # Сохраняем изображение
            if hasattr(reference_image, 'save'):
                reference_image.save(image_path)
            else:
                from PIL import Image
                Image.fromarray(reference_image).save(image_path)
            
            # Здесь должна быть логика inference с моделью
            # Пока что возвращаем заглушку
            return None, f"Модель загружена. Получены файлы:\n- Аудио: {audio_path}\n- Изображение: {image_path}\n- Промпт: {prompt_text}\n\nДля полной реализации нужна интеграция с MultiTalk pipeline."
            
    except Exception as e:
        return None, f"Ошибка генерации: {str(e)}"

def generate_from_text(text, reference_image, voice_type="female"):
    """
    Генерирует видео из текста (сначала TTS, потом lip-sync)
    """
    try:
        if not text.strip():
            return None, "Введите текст для озвучивания"
        
        if reference_image is None:
            return None, "Загрузите референсное изображение"
        
        # Здесь нужно сначала сгенерировать аудио из текста
        # Затем использовать его для генерации видео
        
        return None, f"Функция text-to-video в разработке.\nТекст: {text}\nТип голоса: {voice_type}"
        
    except Exception as e:
        return None, f"Ошибка: {str(e)}"

# Создаем интерфейс Gradio
with gr.Blocks(title="MeiGen-MultiTalk Demo") as demo:
    gr.Markdown("# MeiGen-MultiTalk: Audio-Driven Talking Video Generation")
    gr.Markdown("Создавайте видео с говорящими персонажами из аудио и референсных изображений")
    
    with gr.Tabs():
        with gr.Tab("Аудио → Видео"):
            gr.Markdown("### Загрузите аудио и референсное изображение для создания говорящего видео")
            
            with gr.Row():
                with gr.Column():
                    audio_input = gr.Audio(
                        label="Аудио файл",
                        type="filepath",
                        sources=["upload", "microphone"]
                    )
                    image_input = gr.Image(
                        label="Референсное изображение лица",
                        type="pil"
                    )
                    prompt_input = gr.Textbox(
                        label="Описание видео (опционально)",
                        value="A person is talking naturally",
                        lines=2
                    )
                    generate_btn = gr.Button("Генерировать видео", variant="primary")
                
                with gr.Column():
                    video_output = gr.Video(label="Результат")
                    status_output = gr.Textbox(label="Статус", lines=3)
            
            generate_btn.click(
                fn=generate_talking_video,
                inputs=[audio_input, image_input, prompt_input],
                outputs=[video_output, status_output]
            )
        
        with gr.Tab("Текст → Видео"):
            gr.Markdown("### Введите текст для создания говорящего видео")
            
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(
                        label="Текст для озвучивания",
                        lines=4,
                        placeholder="Введите текст, который должен произнести персонаж..."
                    )
                    image_input2 = gr.Image(
                        label="Референсное изображение лица",
                        type="pil"
                    )
                    voice_select = gr.Dropdown(
                        choices=["female", "male", "child"],
                        value="female",
                        label="Тип голоса"
                    )
                    generate_btn2 = gr.Button("Генерировать видео", variant="primary")
                
                with gr.Column():
                    video_output2 = gr.Video(label="Результат")
                    status_output2 = gr.Textbox(label="Статус", lines=3)
            
            generate_btn2.click(
                fn=generate_from_text,
                inputs=[text_input, image_input2, voice_select],
                outputs=[video_output2, status_output2]
            )
    
    gr.Markdown("""
    ### Инструкция по использованию:
    1. **Аудио → Видео**: Загрузите аудио файл и фото лица для создания lip-sync видео
    2. **Текст → Видео**: Введите текст, выберите тип голоса и загрузите фото
    
    ### Требования:
    - Изображение должно содержать четко видимое лицо
    - Аудио должно содержать речь (не музыку)
    - Рекомендуемая длительность аудио: до 15 секунд
    """)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)