#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#

import os
from typing import List, Dict
from enum import Enum, auto
import logging
from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageContentType, ChatEventType, MessageType

COZE_API_TOKEN = os.getenv("COZE_API_TOKEN")
COZE_CN_BASE_URL = 'https://api.coze.cn'

class ChatServiceType(Enum):
    OPENAI_COMPATIBLE = auto
    COZE_CHAT = auto()

class CozeChat:
    def __init__(self, token, base_url: str):
        self.coze = Coze(auth=TokenAuth(token), base_url=base_url)

    def create(self, bot_id: str, user_id: str, messages: List, custom_variables: Dict):
        response = self.coze.chat.create_and_poll(
            bot_id=bot_id, user_id=user_id, additional_messages=messages,
            custom_variables=custom_variables)
        logging.debug("coze response: {}".format(response.messages))
        for message in response.messages:
            if message.type == MessageType.ANSWER:
                return message.content
        return None

if __name__ == '__main__':
    # Init the Coze client through the access_token.
    coze = Coze(auth=TokenAuth(token=COZE_API_TOKEN), base_url=COZE_CN_BASE_URL)

    # Create a bot instance in Coze, copy the last number from the web link as the bot's ID.
    bot_id = "7479005417885417487"
    # The user id identifies the identity of a user. Developers can use a custom business ID
    # or a random string.
    user_id = "dev_user"

    chat = coze.chat.create_and_poll(
        bot_id=bot_id,
        user_id=user_id,
        additional_messages=[Message.build_user_question_text("北京今天天气怎么样")],
        custom_variables={
            'agent_name': '芳华',
            'agent_age': '25',
            'agent_region': '北京',
            'name': '李明',
            'preferred_nickname': '李叔',
            'age': '70',
            'last_interaction_interval': '12',
            'current_time_period': '上午',
            'if_first_interaction': 'False',
            'if_active_greeting': 'False'
        }
    )

    for message in chat.messages:
        print(message, flush=True)

    if chat.chat.status == ChatStatus.COMPLETED:
        print("token usage:", chat.chat.usage.token_count)