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

import os
import threading
from typing import List, Dict, Optional
from enum import Enum, auto
import logging
from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageType, JWTOAuthApp, JWTAuth
import time

COZE_API_TOKEN = os.getenv("COZE_API_TOKEN")
COZE_CN_BASE_URL = 'https://api.coze.cn'
VOLCENGINE_API_TOKEN = '5e275c38-44fd-415f-abcf-4b59f6377f72'
VOLCENGINE_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"
VOLCENGINE_MODEL_DEEPSEEK_V3 = "ep-20250213194558-rrmr2"
VOLCENGINE_MODEL_DOUBAO_PRO_1_5 = 'ep-20250307150409-4blz9'

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

class CozeChat:
    def __init__(self, base_url: str, auth_token: Optional[str] = None, auth_app: Optional[JWTOAuthApp] = None):
        if not auth_token and not auth_app:
            raise ValueError("Either auth_token or auth_app must be provided.")
        if auth_token:
            self.coze = Coze(auth=TokenAuth(auth_token), base_url=base_url)
        else:
            self.auth_app = auth_app
            oauth_token = auth_app.get_access_token(ttl=12*3600)
            self.coze = Coze(auth=JWTAuth(oauth_app=auth_app), base_url=base_url)
            self.setup_token_refresh()

    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 size: {}".format(len(response.messages)))
        if response.chat.status != ChatStatus.COMPLETED:
            logging.error("Coze chat not completed: {}".format(response.chat.status))
            return None
        final_response = None
        for message in response.messages:
            if message.type == MessageType.ANSWER:
                final_response = message.content
        return final_response

    def setup_token_refresh(self):
        thread = threading.Thread(target=self.refresh_token_loop)
        thread.start()

    def refresh_token_loop(self):
        while True:
            time.sleep(11*3600)
            if self.auth_app:
                self.auth_app.get_access_token(ttl=12*3600)

    @staticmethod
    def get_oauth_app(client_id, private_key_path, public_key_id, base_url=None) -> JWTOAuthApp:
        if not base_url:
            base_url = COZE_CN_BASE_URL
        with open(private_key_path, "r") as f:
            private_key = f.read()
        jwt_oauth_app = JWTOAuthApp(
            client_id=str(client_id),
            private_key=private_key,
            public_key_id=public_key_id,
            base_url=base_url,
        )
        return jwt_oauth_app

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 = "7491250992952999973"
    # 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)