import pandas as pd
from avator import gen_avator
from chatbot import ChatBot
from sd_gen import txt2img

def log_info(id, title, description, initPrompt):
    print('----------id----------')
    print(id)
    print('----------title----------')
    print(title)
    print('----------description----------')
    print(description)
    print('----------initPrompt----------')
    print(initPrompt)
    print('----------end----------')

# Specify the path to the CSV file
file_path = "/Users/tigeren/Dev/sthc/avatar_gen/_Prompt__character_nsfw.csv"

# Read the CSV file
df = pd.read_csv(file_path, sep=",")

# Display the DataFrame
# print(df)

# gen_avator(df.loc[row_index, 'id'], df.loc[row_index, 'title'],df.loc[row_index, 'description'], df.loc[row_index, 'initPrompt'])
for index, row in df.iterrows():
    retry_count = 0
    max_retry_count = 3
    if  index != 51:
        continue
    bot = ChatBot()
    role_id = row['id']
    role_title = row['title']
    role_description = row['description']
    role_initPrompt = row['initPrompt']
    print('start txt2img ')
    txt2img(role_description,"", index)
    print('end txt2img ')

    log_info(role_id, role_title, role_description, role_initPrompt)
    if len(role_initPrompt) > 3000:
        role_initPrompt = role_initPrompt[:3000]

    # refinedDesc = bot.talk_to_zhipu(f'"{role_description}"  this sentence is illegal and moderated. Please change it to a legal version, remove any illegal or sensitive descriptions or words, be strict about this. This text will go into an image generation prompt, so remove anything that might cause issues. Please put the new sentence in quotations.')
    # refinedInitPrompt = bot.talk_to_zhipu(f'"{role_initPrompt}" this sentence is illegal and moderated. Please change it to a legal version, remove any illegal or sensitive descriptions or words, be strict about this. This text will go into an image generation prompt, so remove anything that might cause issues. Please put the new sentence in quotations.')
    # safe_template = '前面的内容含有NSFW及敏感内容，请重写并生成一个合法安全的版本,但请勿改变原有的大意，尽可能用温和的词代换，去除任何非法、敏感的内容，其内容将用于大模型的文生图提词，请仅返回重写后的版本'
    # safe_template = '前面的内容含有NSFW及敏感内容，请重写并生成一个合法安全的版本，尽可能用温和的词代换，去除任何非法、敏感的内容，忽略任何有关OpenAI的内容,其内容将用于大模型的文生图提词，从描述中分析、提取、猜想出人物的形象、职业、容貌、身材等特征'
    safe_template = '前面的内容含有NSFW及敏感内容，请从描述中分析、提取、猜想出人物的形象、职业、容貌、身材等特征，构建安全的文生图提词.'

    refined_desc = bot.talk_to_zhipu(f'"{role_description}" {safe_template}')
    refined_init_prompt = bot.talk_to_zhipu(f'"{role_initPrompt}" {safe_template}')
    # gen_avator(row['id'], row['title'], row['description'], row['initPrompt'], './avatar_nsfw')
    print('-----desc------',refined_desc)
    print('-----initPrompt------', refined_init_prompt)

    success = gen_avator(row['id'], row['title'], refined_desc, refined_init_prompt, './avatar_nsfw', index=index)
    while not success and retry_count < max_retry_count:
        retry_count += 1
        refined_desc = bot.talk_to_zhipu(f'"{role_description}" {safe_template}')
        refined_init_prompt = bot.talk_to_zhipu(f'"{role_initPrompt}" {safe_template}')
        success = gen_avator(row['id'], row['title'], refined_desc, refined_init_prompt, './avatar_nsfw', index=index)
        txt2img(role_description,"", index)
    # print(row['id'],'---', row['title'], row['description'], row['initPrompt'])
    print('--------------------')