from abc import ABC
from openai import OpenAI
import os
import base64
from io import BytesIO

class Base(ABC):
    def describe(self, image, max_tokens=300):
        raise NotImplementedError("Please implement encode method!")


class GptV4(Base):
    def __init__(self):
        import openapi
        openapi.api_key = os.environ["OPENAPI_KEY"]
        self.client = OpenAI()

    def describe(self, image, max_tokens=300):
        buffered = BytesIO()
        try:
            image.save(buffered, format="JPEG")
        except Exception as e:
            image.save(buffered, format="PNG")
        b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

        res = self.client.chat.completions.create(
            model="gpt-4-vision-preview",
            messages=[
              {
                "role": "user",
                "content": [
                  {
                    "type": "text",
                    "text": "请用中文详细描述一下图中的内容，比如时间，地点，人物，事情，人物心情等。",
                  },
                  {
                    "type": "image_url",
                    "image_url": {
                      "url": f"data:image/jpeg;base64,{b64}"
                    },
                  },
                ],
              }
            ],
            max_tokens=max_tokens,
        )
        return res.choices[0].message.content.strip()


class QWen(Base):
    def chat(self, system, history, gen_conf):
        from http import HTTPStatus
        from dashscope import Generation
        from dashscope.api_entities.dashscope_response import Role
        # export DASHSCOPE_API_KEY=YOUR_DASHSCOPE_API_KEY
        response = Generation.call(
                    Generation.Models.qwen_turbo,
                    messages=messages,
                    result_format='message'
        )
        if response.status_code == HTTPStatus.OK:
            return response.output.choices[0]['message']['content']
        return response.message