import json
from evaluation.hardConstriant import get_total_cost
from tqdm import tqdm
import math

def load_line_json_data(filename):
    data = []
    with open(filename, 'r', encoding='utf-8') as f:
        for line in f.read().strip().split('\n'):
            unit = json.loads(line)
            data.append(unit)
    return data

def round_up_to_nearest_hundred(number):
    return math.ceil(number / 100) * 100

if __name__ == '__main__':
    query = load_line_json_data('/home/xj/toolAugEnv/code/toolConstraint/data/query/all_reset_budget.jsonl')
    rewritten_query = open('/home/xj/toolAugEnv/code/toolConstraint/data/api_request/query_rewritten_reset_budget_gpt4.txt').readlines()
    assert len(rewritten_query) == len(query)
    for idx in tqdm(range(len(query))):
        # tested_data = json.load(open(f'/home/xj/toolAugEnv/code/toolConstraint/data/annotation/all/annotation_{idx+1}.json'))
        # query[idx]['reset_budget'] = round_up_to_nearest_hundred(get_total_cost(query[idx], tested_data))
        # print(query[idx])
        # break
        query[idx]['query'] = "".join(x for x in rewritten_query[idx].strip().split('\t')[1:])
        query[idx]['budget'] = query[idx]['reset_budget']
        query[idx].pop('reset_budget')


    with open('/home/xj/toolAugEnv/code/toolConstraint/data/query/all.jsonl', 'w') as f:
            for query_unit in query:
                # print(query)
                json.dump(query_unit, f)
                f.write('\n')
            f.close()