from time import strftime
from src.utils.preprocess import CropAndExtract
from src.test_audio2coeff import Audio2Coeff  
from src.facerender.animate_onnx import AnimateFromCoeff
from src.generate_batch import get_data
from src.generate_facerender_batch import get_facerender_data
from src.utils.init_path import init_path


import os
import asyncio
import random
import logging  

pic_path ="./sadtalker_default.jpeg" #图片文件路径
audio_path = "/root/audio/001.mp3" #音频文件路径
image_directory_path = "/root/baby/" #图片目录
audio_directory_path = "/root/audio/" #音频文件目录
results_folder = "/root/SadTalker/results" #存储结果的文件夹
facerender_batch_size = 4
sadtalker_paths = init_path("./checkpoints", os.path.join("/root/SadTalker", 'src/config'), "256", False, "full")

# 配置日志  
logging.basicConfig(filename=f'./log/{strftime("%Y_%m_%d_%H.%M.%S")}.log', level=logging.DEBUG) 

preprocess_model = CropAndExtract(sadtalker_paths, "cuda")
audio_to_coeff = Audio2Coeff(sadtalker_paths, "cuda")
animate_from_coeff = AnimateFromCoeff(sadtalker_paths, "cuda")

def get_image_files(directory):
    '''获取指定文件夹内的所有图片文件'''
    image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp']  # 添加你认为是图像文件的扩展名

    image_files = []

    for filename in os.listdir(directory):
        if any(filename.lower().endswith(ext) for ext in image_extensions):
            image_files.append(filename)

    return image_files

def get_audio_files(directory):
    '''获取指定文件夹内的所有音频文件'''
    audio_extensions = ['.mp3', '.wav', '.ogg', '.flac', '.aac', '.m4a']  # 添加你认为是音频文件的扩展名

    audio_files = []

    for filename in os.listdir(directory):
        if any(filename.lower().endswith(ext) for ext in audio_extensions):
            audio_files.append(filename)

    return audio_files

async def predict_image(pic_path,audio_path):
    save_dir = os.path.join(results_folder, strftime("%Y_%m_%d_%H.%M.%S"))
    """
    根据给定的图像和音频文件生成视频
    """

    first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
    os.makedirs(first_frame_dir, exist_ok=True)
    first_coeff_path, crop_pic_path, crop_info =  preprocess_model.generate(pic_path, first_frame_dir, "full", source_image_flag=True)
    ref_eyeblink_coeff_path=None
    ref_pose_coeff_path=None
    batch = get_data(first_coeff_path, audio_path, "cuda", ref_eyeblink_coeff_path, still=True)
    coeff_path = audio_to_coeff.generate(batch, save_dir, 0, ref_pose_coeff_path)

    data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, 
                                facerender_batch_size, None, None, None,
                                expression_scale=1, still_mode=True, preprocess="full")
    video_path = animate_from_coeff.generate_deploy(data, save_dir, pic_path, crop_info, \
                                enhancer="gfpgan", background_enhancer=None, preprocess="full")
    logging.info(f'将生成的文件{video_path}拷贝到上级目录: {results_folder}')
    os.system(f'mv {video_path} {results_folder}')
    return video_path

if __name__ == '__main__':
    # 获取图片
    image_files = get_image_files(image_directory_path)

    print("Image Files:")
    for file in image_files:
        print(file)
    
    # 获取音频
    audio_files = get_audio_files(audio_directory_path)

    print("Audio Files:")
    for file in audio_files:
        print(file)

    # 将所有音频转化为视频
    for audio in audio_files:
        # 随机挑选一个图片
        random_image = random.choice(image_files)
        # 图片完整路径
        image_path = os.path.join(image_directory_path, random_image)
        # 音频文件完整路径
        audio_path = os.path.join(audio_directory_path,audio)
        logging.info(f'开始转换：图片：{random_image}，音频：{audio}')
        asyncio.run(predict_image(image_path,audio_path))
        logging.info(f'结束转换')
    # 执行关机命令
    logging.info('执行关机命令')
    os.system('shutdown')