#encoding = utf8
import logging
import os
import queue
import time
from queue import Queue
from threading import Event, Thread

import numpy as np
import torch

from human_handler import AudioHandler
from utils import load_model, mirror_index, get_device

logger = logging.getLogger(__name__)
current_file_path = os.path.dirname(os.path.abspath(__file__))


class AudioInferenceHandler(AudioHandler):
    def __init__(self, context, handler):
        super().__init__(context, handler)

        self._mal_queue = Queue()
        self._audio_queue = Queue()

        self._exit_event = Event()
        self._run_thread = Thread(target=self.__on_run)
        self._exit_event.set()
        self._run_thread.start()
        logger.info("AudioInferenceHandler init")

    def on_handle(self, stream, type_):
        if type_ == 1:
            self._mal_queue.put(stream)
        elif type_ == 0:
            self._audio_queue.put(stream)

    def __on_run(self):
        wav2lip_path = os.path.join(current_file_path, '..', 'checkpoints', 'wav2lip.pth')
        logger.info(f'AudioInferenceHandler init, path:{wav2lip_path}')
        model = load_model(wav2lip_path)
        logger.info("Model loaded")

        face_list_cycle = self._context.face_list_cycle

        length = len(face_list_cycle)
        index = 0
        count = 0
        count_time = 0
        logger.info('start inference')

        device = get_device()
        logger.info(f'use device:{device}')

        while True:
            if self._exit_event.is_set():
                start_time = time.perf_counter()
                batch_size = self._context.batch_size
                try:
                    mel_batch = self._mal_queue.get(block=True, timeout=0.1)
                except queue.Empty:
                    continue

                is_all_silence = True
                audio_frames = []
                for _ in range(batch_size * 2):
                    frame, type_ = self._audio_queue.get()
                    audio_frames.append((frame, type_))
                    if type_ == 0:
                        is_all_silence = False
                if is_all_silence:
                    for i in range(batch_size):
                        self.on_next_handle((None, mirror_index(length, index), audio_frames[i * 2:i * 2 + 2]),
                                            0)
                        index = index + 1
                else:
                    logger.info('infer=======')
                    t = time.perf_counter()
                    img_batch = []
                    for i in range(batch_size):
                        idx = mirror_index(length, index + i)
                        face = face_list_cycle[idx]
                        img_batch.append(face)

                    img_batch = np.asarray(img_batch)
                    mel_batch = np.asarray(mel_batch)
                    img_masked = img_batch.copy()
                    img_masked[:, face.shape[0] // 2:] = 0

                    img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
                    mel_batch = np.reshape(mel_batch,
                                           [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])

                    img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
                    mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)

                    with torch.no_grad():
                        pred = model(mel_batch, img_batch)

                    pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.

                    count_time += (time.perf_counter() - t)
                    count += batch_size

                    if count >= 100:
                        logger.info(f"------actual avg infer fps:{count / count_time:.4f}")
                        count = 0
                        count_time = 0

                    for i, res_frame in enumerate(pred):
                        self.on_next_handle(
                            (res_frame, mirror_index(length, index), audio_frames[i * 2:i * 2 + 2]),
                            0)
                        index = index + 1
                    logger.info(f'total batch time: {time.perf_counter() - start_time}')
            else:
                time.sleep(1)
                break
        logger.info('AudioInferenceHandler inference processor stop')

    def stop(self):
        self._exit_event.clear()
        self._run_thread.join()

    def pause_talk(self):
        self._audio_queue.queue.clear()
        self._mal_queue.queue.clear()
