import os
import openai
import re
import time


class OpenAiApi:

    def __init__(self):
        openai.api_base = "https://api.chatanywhere.tech/v1"
        # openai.api_base = "https://api.chatanywhere.com.cn/v1"
        # openai.api_base = "https://api.chatanywhere.tech"
        # openai.organization = "org-hPyWJWxTelSVBObp7A7GWb2H"
        openai.api_key = "sk-yPoMGUTiAXfV3BZPZK6W1h6MpsImRJybKMMAMI6T2foQpZdm"
        # openai.proxy = {
        #     "http":"http://127.0.0.1:1080",
        #     "https": "http://127.0.0.1:1080"
        # }
        openai.Model.list()

    def chat(self, inputs=None, model_name="gpt-4"):
        MODEL = model_name
        messages = []
        for input in inputs:
            messages.append({"role": "system", "content": input})

        response = openai.ChatCompletion.create(
            model=MODEL,
            messages=messages,
            temperature=0,
        )
        story = response['choices'][0]['message']['content'];
        return story
        # return {"name": self.get_story_name(story), "story": story}

    def chat_dist(self, inputs=None, model_name="gpt-4"):
        MODEL = model_name
        messages = []
        for input in inputs:
            messages.append({"role": "system", "content": input})
        startTime = time.time()
        response = openai.ChatCompletion.create(
            model=MODEL,
            messages=messages,
            temperature=0,
        )

        print(response)
        story = response['choices'][0]['message']['content']
        t = time.time() - startTime
        print("------------")
        print("%s 生成完成\n使用时间：%s \n生成内容：%s \n" % (model_name, t, story))
        print("------------")
        return {"content": story, "model_name": model_name, "executeTime": t}

    def chat_outline(self, outline=None):
        # MODEL = "gpt-3.5-turbo"
        MODEL = "gpt-4"
        response = openai.ChatCompletion.create(
            model=MODEL,
            messages=[
                {"role": "system",
                 "content": "你现在是一个优秀的故事创作者，你总是能产生一些天马行空的想法，你可以再知道故事大纲的情况下，很好的补充一个完整的故事"},
                {"role": "user", "content": "故事大纲：" + outline},
            ],
            temperature=0,
        )
        story = response['choices'][0]['message']['content']
        return story

    def get_story_name(self, story):
        pattern = r"《(.*?)》"
        match = re.search(pattern, story)

        if match:
            story_name = match.group(1)
            print("故事名称：", story_name)
        else:
            print("未找到故事名称。")

    def chat_storyboard(self, story=None):
        # MODEL = "gpt-3.5-turbo"
        MODEL = "gpt-4"
        response = openai.ChatCompletion.create(
            model=MODEL,
            messages=[
                {"role": "system",
                 "content": "你的角色是一个分镜大师，总能把一个故事，拆分成不同的画面，并给只有视觉元素的描述和画面感的文字，接下来我会给以一段故事，你要把这个故事分成四个画面，并用简练的文字描述出来,每个镜头用 || 分割"},
                {"role": "user", "content": story},
            ],
            temperature=0,
        )
        storyboard = response['choices'][0]['message']['content']
        print(storyboard)
        return self.split_storyboard(storyboard)

    def storyboard_to_prompt(self, storyboard=None):

        text = ""
        for i, scene in enumerate(storyboard, 1):
            text += f"{scene} \n"

        MODEL = "gpt-4"
        response = openai.ChatCompletion.create(
            model=MODEL,
            messages=[
                {"role": "system",
                 "content": "你是一个专业的中英文翻译，你的工作是把中文翻译成英文，你现在需要把下面的中文翻译成英文，不需要除了翻译之外的其他话语，每段翻译用 || 分割"},
                {"role": "user", "content": text},
            ],
            temperature=0,
        )
        prompt = response['choices'][0]['message']['content']
        print(prompt)
        return prompt.split("||")

    def split_storyboard(self, storyboard=None):

        ret = []
        scenes = storyboard.split("||")

        # Trim leading and trailing whitespaces from each scene
        scenes = [scene.strip() for scene in scenes]
        for scene in scenes:
            colon_index = scene.find("：")
            if colon_index != -1:  # 仅当找到冒号时才截取内容
                ret.append(scene[colon_index + 1:].strip())

        # Print the list of scenes
        for i, scene in enumerate(ret, 1):
            print(f"镜头{i}：{scene}")

        return ret

# api = OpenAiApi()
# list = ["毛毛背着小书包，离开繁华的城市，远远的可以看到山脚下的小村子，夏天的阳光洒在她的身上。",
#         "毛毛和外公手挽手走在果园里，篮子里装满了各种各样的水果，毛毛一边摘水果，一边听着外公讲述关于果园的故事。",
#         "毛毛在草莓地里快乐的摘着草莓，一边吃一边摘，甜甜的草莓汁液沾在她的嘴角，她的脸上洋溢着幸福的笑容。",
#         "毛毛和村子里的小朋友们在田野里追逐，小溪边嬉戏，树荫下讲故事，夏天的阳光洒在他们身上，那些无忧无虑的日子，让毛毛觉得非常快乐。"]
#
# api.storyboard_to_prompt(list)
#
#
# api.split_storyboard("镜头一：小明站在城市的高楼大厦中，远眺窗外，脸上带着期待的微笑。背景是繁华的城市景象，高楼大厦，车水马龙。||镜头二：小明在果园中，爬在一棵大桃树上，手中拿着一个熟透的桃子，咬一口，汁水流下嘴角。背景是满园的各种颜色的水果，红的、黄的、绿的，看得人眼花缭乱。|| 镜头三：小明和村里的小朋友们在果园里玩耍，他们在树下搭建小木屋，树间挂起吊床，树上挂起秋千，一片欢乐的气氛。|| 镜头四：小明和外公在果园里劳作，外公教他如何修剪树枝，如何施肥，如何挑选熟透的水果。小明满脸的汗水，但脸上却洋溢着满足的笑容。")
#
# story = api.chat_outline()
# storyboard = api.chat_storyboard(story["story"])
#
# print(storyboard)