import os
from PIL import Image
import random
import shutil
import datetime
import torchvision.transforms.functional as f
import torch

from typing import Optional, Tuple
from threading import Lock
from langchain import ConversationChain

from chat_anything.tts_talker.tts_edge import TTSTalker
from chat_anything.sad_talker.sad_talker import SadTalker
from chat_anything.chatbot.chat import load_chain
from chat_anything.chatbot.select import model_selection_chain
from chat_anything.chatbot.voice_select import voice_selection_chain
import gradio as gr


TALKING_HEAD_WIDTH = "350"
sadtalker_checkpoint_path = "MODELS/SadTalker"
config_path = "chat_anything/sad_talker/config"

class ChatWrapper:
    def __init__(self):
        self.lock = Lock()
        self.sad_talker = SadTalker(
            sadtalker_checkpoint_path, config_path, lazy_load=True)

    def __call__(
            self,
            api_key: str,
            inp: str,
            history: Optional[Tuple[str, str]],
            chain: Optional[ConversationChain],
            speak_text: bool, talking_head: bool,
            uid: str,
            talker : None,
            fullbody : str,
    ):
        """Execute the chat functionality."""
        self.lock.acquire()
        if chain is None:
            history.append((inp, "Please register with your API key first!"))
        else:
            try:
                print("\n==== date/time: " + str(datetime.datetime.now()) + " ====")
                print("inp: " + inp)
                print("speak_text: ", speak_text)
                print("talking_head: ", talking_head)
                history = history or []
                # If chain is None, that is because no API key was provided.
                output = "Please paste your OpenAI key from openai.com to use this app. " + \
                    str(datetime.datetime.now())

                output = chain.predict(input=inp).strip()
                output = output.replace("\n", "\n\n")

                text_to_display = output

                # #预定义一个talker
                # talker = MaleEn()
                history.append((inp, text_to_display))

                html_video, temp_file, html_audio, temp_aud_file = None, None, None, None
                if speak_text:
                    if talking_head:
                        html_video, temp_file = self.do_html_video_speak(
                         talker, output, fullbody, uid)
                    else:
                        html_audio, temp_aud_file = self.do_html_audio_speak(
                         talker,  output,uid)
                else:
                    if talking_head:
                        temp_file = os.path.join('tmp', uid, 'videos')
                        html_video = create_html_video(
                            temp_file, TALKING_HEAD_WIDTH)
                    else:
                        pass

            except Exception as e:
                raise e
            finally:
                self.lock.release()
        return history, history, html_video, temp_file, html_audio, temp_aud_file, ""
    

    def do_html_audio_speak(self,talker, words_to_speak, uid):
        audio_path = os.path.join('tmp', uid, 'audios')
        print('uid:', uid, ":", words_to_speak)
        audo_file_path = talker.test(text=words_to_speak, audio_path=audio_path)
        html_audio = '<pre>no audio</pre>'
        try:
            temp_aud_file = gr.File(audo_file_path)
            # print("audio-----------------------------------------------------success")
            temp_aud_file_url = "/file=" + temp_aud_file.value['name']
            html_audio = f'<audio autoplay><source src={temp_aud_file_url} type="audio/mp3"></audio>'
        except IOError as error:
            # Could not write to file, exit gracefully
            print(error)
            return None, None

        return html_audio, audo_file_path

    def do_html_video_speak(self,talker,words_to_speak,fullbody, uid):
        if fullbody:
            # preprocess='somthing'
            preprocess='full'
        else:
            preprocess='crop'
        print("success")
        video_path = os.path.join('tmp', uid, 'videos')
        if not os.path.exists(video_path):
            os.makedirs(video_path)
        video_file_path = os.path.join(video_path, 'tempfile.mp4')
        _, audio_path = self.do_html_audio_speak(
            talker,words_to_speak,uid)
        face_file_path = os.path.join('tmp', uid, 'images', 'test.jpg')
        
        video = self.sad_talker.test(face_file_path, audio_path,preprocess, uid=uid) #video_file_path
        # print("---------------------------------------------------------success")
        # print(f"moving {video} -> {video_file_path}")
        shutil.move(video, video_file_path)

        return video_file_path, video_file_path


    def generate_init_face_video(self,class_concept="clock", llm=None,uid=None,fullbody=None, ref_image=None, seed=None):
        """
        """
        print('generate concept of', class_concept)
        print("=================================================")
        print('fullbody:', fullbody)
        print('uid:', uid)
        print("==================================================")
        chain, memory, personality_text = load_chain(llm, class_concept)
        model_conf, selected_model = model_selection_chain(llm, class_concept, conf_file='resources/models.yaml') # use class concept to choose a generating model, otherwise crack down
        # model_conf, selected_model = model_selection_chain(llm, personality_text, conf_file='resources/models_personality.yaml') # use class concept to choose a generating model, otherwise crack down
        voice_conf, selected_voice = model_selection_chain(llm, personality_text, conf_file='resources/voices_edge.yaml')

        # added for safe face generation
        print('generate concept of', class_concept)
        augment_word_list = ["Female ", "female ", "beautiful ", "small ", "cute "]
        first_sentence = "Hello, how are you doing today?"
        voice_conf, selected_voice = model_selection_chain(llm, personality_text, conf_file='resources/voices_edge.yaml')
        talker = TTSTalker(selected_voice=selected_voice, gender=voice_conf['gender'], language=voice_conf['language'])
        model_conf, selected_model = model_selection_chain(llm, class_concept, conf_file='resources/models.yaml') # use class concept to choose a generating model, otherwise crack down
        retry_cnt = 4
        if ref_image is None:
            face_files = os.listdir(FACE_DIR)
            face_img_path = os.path.join(FACE_DIR, random.choice(face_files))
            ref_image = Image.open(face_img_path)

        print('loading face generating model')
        anything_facemaker = load_face_generator(
            model_dir=model_conf['model_dir'],                                                                                           
            lora_path=model_conf['lora_path'],                                                                                           
            prompt_template=model_conf['prompt_template'],                                                                               
            negative_prompt=model_conf['negative_prompt'],    
        )
        retry_cnt = 0                                                                                                                                  
        has_face = anything_facemaker.has_face(ref_image)
        init_strength = 1.0 if has_face else 0.85                                                                                       
        strength_retry_step = -0.04 if has_face else 0.04
        while retry_cnt < 8:                                                                                                
            try:                                                                                                                                 
                generate_face_image(                                                                                                             
                    anything_facemaker,
                    class_concept,
                    ref_image,
                    uid=uid,                                                                                                  
                    strength=init_strength if (retry_cnt==0 and has_face) else init_strength + retry_cnt * strength_retry_step,                                          
                    controlnet_conditioning_scale=0.5 if retry_cnt == 8 else 0.3,
                    seed=seed,                                                                                                                              
                )                                                                                                                                
                self.do_html_video_speak(talker, first_sentence, fullbody, uid=uid)                                                                   
                video_file_path = os.path.join('tmp', uid, 'videos/tempfile.mp4')                                                                
                htm_video = create_html_video(                                                                                                   
                    video_file_path, TALKING_HEAD_WIDTH)                                                                                                                                                                                                     
                break                                                                                                                            
            except Exception as e:                                                                                                               
                retry_cnt += 1                                                                                                                
                class_concept = random.choice(augment_word_list) + class_concept                                                                                                                                                                            
                print(e)         
        # end of repeat block       

        return chain, memory, htm_video, talker


    def update_talking_head(self, widget, uid, state):
        # print("success----------------")
        if widget:
            state = widget
            temp_file = os.path.join('tmp', uid, 'videos')
            video_html_talking_head = create_html_video(
                temp_file, TALKING_HEAD_WIDTH)
            return state, video_html_talking_head
        else:
            return None, "<pre></pre>"


def reset_memory(history, memory):
    memory.clear()
    history = []
    return history, history, memory
            

def create_html_video(file_name, width):
    return file_name


def create_html_audio(file_name):
    if os.path.exists(file_name):
        tmp_audio_file = gr.File(file_name, visible=False)
        tmp_aud_file_url = "/file=" + tmp_audio_file.value['name']
        html_audio = f'<audio><source src={tmp_aud_file_url} type="audio/mp3"></audio>'
        del tmp_aud_file_url
    else:
       html_audio = f'' 
    
    return html_audio


def update_foo(widget, state):
    if widget:
        state = widget
        return state


# Pertains to question answering functionality
def update_use_embeddings(widget, state):
    if widget:
        state = widget
        return state

# This is the code for image generating.


def load_face_generator(model_dir, lora_path, prompt_template, negative_prompt):
    from chat_anything.face_generator.long_prompt_control_generator import LongPromptControlGenerator
    # # using local
    model_zoo = "MODELS"
    face_control_dir = os.path.join(
        model_zoo, "Face-Landmark-ControlNet", "models_for_diffusers")
    face_detect_path = os.path.join(
        model_zoo, "SadTalker", "shape_predictor_68_face_landmarks.dat")
    # use remote, hugginface auto-download.
    # use your model path, has to be a model derived from stable diffusion v1-5
    anything_facemaker = LongPromptControlGenerator(
        model_dir=model_dir,
        lora_path=lora_path,
        prompt_template=prompt_template,
        negative_prompt=negative_prompt,
        face_control_dir=face_control_dir,
        face_detect_path=face_detect_path,
    )
    anything_facemaker.load_model(safety_checker=None)
    return anything_facemaker



FACE_DIR="resources/images/faces"
def generate_face_image(
        anything_facemaker,
        class_concept, 
        face_img_pil,
        uid=None,
        controlnet_conditioning_scale=1.0,
        strength=0.95,
        seed=42,
    ):
    face_img_pil = f.center_crop(
        f.resize(face_img_pil, 512), 512).convert('RGB')
    prompt = anything_facemaker.prompt_template.format(class_concept)
    # # There are four ways to generate a image by now.
    # pure_generate = anything_facemaker.generate(prompt=prompt, image=face_img_pil, do_inversion=False)
    # inversion = anything_facemaker.generate(prompt=prompt, image=face_img_pil, strength=strength, do_inversion=True)

    print('USING SEED:', seed)
    generator = torch.Generator(device=anything_facemaker.face_control_pipe.device)
    generator.manual_seed(seed)
    if strength is None:
        pure_control = anything_facemaker.face_control_generate(prompt=prompt, face_img_pil=face_img_pil, do_inversion=False,
                                                                 controlnet_conditioning_scale=controlnet_conditioning_scale, generator=generator)
        init_face_pil = pure_control
    else:
        control_inversion = anything_facemaker.face_control_generate(prompt=prompt, face_img_pil=face_img_pil, do_inversion=True, 
                                                                 strength=strength,
                                                                 controlnet_conditioning_scale=controlnet_conditioning_scale, generator=generator)
        init_face_pil = control_inversion
    print('succeeded generating face image')
    face_path = os.path.join('tmp', uid, 'images')
    if not os.path.exists(face_path):
        os.makedirs(face_path)
    # TODO: reproduce the images for return, shouldn't use the filesystem
    face_file_path = os.path.join(face_path, 'test.jpg')
    init_face_pil.save(face_file_path)
    return init_face_pil