import json
import os
import random
from typing import AnyStr

# parsing OA data files with oasst_data helpers
from oasst_data import read_message_trees, ExportMessageNode

messages: list[ExportMessageNode] = []

this_dir = os.path.dirname(os.path.abspath(__file__))
input_file_path = this_dir + '/2023-11-05_oasst2_all.trees.jsonl.gz'

role_dict = {'prompter': 'user', 'assistant': 'assistant'}
user_names = ['Adam', 'Alice', 'Anne', 'Bob', 'Charlie', 'Cody', 'Corinna', 'Cynthia', 'Fred', 'Grace', 'Jane', 'Paul',
              'Rachel', 'Ramesh']
conversations = []


def visit(node: ExportMessageNode, parents: [ExportMessageNode], user: AnyStr):
    new_parents = parents + [node]
    if not node.replies:  # end of conversation
        conversations.append({'messages': [{'role': role_dict[p.role],
                                            'name': user if role_dict[p.role] != 'assistant' else 'assistant',
                                            'content': p.text
                                            } for p in new_parents]})
    else:
        for reply in node.replies:
            visit(reply, new_parents, user)


for tree in read_message_trees(input_file_path):
    if tree.prompt.lang not in ['en']:  # filtering by language tag (optional)
        continue

    visit(tree.prompt, [], user_names[random.randint(0, len(user_names) - 1)])

print(conversations[0])


def mkdir(path):
    if not os.path.isdir(path):
        os.mkdir(path)


mkdir(this_dir + '/../../data')
mkdir(this_dir + '/../../data/oasst')
with open(this_dir + '/../../data/oasst/oasst_all.jsonl', 'w') as f:
    f.writelines([json.dumps(conv) + '\n' for conv in conversations])