import requests
import time
import json
import re
import tiktoken
from langchain_text_splitters import TokenTextSplitter

ACCESS_TOKEN = "***"
BOT_ID = "***"
USER_ID = "***"

def call_coze_api(query):
    
    url = 'https://api.coze.cn/v3/chat'
    headers = {
        'Authorization': f'Bearer {ACCESS_TOKEN}',
        'Content-Type': 'application/json'
    }
    payload = {
        "bot_id": BOT_ID,
        "user_id": USER_ID,
        "stream": False,
        "auto_save_history": True,
        "additional_messages": [
            {
                "role": "user",
                "content": query,
                "content_type": "text"
            }
        ]
    }

    response = requests.post(url, headers=headers, json=payload)
    
    if response.status_code == 200:

        response_data = response.json().get('data')
        id = response_data['id']
        conversation_id = response_data['conversation_id']

        retrieve_url = f'https://api.coze.cn/v3/chat/retrieve?chat_id={id}&conversation_id={conversation_id}'
        while True:

            print('Starting retrieve the response...')
            retrieve_response = requests.get(retrieve_url, headers=headers)
            if retrieve_response.status_code == 200:
                retrieve_data = retrieve_response.json().get('data')
                if retrieve_data['status'] == 'completed':

                    result_url = f'https://api.coze.cn/v3/chat/message/list?chat_id={id}&conversation_id={conversation_id}'
                    result_response = requests.get(result_url, headers=headers)
                    if result_response.status_code == 200:
                        result_data = result_response.json().get('data')
                        result_dict = next((item for item in result_data if item['type'] == 'answer'), None)
                        return result_dict.get('content')
                    else:
                        print(f"Error: {result_response.status_code}, {result_response.text}")
                        return {"error": result_response.status_code, "message": result_response.text}
                    
            else:
                print(f"Error: {retrieve_response.status_code}, {retrieve_response.text}")
                return {"error": retrieve_response.status_code, "message": retrieve_response.text}
            
            time.sleep(10) # Add a delay to avoid excessive API calls
    else:
        print(f"Error: {response.status_code}, {response.text}")
        return {"error": response.status_code, "message": response.text}

def token_count(text):
    encoding = tiktoken.get_encoding("gpt2")
    tokens = encoding.encode(text)
    token_count = len(tokens)
    return token_count

def calculate_chunk_size(token_count: int, token_limit: int) -> int:
    if token_count <= token_limit:
        return token_count
    num_chunks = (token_count + token_limit - 1) // token_limit
    chunk_size = token_count // num_chunks
    remaining_tokens = token_count % token_limit
    if remaining_tokens > 0:
        chunk_size += remaining_tokens // num_chunks
    return chunk_size

def get_text_chunks(text,chunk_size) : 
    splitter = TokenTextSplitter(
        chunk_size=int(chunk_size),
        chunk_overlap=0,
    )
    initial_chunks = splitter.split_text(text)  
    # 定义不同语言的句子分隔符
    sentence_delimiters = re.compile(r'[。！？.!?]')
    # 进一步处理每个初步分割块
    output = []
    current_chunk = initial_chunks[0] if initial_chunks else ''
    for i in range(1, len(initial_chunks)):
        sentences = sentence_delimiters.split(initial_chunks[i])
        if sentences:
            current_chunk += sentences[0] # 拼接第一个句子到当前块
            output.append(current_chunk.strip()) # 将当前块加入输出数组
            current_chunk = ''.join(sentences[1:]) # 剩余的句子作为新的当前块
    # 将最后一个块加入输出数组
    if current_chunk.strip():
        output.append(current_chunk.strip())       
    return output

# 示例调用
with open('source.txt', 'r') as file:
    query = file.read().strip()

#tokens计算、chunk_size计算、分割chunks
token_count = token_count(query)
chunk_size = calculate_chunk_size(token_count, 800)
chunks = get_text_chunks(query,chunk_size)

results = []
for index in range(len(chunks)):
    tagged_text = chunks[index]
    """
    如果需要增加上下文，采用如下方式；注意需要提示词配合
    tagged_text = (
            "".join(chunks[0:index])
            + "<TRANSLATE_THIS>"
            + chunks[index]
            + "</TRANSLATE_THIS>"
            + "".join(chunks[index + 1 :])
    )
    """
    result = call_coze_api(tagged_text)
    result_dict =json.loads(result)
    results.append(result_dict.get('output'))

output = " ".join(results)
with open('target.txt', 'w') as file:
    file.write(output)