import os
from dotenv import load_dotenv
import google.generativeai as genai
import json

data = None
with open('problmes.json', 'rt') as f:
    data = json.load(f)

problems = [d['problem'].replace('\n\n', '\n') for d in data]

print(len(problems))
print(len(problems)//50)

# .envファイルの読み込み
load_dotenv()

# API-KEYの設定
GOOGLE_API_KEY=os.getenv('GOOGLE_API_KEY')
genai.configure(api_key=GOOGLE_API_KEY)
gemini_pro = genai.GenerativeModel("gemini-1.5-flash")


for i in range(len(problems)//50):
    start = i*50
    end = (i+1)*50
    text_problems = '\n\n'.join(problems[start:end])
    prompt = "以下のテキストを行ごとに日本語に翻訳して，1行ずつ空けて出力してください．ただし，LATEX文字列には変更を加えずそのまま出力してください．\n\n" + text_problems

    print('prompt:\n'+ prompt + '\n\n')
    print("total_tokens: ", gemini_pro.count_tokens(prompt))
    response = gemini_pro.generate_content(prompt)
    print(response.text)
    print('----- meta data -----')
    print(response.usage_metadata)
    print('---------------------')

    output_dir = 'gens'
    file_name = f'gens{start}-{end}.txt'
    with open(os.path.join(output_dir, file_name), 'wt') as f:
        f.write(response.text)