import gradio as gr
import subprocess
import os
import time

def run_ttv(prompt, progress=gr.Progress()):
    try:
        # Kiểm tra file input.jpg
        if not os.path.exists("input.jpg"):
            return "Lỗi: Không tìm thấy file input.jpg", None
            
        progress(0, desc="Đang khởi động mô hình...")
        
        # Chạy file ttv.py với prompt được truyền vào
        env = os.environ.copy()
        env["PROMPT_TEXT"] = prompt  # Truyền prompt qua biến môi trường
        
        process = subprocess.Popen(
            ['python', 'ttv.py'],
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True
        )
        
        # Theo dõi output và cập nhật progress
        output_lines = []
        while True:
            line = process.stdout.readline()
            if not line and process.poll() is not None:
                break
            if line:
                output_lines.append(line.strip())
                if "Starting..." in line:
                    progress(0.2, desc="Đang xử lý...")
                elif "Loading model..." in line:
                    progress(0.4, desc="Đang tải mô hình...")
                elif "Generating..." in line:
                    progress(0.6, desc="Đang tạo video...")
                
        progress(0.8, desc="Đang hoàn thiện...")
        
        # Kiểm tra và trả về video mới nhất
        video_files = [f for f in os.listdir('.') if f.endswith(('.mp4', '.avi', '.mov'))]
        if video_files:
            latest_video = video_files[-1]
            progress(1.0, desc="Hoàn thành!")
            return "\n".join(output_lines), latest_video
            
        return "\n".join(output_lines), None
    except Exception as e:
        return f"Có lỗi xảy ra: {str(e)}", None

# Tạo giao diện
with gr.Blocks() as demo:
    gr.Markdown("# Ứng dụng Text-to-Video")
    
    with gr.Row():
        # Thêm input cho prompt
        text_input = gr.Textbox(
            label="Nhập mô tả cho video",
            placeholder="Ví dụ: A little girl is riding a bicycle at high speed...",
            lines=3
        )
    
    with gr.Row():
        # Thêm preview cho input image
        input_image = gr.Image(
            label="Ảnh input.jpg", 
            value="input.jpg" if os.path.exists("input.jpg") else None,
            interactive=False
        )
        video_output = gr.Video(label="Video đã tạo")
    
    with gr.Row():
        run_button = gr.Button("Tạo Video", variant="primary")
        output = gr.Textbox(label="Trạng thái")
    
    # Thêm các thông tin hướng dẫn
    gr.Markdown("""
    ### Hướng dẫn sử dụng:
    1. Đặt ảnh nguồn với tên `input.jpg` vào thư mục chứa code
    2. Nhập mô tả chi tiết cho video bạn muốn tạo
    3. Nhấn "Tạo Video" và đợi quá trình xử lý hoàn tất
    
    ### Lưu ý:
    - Quá trình tạo video có thể mất vài phút
    - Đảm bảo máy tính có đủ RAM và GPU để xử lý
    """)
    
    run_button.click(
        fn=run_ttv, 
        inputs=text_input,
        outputs=[output, video_output]
    )

# Chạy ứng dụng
if __name__ == "__main__":
    demo.launch()