import os
import time
import sys
import subprocess
import numpy as np
import requests

import argparse
from omegaconf import OmegaConf
import cv2
import torch
import glob
import pickle
import copy
from argparse import Namespace
import shutil
import gdown
import tempfile

from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware

from musetalk.utils.utils import get_file_type, get_video_fps, datagen
from musetalk.utils.preprocessing import get_landmark_and_bbox, read_imgs, coord_placeholder, get_bbox_range
from musetalk.utils.blending import get_image
from musetalk.utils.utils import load_all_model

ProjectDir = os.path.abspath(os.path.dirname(__file__))
CheckpointsDir = os.path.join(ProjectDir, "models")


def print_directory_contents(path):
    for child in os.listdir(path):
        child_path = os.path.join(path, child)
        if os.path.isdir(child_path):
            print(child_path)


def download_model():
    if not os.path.exists(CheckpointsDir):
        os.makedirs(CheckpointsDir)
        print("Checkpoint Not Downloaded, start downloading...")
        tic = time.time()
        # Download MuseTalk model
        from huggingface_hub import snapshot_download
        snapshot_download(
            repo_id="TMElyralab/MuseTalk",
            local_dir=CheckpointsDir,
            max_workers=8,
            local_dir_use_symlinks=True,
        )
        # Download sd-vae-ft-mse model
        os.makedirs(f"{CheckpointsDir}/sd-vae-ft-mse/", exist_ok=True)
        snapshot_download(
            repo_id="stabilityai/sd-vae-ft-mse",
            local_dir=CheckpointsDir + '/sd-vae-ft-mse',
            max_workers=8,
            local_dir_use_symlinks=True,
        )
        # Download dwpose model
        os.makedirs(f"{CheckpointsDir}/dwpose/", exist_ok=True)
        snapshot_download(
            repo_id="yzd-v/DWPose",
            local_dir=CheckpointsDir + '/dwpose',
            max_workers=8,
            local_dir_use_symlinks=True,
        )
        # Download whisper model
        url = "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt"
        response = requests.get(url)
        if response.status_code == 200:
            file_path = f"{CheckpointsDir}/whisper/tiny.pt"
            os.makedirs(f"{CheckpointsDir}/whisper/", exist_ok=True)
            with open(file_path, "wb") as f:
                f.write(response.content)
        else:
            print(f"请求失败，状态码：{response.status_code}")
        # Download face parse model
        import gdown
        url = "https://drive.google.com/uc?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812"
        os.makedirs(f"{CheckpointsDir}/face-parse-bisent/", exist_ok=True)
        file_path = f"{CheckpointsDir}/face-parse-bisent/79999_iter.pth"
        gdown.download(url, file_path, quiet=False)
        # Download resnet model
        url = "https://download.pytorch.org/models/resnet18-5c106cde.pth"
        response = requests.get(url)
        if response.status_code == 200:
            file_path = f"{CheckpointsDir}/face-parse-bisent/resnet18-5c106cde.pth"
            with open(file_path, "wb") as f:
                f.write(response.content)
        else:
            print(f"请求失败，状态码：{response.status_code}")

        toc = time.time()
        print(f"download cost {toc - tic} seconds")
        print_directory_contents(CheckpointsDir)

    else:
        print("Already download the model.")


download_model()  # for huggingface deployment.

# Load model weights
audio_processor, vae, unet, pe = load_all_model()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
timesteps = torch.tensor([0], device=device)

app = FastAPI()

# 添加跨域中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 可根据需要指定允许的源
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


def trim_video_to_audio_length(video_path, audio_path, output_path):
    if get_file_type(video_path) == "video":
        # 获取视频时间长度
        cmd = ['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
               '-of', 'default=noprint_wrappers=1:nokey=1', video_path]
        video_duration = float(subprocess.check_output(cmd).decode().strip())

        # 获取音频时间长度
        cmd = ['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
               '-of', 'default=noprint_wrappers=1:nokey=1', audio_path]
        audio_duration = float(subprocess.check_output(cmd).decode().strip())

        # 如果音频时间长度小于视频时间长度,则裁剪视频
        if audio_duration < video_duration:
            cmd = ['ffmpeg', '-y', '-i', video_path, '-t',
                   str(audio_duration), '-c', 'copy', output_path]
            subprocess.call(cmd)
            print(f"视频已裁切至与音频相同的时间长度: {audio_duration}秒")
        else:
            output_path = video_path

        return output_path
    else:
        return video_path


def get_trimmed_video_path(video_path, result_dir):
    # 提取视频文件的名称部分(不包括扩展名)
    video_name = os.path.splitext(os.path.basename(video_path))[0]

    # 获取视频文件的扩展名
    video_ext = os.path.splitext(video_path)[1]

    # 生成裁剪后的视频文件名
    trimmed_video_name = f"{video_name}_trim{video_ext}"

    # 拼接裁剪后的视频文件的完整路径
    trimmed_video_path = os.path.join(result_dir, trimmed_video_name)

    return trimmed_video_path


@torch.no_grad()
def inference(audio_path, video_path, bbox_shift, batch_size_args, use_saved_coord_args, trim_video, result_dir):
    start_time = time.time()
    args_dict = {
        "result_dir": result_dir,
        "fps": 25,
        "batch_size": batch_size_args,
        "output_vid_name": '',
        "use_saved_coord": "store_true"
    }
    args = Namespace(**args_dict)

    # 如果开始裁切视频
    if trim_video:
        video_path = trim_video_to_audio_length(
            video_path, audio_path, get_trimmed_video_path(video_path, result_dir))
        print("裁切后的视频路径：", video_path)

    input_basename = os.path.basename(video_path).split('.')[0]
    audio_basename = os.path.basename(audio_path).split('.')[0]
    output_basename = f"{input_basename}_{audio_basename}"
    result_img_save_path = os.path.join(
        args.result_dir, output_basename)  # related to video & audio inputs
    input_img_save_path = os.path.join(
        args.result_dir, input_basename)  # related to video & audio inputs
    crop_coord_save_path = os.path.join(
        input_img_save_path, input_basename+".pkl")  # only related to video input
    print("crop_coord_save_path ", crop_coord_save_path)
    os.makedirs(result_img_save_path, exist_ok=True)

    if args.output_vid_name == "":
        output_vid_name = os.path.join(
            args.result_dir, output_basename+".mp4")
    else:
        output_vid_name = os.path.join(
            args.result_dir, args.output_vid_name)
    ############################################## extract frames from source video ##############################################
    if get_file_type(video_path) == "video":
        save_dir_full = os.path.join(args.result_dir, input_basename)
        os.makedirs(save_dir_full, exist_ok=True)
        cmd = f"ffmpeg -v fatal -i {video_path} -start_number 0 {save_dir_full}/%08d.png"
        os.system(cmd)
        input_img_list = sorted(
            glob.glob(os.path.join(save_dir_full, '*.[jpJP][pnPN]*[gG]')))
        fps = get_video_fps(video_path)
    elif get_file_type(video_path) == "image":
        input_img_list = [video_path, ]
        fps = args.fps
    elif os.path.isdir(video_path):  # input img folder
        input_img_list = glob.glob(
            os.path.join(video_path, '*.[jpJP][pnPN]*[gG]'))
        input_img_list = sorted(
            input_img_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
        fps = args.fps
    else:
        raise ValueError(
            f"{video_path} should be a video file, an image file or a directory of images")
    # print(input_img_list)
    ############################################## extract audio feature ##############################################
    whisper_feature = audio_processor.audio2feat(audio_path)
    whisper_chunks = audio_processor.feature2chunks(
        feature_array=whisper_feature, fps=fps)
    ############################################## preprocess input image  ##############################################
    bbox_shift_text = None
    if os.path.exists(crop_coord_save_path) and use_saved_coord_args:
        print("使用提取的landmark")
        with open(crop_coord_save_path, 'rb') as f:
            coord_list = pickle.load(f)
        frame_list = read_imgs(input_img_list)
    else:
        print("重新提取landmark，很耗时")
        coord_list, frame_list, bbox_shift_text = get_landmark_and_bbox(
            input_img_list, bbox_shift)
        os.makedirs(os.path.dirname(crop_coord_save_path), exist_ok=True)
        with open(crop_coord_save_path, 'wb') as f:
            pickle.dump(coord_list, f)
    # bbox_shift_text=get_bbox_range(input_img_list, bbox_shift)
    i = 0
    input_latent_list = []
    for bbox, frame in zip(coord_list, frame_list):
        if bbox == coord_placeholder:
            continue
        x1, y1, x2, y2 = bbox
        crop_frame = frame[y1:y2, x1:x2]
        crop_frame = cv2.resize(
            crop_frame, (256, 256), interpolation=cv2.INTER_LANCZOS4)
        latents = vae.get_latents_for_unet(crop_frame)
        input_latent_list.append(latents)

    # to smooth the first and the last frame
    frame_list_cycle = frame_list + frame_list[::-1]
    coord_list_cycle = coord_list + coord_list[::-1]
    input_latent_list_cycle = input_latent_list + input_latent_list[::-1]
    ############################################## inference batch by batch ##############################################
    print("start inference")
    video_num = len(whisper_chunks)
    batch_size = batch_size_args
    gen = datagen(whisper_chunks, input_latent_list_cycle, batch_size)
    res_frame_list = []
    for i, (whisper_batch, latent_batch) in enumerate(tqdm(gen, total=int(np.ceil(float(video_num)/batch_size)))):

        tensor_list = [torch.FloatTensor(arr) for arr in whisper_batch]
        audio_feature_batch = torch.stack(
            tensor_list).to(unet.device)  # torch, B, 5*N,384
        audio_feature_batch = pe(audio_feature_batch)

        pred_latents = unet.model(latent_batch, timesteps,
                                  encoder_hidden_states=audio_feature_batch).sample
        recon = vae.decode_latents(pred_latents)
        for res_frame in recon:
            res_frame_list.append(res_frame)

    ############################################## pad to full image ##############################################
    print("pad talking image to original video")
    for i, res_frame in enumerate(tqdm(res_frame_list)):
        bbox = coord_list_cycle[i % (len(coord_list_cycle))]
        ori_frame = copy.deepcopy(
            frame_list_cycle[i % (len(frame_list_cycle))])
        x1, y1, x2, y2 = bbox
        try:
            res_frame = cv2.resize(
                res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
        except:
            continue

        combine_frame = get_image(ori_frame, res_frame, bbox)
        cv2.imwrite(
            f"{result_img_save_path}/{str(i).zfill(8)}.png", combine_frame)
    cmd_img2video = f"ffmpeg -y -v warning -r {fps} -f image2 -i {result_img_save_path}/%08d.png -vcodec libx264 -vf format=rgb24,scale=out_color_matrix=bt709,format=yuv420p -crf 18 {result_dir}/temp.mp4"
    print(cmd_img2video)
    os.system(cmd_img2video)

    cmd_combine_audio = f"ffmpeg -y -v warning -i {audio_path} -i {result_dir}/temp.mp4 {output_vid_name}"
    print(cmd_combine_audio)
    os.system(cmd_combine_audio)

    os.remove(f"{result_dir}/temp.mp4")
    shutil.rmtree(result_img_save_path)
    print(f"result is save to {output_vid_name}")
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"代码块执行耗时：{elapsed_time}秒")
    return output_vid_name, bbox_shift_text


@app.post("/inference")
async def inference_endpoint(
    audio_file: UploadFile = File(...),
    video_file: UploadFile = File(...),
    bbox_shift: float = Form(0),
    batch_size_args: int = Form(8),
    use_saved_coord_args: bool = Form(False),
    trim_video: bool = Form(False),
):
    # 保存上传的文件到临时目录
    with tempfile.TemporaryDirectory() as tmp_dir:
        audio_temp_path = os.path.join(
            tmp_dir, f"audio{os.path.splitext(audio_file.filename)[1]}")
        with open(audio_temp_path, 'wb') as f:
            content = await audio_file.read()
            f.write(content)

        video_temp_path = os.path.join(
            tmp_dir, f"video{os.path.splitext(video_file.filename)[1]}")
        with open(video_temp_path, 'wb') as f:
            content = await video_file.read()
            f.write(content)

        # 调用推理函数
        output_video_path, bbox_shift_text = inference(
            audio_path=audio_temp_path,
            video_path=video_temp_path,
            bbox_shift=bbox_shift,
            batch_size_args=batch_size_args,
            use_saved_coord_args=use_saved_coord_args,
            trim_video=trim_video,
            result_dir=tmp_dir
        )

        # 将生成的视频作为响应返回
        return FileResponse(
            output_video_path,
            media_type='video/mp4',
            filename=os.path.basename(output_video_path)
        )
          这个代码提示以下错误：crop_coord_save_path  /tmp/tmpqllbj3yy/video/video.pkl
video in 30.0 FPS, audio idx in 50FPS
重新提取landmark，很耗时
reading images...
100%|████████████████████████████████████████████████████████████████████████████████████████████████| 901/901 [00:01<00:00, 681.34it/s]
get key_landmark and face bounding boxes with the default value
100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 901/901 [00:30<00:00, 29.76it/s]
********************************************bbox_shift parameter adjustment**********************************************************
INFO:     127.0.0.1:32948 - "POST /inference HTTP/1.1" 500 Internal Server Error
ERROR:    Exception in ASGI application
Traceback (most recent call last):
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/uvicorn/protocols/http/h11_impl.py", line 407, in run_asgi
    result = await app(  # type: ignore[func-returns-value]
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/uvicorn/middleware/proxy_headers.py", line 69, in __call__
    return await self.app(scope, receive, send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/fastapi/applications.py", line 1054, in __call__
    await super().__call__(scope, receive, send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/applications.py", line 123, in __call__
    await self.middleware_stack(scope, receive, send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/middleware/errors.py", line 186, in __call__
    raise exc
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/middleware/errors.py", line 164, in __call__
    await self.app(scope, receive, _send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/middleware/cors.py", line 93, in __call__
    await self.simple_response(scope, receive, send, request_headers=headers)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/middleware/cors.py", line 148, in simple_response
    await self.app(scope, receive, send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/middleware/exceptions.py", line 65, in __call__
    await wrap_app_handling_exceptions(self.app, conn)(scope, receive, send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/_exception_handler.py", line 64, in wrapped_app
    raise exc
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
    await app(scope, receive, sender)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/routing.py", line 756, in __call__
    await self.middleware_stack(scope, receive, send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/routing.py", line 776, in app
    await route.handle(scope, receive, send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/routing.py", line 297, in handle
    await self.app(scope, receive, send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/routing.py", line 77, in app
    await wrap_app_handling_exceptions(app, request)(scope, receive, send)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/_exception_handler.py", line 64, in wrapped_app
    raise exc
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
    await app(scope, receive, sender)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/starlette/routing.py", line 72, in app
    response = await func(request)
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/fastapi/routing.py", line 278, in app
    raw_response = await run_endpoint_function(
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/fastapi/routing.py", line 191, in run_endpoint_function
    return await dependant.call(**values)
  File "/root/MuseTalk/api.py", line 333, in inference_endpoint
    output_video_path, bbox_shift_text = inference(
  File "/root/miniconda3/envs/musetalk/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/MuseTalk/api.py", line 235, in inference
    coord_list, frame_list, bbox_shift_text = get_landmark_and_bbox(
  File "/root/MuseTalk/musetalk/utils/preprocessing.py", line 134, in get_landmark_and_bbox
    print(f"Total frame:「{len(frames)}」 Manually adjust range : [ -{int(sum(average_range_minus) / len(average_range_minus))}~{int(sum(average_range_plus) / len(average_range_plus))} ] , the current value: {upperbondrange}")
ZeroDivisionError: division by zero，请检查并且修正代码逻辑