import os
import time
import threading
import queue
import logging
from io import BytesIO

from tqdm import tqdm
from scipy.interpolate import interp1d
from pydub import AudioSegment
from pydub.silence import detect_silence
import soundfile as sf
import wave
import cv2
import torch
from torchvision import transforms
import numpy as np


def geneHeadInfo(sampleRate, bits, sampleNum):
    import struct

    rHeadInfo = b"\x52\x49\x46\x46"
    fileLength = struct.pack("i", sampleNum + 36)
    rHeadInfo += fileLength
    rHeadInfo += b"\x57\x41\x56\x45\x66\x6d\x74\x20\x10\x00\x00\x00\x01\x00\x01\x00"
    rHeadInfo += struct.pack("i", sampleRate)
    rHeadInfo += struct.pack("i", int(sampleRate * bits / 8))
    rHeadInfo += b"\x02\x00"
    rHeadInfo += struct.pack("H", bits)
    rHeadInfo += b"\x64\x61\x74\x61"
    rHeadInfo += struct.pack("i", sampleNum)
    return rHeadInfo


class liteAvatar(object):
    def __init__(
        self,
        data_dir=None,  # avatar data dir (mouth pos)
        num_threads=1,
        use_bg_as_idle=False,
        fps=30,
        generate_offline=False,
        use_gpu=False,
        weight_dir="./weights",
    ):
        logging.info("liteAvatar init start...")

        self.data_dir = data_dir
        self.fps = fps
        self.use_bg_as_idle = use_bg_as_idle
        self.use_gpu = use_gpu
        self.device = "cuda" if use_gpu else "cpu"

        s = time.time()
        from audio2mouth import Audio2Mouth

        self.audio2mouth = Audio2Mouth(use_gpu=use_gpu, weight_dir=weight_dir)
        logging.info(f"audio2mouth init over in {time.time() - s}s")

        self.input_queue = queue.Queue()
        self.output_queue = queue.Queue()
        self.load_data_thread: threading.Thread = None

        logging.info("liteAvatar init over")
        self._generate_offline = generate_offline
        if generate_offline:
            self.load_dynamic_model(data_dir)

            self.threads_prep = []
            barrier_prep = threading.Barrier(num_threads, action=None, timeout=None)
            for i in range(num_threads):
                t = threading.Thread(
                    target=self.face_gen_loop,
                    args=(i, barrier_prep, self.input_queue, self.output_queue),
                )
                self.threads_prep.append(t)

            for t in self.threads_prep:
                t.daemon = True
                t.start()

    def stop_algo(self):
        pass

    def load_dynamic_model(self, data_dir):
        logging.info("start to load dynamic data")
        start_time = time.time()
        self.encoder = torch.jit.load(f"{data_dir}/net_encode.pt").to(self.device)
        self.generator = torch.jit.load(f"{data_dir}/net_decode.pt").to(self.device)

        self.load_pos_bg_video_image_frames(data_dir=data_dir, bg_frame_cnt=150)
        self.load_ref_mouth_images(data_dir=data_dir, bg_frame_cnt=150)
        self.ref_data_list = [0 for x in range(150)]
        self.input_queue = queue.Queue()
        self.output_queue = queue.Queue()
        logging.info(f"load dynamic model in {time.time() - start_time:.3f}s")

    def unload_dynamic_model(self):
        pass

    def load_pos_bg_video_image_frames(self, data_dir, bg_frame_cnt=None):
        """load pos np.ndarray and bg video image frames"""
        t = time.time()
        self.neutral_pose = np.load(f"{data_dir}/neutral_pose.npy")
        self.mouth_scale = None

        self.bg_data_list = []
        bg_video = cv2.VideoCapture(f"{data_dir}/bg_video.mp4")
        while True:
            ret, img = bg_video.read()
            self.bg_data_list.append(img)
            if ret is False:
                break
        self.bg_video_frame_count = (
            len(self.bg_data_list)
            if bg_frame_cnt is None
            else min(bg_frame_cnt, len(self.bg_data_list))
        )

        y1, y2, x1, x2 = open(f"{data_dir}/face_box.txt", "r").readlines()[0].split()
        self.y1, self.y2, self.x1, self.x2 = int(y1), int(y2), int(x1), int(x2)

        self.merge_mask = (np.ones((self.y2 - self.y1, self.x2 - self.x1, 3)) * 255).astype(
            np.uint8
        )
        self.merge_mask[10:-10, 10:-10, :] *= 0
        self.merge_mask = cv2.GaussianBlur(self.merge_mask, (21, 21), 15)
        self.merge_mask = self.merge_mask / 255

        self.frame_vid_list = []

        self.image_transforms = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5]),
            ]
        )
        logging.info(f"load data sync in {time.time() - t:.3f}s")

    def load_ref_mouth_images(self, data_dir, bg_frame_cnt=None):
        """load ref mouth images from bg_video.mp4"""
        logging.info(f"loading data from {data_dir}")
        s = time.time()

        self.ref_img_list = []
        for ii in tqdm(range(bg_frame_cnt)):
            img_file_path = os.path.join(f"{data_dir}", "ref_frames", f"ref_{ii:05d}.jpg")
            image = cv2.cvtColor(cv2.imread(img_file_path)[:, :, 0:3], cv2.COLOR_BGR2RGB)
            image = cv2.resize(image, (384, 384), interpolation=cv2.INTER_LINEAR)
            ref_img = self.image_transforms(np.uint8(image))
            encoder_input = ref_img.unsqueeze(0).float().to(self.device)
            x = self.encoder(encoder_input)
            self.ref_img_list.append(x)

        logging.info(f"load data over in {time.time() - s}s")

    def face_gen_loop(self, thread_id, barrier, in_queue, out_queue):
        while True:
            try:
                data = in_queue.get()
            except queue.Empty:
                break

            if data is None:
                in_queue.put(None)
                break

            s = time.time()

            param_res = data[0]
            bg_frame_id = data[1]
            global_frame_id = data[2]

            mouth_img = self.param2img(param_res, bg_frame_id)
            full_img, mouth_img = self.merge_mouth_to_bg(mouth_img, bg_frame_id)

            logging.info(
                "global_frame_id: {} in {}s".format(global_frame_id, round(time.time() - s, 3))
            )

            out_queue.put((global_frame_id, full_img, mouth_img))

        barrier.wait()
        if thread_id == 0:
            out_queue.put(None)

    def param2img(
        self, param_res: dict, bg_frame_id: int, global_frame_id=0, is_idle=False
    ) -> torch.Tensor:
        param_val = []
        for key in param_res:
            val = param_res[key]
            param_val.append(val)
        param_val = np.asarray(param_val)

        source_img = self.generator(
            self.ref_img_list[bg_frame_id],
            torch.from_numpy(param_val).unsqueeze(0).float().to(self.device),
        )
        source_img = source_img.detach().to("cpu")

        return source_img

    def get_idle_param(self):
        bg_param = self.neutral_pose
        tmp_json = {}
        for ii in range(len(self.audio2mouth.p_list)):
            tmp_json[str(ii)] = float(bg_param[ii])
        return tmp_json

    def merge_mouth_to_bg(
        self, mouth_image: torch.Tensor, bg_frame_id: int, use_bg: bool = False
    ) -> tuple[np.ndarray, np.ndarray]:
        mouth_image = (mouth_image / 2 + 0.5).clamp(0, 1)
        mouth_image = mouth_image[0].permute(1, 2, 0) * 255

        mouth_image = mouth_image.numpy().astype(np.uint8)
        mouth_image = cv2.resize(mouth_image, (self.x2 - self.x1, self.y2 - self.y1))
        mouth_image = mouth_image[:, :, ::-1]
        full_img = self.bg_data_list[bg_frame_id].copy()
        if not use_bg:
            # image no bg, use default bg
            full_img[self.y1 : self.y2, self.x1 : self.x2, :] = (
                mouth_image * (1 - self.merge_mask)
                + full_img[self.y1 : self.y2, self.x1 : self.x2, :] * self.merge_mask
            )
        full_img = full_img.astype(np.uint8)
        return full_img, mouth_image.astype(np.uint8)

    def interp_param(self, param_res, fps=25):
        old_len = len(param_res)
        new_len = int(old_len / 30 * fps + 0.5)

        interp_list = {}
        for key in param_res[0]:
            tmp_list = []
            for ii in range(len(param_res)):
                tmp_list.append(param_res[ii][key])
            tmp_list = np.asarray(tmp_list)

            x = np.linspace(0, old_len - 1, num=old_len, endpoint=True)
            newx = np.linspace(0, old_len - 1, num=new_len, endpoint=True)
            f = interp1d(x, tmp_list)
            y = f(newx)
            interp_list[key] = y

        new_param_res = []
        for ii in range(new_len):
            tmp_json = {}
            for key in interp_list:
                tmp_json[key] = interp_list[key][ii]
            new_param_res.append(tmp_json)

        return new_param_res

    def padding_last(self, param_res, last_end=None):
        bg_param = self.neutral_pose

        if last_end is None:
            last_end = len(param_res)

        padding_cnt = 5
        final_end = max(last_end + 5, len(param_res))
        param_res = param_res[:last_end]
        padding_list = []
        for ii in range(last_end, final_end):
            tmp_json = {}
            for key in param_res[-1]:
                kk = ii - last_end
                scale = max((padding_cnt - kk - 1) / padding_cnt, 0.0)

                end_value = bg_param[int(key)]
                tmp_json[key] = (param_res[-1][key] - end_value) * scale + end_value
            padding_list.append(tmp_json)

        logging.info(f"padding_cnt:{len(padding_list)}")
        param_res = param_res + padding_list
        return param_res

    def audio2param(
        self, input_audio_byte, prefix_padding_size=0, is_complete=False, audio_status=-1
    ):
        headinfo = geneHeadInfo(16000, 16, len(input_audio_byte))
        input_audio_byte = headinfo + input_audio_byte
        input_audio, sr = sf.read(BytesIO(input_audio_byte))

        param_res, _, _ = self.audio2mouth.inference(subtitles=None, input_audio=input_audio)

        sil_scale = np.zeros(len(param_res))
        sound = AudioSegment.from_raw(
            BytesIO(input_audio_byte), sample_width=2, frame_rate=16000, channels=1
        )
        start_end_list = detect_silence(sound, 500, -50)
        if len(start_end_list) > 0:
            for start, end in start_end_list:
                start_frame = int(start / 1000 * 30)
                end_frame = int(end / 1000 * 30)
                logging.info(f"silence part: {start_frame}-{end_frame} frames")
                sil_scale[start_frame:end_frame] = 1
        sil_scale = np.pad(sil_scale, 2, mode="edge")
        kernel = np.array([0.1, 0.2, 0.4, 0.2, 0.1])
        sil_scale = np.convolve(sil_scale, kernel, "same")
        sil_scale = sil_scale[2:-2]
        self.make_silence(param_res, sil_scale)
        if self.fps != 30:
            param_res = self.interp_param(param_res, fps=self.fps)

        if is_complete:
            param_res = self.padding_last(param_res)

        return param_res

    def make_silence(self, param_res, sil_scale):
        bg_param = self.neutral_pose

        for ii in range(len(param_res)):
            for key in param_res[ii]:
                neu_value = bg_param[int(key)]
                param_res[ii][key] = (
                    param_res[ii][key] * (1 - sil_scale[ii]) + neu_value * sil_scale[ii]
                )
        return param_res

    def handle(self, audio_file_path, result_dir, param_res=None):
        audio_data = self.read_wav_to_bytes(audio_file_path)

        if param_res is None:
            param_res = self.audio2param(audio_data)

        for ii in range(len(param_res)):
            s = time.time()
            frame_id = ii
            if int(frame_id / self.bg_video_frame_count) % 2 == 0:
                frame_id = frame_id % self.bg_video_frame_count
            else:
                frame_id = self.bg_video_frame_count - 1 - frame_id % self.bg_video_frame_count
            self.input_queue.put((param_res[ii], frame_id, ii))

        self.input_queue.put(None)

        tmp_frame_dir = os.path.join(result_dir, "tmp_frames")
        if os.path.exists(tmp_frame_dir):
            os.system(f"rm -rf {tmp_frame_dir}")
        os.mkdir(tmp_frame_dir)

        while True:
            res_data = self.output_queue.get()
            if res_data is None:
                break
            global_frame_index = res_data[0]
            target_path = f"{tmp_frame_dir}/{str(global_frame_index + 1).zfill(5)}.jpg"
            cv2.imwrite(target_path, res_data[1])

        for p in self.threads_prep:
            p.join()

        cmd = "`which ffmpeg` -r 30 -i {}/%05d.jpg -i {} -framerate 30 -c:v libx264 -pix_fmt yuv420p -b:v 5000k -strict experimental -loglevel error {}/test_demo.mp4 -y".format(
            tmp_frame_dir, audio_file_path, result_dir
        )
        os.system(cmd)

    @staticmethod
    def read_wav_to_bytes(file_path):
        try:
            # 打开WAV文件
            with wave.open(file_path, "rb") as wav_file:
                # 获取WAV文件的参数
                params = wav_file.getparams()
                logging.info(
                    f"Channels: {params.nchannels}, Sample Width: {params.sampwidth}, Frame Rate: {params.framerate}, Number of Frames: {params.nframes}"
                )

                # 读取所有帧
                frames = wav_file.readframes(params.nframes)
                return frames
        except wave.Error as e:
            logging.warning(f"Error reading WAV file: {e}")
            return None


"""
huggingface-cli download weege007/liteavatar --local-dir ./models/weege007/liteavatar
modelscope download --model HumanAIGC-Engineering/LiteAvatarGallery 20250408/sample_data.zip --local_dir ./resources/avatar/liteavatar
unzip ./resources/avatar/liteavatar/20250408/sample_data.zip -d ./resources/avatar/liteavatar/20250408
PYTHONPATH=$PYTHONPATH:./deps/LiteAvatar python -m deps.LiteAvatar.lite_avatar \
    --weight_dir ./models/weege007/liteavatar \
    --data_dir ./resources/avatar/liteavatar/20250408/sample_data \
    --audio_file ./test/audio_files/asr_example_zh.wav \
    --result_dir ./

modelscope download --model HumanAIGC-Engineering/LiteAvatarGallery 20250612/P1-64AzfrJY037WpS69RiUMw.zip --local_dir ./resources/avatar/liteavatar
unzip ./resources/avatar/liteavatar/20250612/P1-64AzfrJY037WpS69RiUMw.zip -d ./resources/avatar/liteavatar/20250612
PYTHONPATH=$PYTHONPATH:./deps/LiteAvatar python -m deps.LiteAvatar.lite_avatar \
    --weight_dir ./models/weege007/liteavatar \
    --data_dir ./resources/avatar/liteavatar/20250612/P1-64AzfrJY037WpS69RiUMw \
    --audio_file ./test/audio_files/asr_example_zh.wav \
    --result_dir ./
"""
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", type=str)
    parser.add_argument("--audio_file", type=str)
    parser.add_argument("--result_dir", type=str)
    parser.add_argument("--weight_dir", type=str)
    parser.add_argument("--log_level", type=str, default="INFO")
    args = parser.parse_args()

    log_format = "%(asctime)s - %(name)s - %(levelname)s - %(pathname)s:%(lineno)d - %(funcName)s - %(message)s"
    logging.basicConfig(
        level=args.log_level,
        format=log_format,
    )

    audio_file = args.audio_file
    tmp_frame_dir = args.result_dir
    weight_dir = args.weight_dir
    use_gpu = True if torch.cuda.is_available() else False

    lite_avatar = liteAvatar(
        data_dir=args.data_dir,
        num_threads=1,
        generate_offline=True,
        weight_dir=weight_dir,
        use_gpu=use_gpu,
    )

    lite_avatar.handle(audio_file, tmp_frame_dir)
