#encoding = utf8

import logging
import time

from nlp.nlp_base import NLPBase
from volcenginesdkarkruntime import AsyncArk

logger = logging.getLogger(__name__)


class DouBao(NLPBase):
    def __init__(self, split, callback=None):
        super().__init__(split, callback)
        # Access Key ID
        # AKLTYTdmOTBmNWFjODkxNDE2Zjk3MjU0NjRhM2JhM2IyN2Y
        # AKLTNDZjNTdhNDlkZGE3NDZjMDlkMzk5YWQ3MDA4MTY1ZDc
        # Secret Access Key
        # WmpRelltRXhNbVkyWWpnNU5HRmpNamc0WTJZMFpUWmpOV1E1TTJFME1tTQ==
        # TkRJMk1tTTFZamt4TkRVNE5HRTNZMkUyTnpFeU5qQmxNMkUwWXpaak1HRQ==
        # endpoint_id
        # ep-20241008152048-fsgzf
        # api_key
        # c9635f9e-0f9e-4ca1-ac90-8af25a541b74
        # api_ky
        # eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJhcmstY29uc29sZSIsImV4cCI6MTczMDk2NTMxOSwiaWF0IjoxNzI4MzczMzE5LCJ0IjoidXNlciIsImt2IjoxLCJhaWQiOiIyMTAyMjc3NDc1IiwidWlkIjoiMCIsImlzX291dGVyX3VzZXIiOnRydWUsInJlc291cmNlX3R5cGUiOiJlbmRwb2ludCIsInJlc291cmNlX2lkcyI6WyJlcC0yMDI0MTAwODE1MjA0OC1mc2d6ZiJdfQ.BHgFj-UKeu7IGG5VL2e6iPQEMNMkQrgmM46zYmTpoNG_ySgSFJLWYzbrIABZmqVDB4Rt58j8kvoORs-RHJUz81rXUlh3BYl9-ZwbggtAU7Z1pm54_qZ00jF0jQ6r-fUSXZo2PVCLxb_clNuEh06NyaV7ullZwUCyLKx3vhCsxPAuEvQvLc_qDBx-IYNT-UApVADaqMs-OyewoxahqQ7RvaHFF14R6ihmg9H0uvl00_JiGThJveszKvy_T-Qk6iPOy-EDI2pwJxdHMZ7By0bWK5EfZoK2hOvOSRD0BNTYnvrTfI0l2JgS0nwCVEPR4KSTXxU_oVVtuUSZp1UHvvkhvA
        self.__token = 'c9635f9e-0f9e-4ca1-ac90-8af25a541b74'
        self.__client = AsyncArk(api_key=self.__token)

    async def _request(self, question):
        t = time.time()
        logger.info(f'_request:{question}')
        print(f'-------dou_bao ask:', question)
        try:
            stream = await self.__client.chat.completions.create(
                model="ep-20241008152048-fsgzf",
                messages=[
                    {"role": "system", "content": "你是测试客服，是由字节跳动开发的 AI 人工智能助手"},
                    {"role": "user", "content": question},
                ],
                stream=True
            )
            sec = ''
            async for completion in stream:
                sec = sec + completion.choices[0].delta.content
                sec, message = self._split_handle.handle(sec)
                if len(message) > 0:
                    self._on_callback(message)
            self._on_callback(sec)
            await stream.close()
        except Exception as e:
            print(e)
        logger.info(f'_request:{question}, time:{time.time() - t:.4f}s')
        print(f'-------dou_bao nlp time:{time.time() - t:.4f}s')

    async def _on_close(self):
        print('AsyncArk close')
        if self.__client is not None and not self.__client.is_closed():
            await self.__client.close()

'''
if __name__ == "__main__":
    # print(get_dou_bao_api())
    dou_bao = DouBao()
    dou_bao.ask('你好。')
    dou_bao.ask('你好，你是谁？')
    dou_bao.ask('你能做什么？')
    dou_bao.ask('介绍一下，我自己。')
    count = 1000
    sec = ''
    while count >= 0:
        count = count - 1
        if nlp_queue.empty():
            time.sleep(0.1)
            continue
        sec = sec + nlp_queue.get(block=True, timeout=0.01)

        pattern = r'[，。、；？！,.!?]'
        match = re.search(pattern, sec)
        if match:
            pos = match.start() + 1
            print(sec[:pos])
            sec = sec[pos:]
    print(sec)


    dou_bao.stop()
'''