import os
import json
from dotenv import load_dotenv
import argparse
import yaml

load_dotenv()
RESULTS_STATION_PATH = os.getenv("RESULTS_STATION_PATH")

data_file_map = {
    'ifeval': 'IFEval',


}

data_prefix_map = {


}

with open('dataset-index.yml', 'r') as f1:
    data_list = yaml.load(f1, Loader=yaml.FullLoader)
f1.close()
data_searchable_list = [next(iter(i.keys())) for i in data_list]

def parse_args():
    parser = argparse.ArgumentParser(description='connect to results station')

    parser.add_argument('-sp', '--station-path', type=str, default=None, help='if no env path, use this.')
    parser.add_argument('-p', '--my-path', type=str, default=None, help='your operation path.')
    parser.add_argument('-op', '--operation', type=str, default='d', help='u:update, d:download, ls: show dataset and model options')
    parser.add_argument('-d', '--dataset', type=str, default='mmlu_pro', help='target dataset name')
    parser.add_argument('-m', '--model', type=str, default='deepseek-v2_5-turbomind', help='target model name')
    # parser.add_argument('-all', '--all-transfer', action='store_true', default=False, help='transfer all files under the path')

    args = parser.parse_args()
    return args


def read_json(path):
    results = []
    for i in path:
        with open(i, 'r') as f:
            results.append(json.load(f))
        f.close()
    return results


def load_json_files_by_prefix(prefix, target_path):
    if prefix in data_file_map.keys():
        prefix = data_file_map[prefix]
    result_dict = {}
    for filename in os.listdir(target_path):
        if filename.startswith(prefix) and filename.endswith(".json"):
            file_path = os.path.join(target_path, filename)
            with open(file_path, "r", encoding="utf-8") as file:
                json_data = json.load(file)
            result_dict[os.path.splitext(filename)[0]] = json_data
    return result_dict


def main(path, mypath, args):

    if args.dataset not in data_searchable_list:
        raise ValueError('invalid dataset input!')

    update_path = path + args.dataset if path[-1] == '/' else path + '/' + args.dataset
    update_filename = args.dataset + '_' + args.model + '.json'
    update_goal = update_path + '/' + update_filename

    # update from your path to result station
    if args.operation == 'u':
        mypath_prediction = (mypath + 'predictions/' + args.model) if mypath[-1] == '/' else (mypath + '/predictions/' + args.model)
        mypath_result = (mypath + 'results/' + args.model) if mypath[-1] == '/' else (mypath + '/results/' + args.model)

        if os.path.exists(mypath_prediction) and os.path.exists(mypath_result):

            result_dict = load_json_files_by_prefix(args.dataset, mypath_result)
            prediction_list = []
            for i in result_dict.keys():
                prediction_dict = load_json_files_by_prefix(i, mypath_prediction)
                for j in range(len(prediction_dict)):
                    for k in prediction_dict[i + '_' + str(j)].keys():
                        prediction_list.append(
                            {
                                'prediction': prediction_dict[i + '_' + str(j)][k],
                                'sub_category': i
                            }
                        )
            update_dict = {
                'predictions': prediction_list,
                'results': result_dict,
            }

            if not os.path.exists(update_path):
                os.makedirs(update_path)
            if os.path.exists(update_goal):
                input("This result exists! Press any key to continue...")
            with open(update_goal, 'w', encoding="utf-8") as f:
                json.dump(update_dict, f, ensure_ascii=False, indent=4)
            f.close()


    # read from result station to your path
    if args.operation == 'd':
        if not os.path.exists(update_goal):
            raise ValueError('This result does not exist!')
        with open(update_goal, 'r', encoding="utf-8") as f:
            results = json.load(f)
        f.close()
        legal_key_set = {'predictions', 'results'}
        if set(results.keys()) == legal_key_set and isinstance(results['predictions'], list) and isinstance(results['results'], dict):
            print("Successfully download result from station! you've got a dict with format as follows: \n content['precitions', 'results']")
        else:
            raise ValueError('illegal format of the result!')
        save_path = args.my_path if args.my_path[-1] == '/' else args.my_path + '/'
        save_path += args.dataset + '/'
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        with open(save_path + update_filename, 'w', encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=4)
        f.close()


if __name__ == '__main__':
    args = parse_args()

    if args.operation == 'ls':
        print("----DATASET LIST----")
        print(data_searchable_list)
        print("----MODEL LIST----")

    else:
        if RESULTS_STATION_PATH is not None:
            path = RESULTS_STATION_PATH
        else:
            path = args.station_path
        if path is None:
            raise ValueError('Please appoint the path of results station!')
        if not os.path.exists(path):
            raise ValueError('Not a valid path of results station!')
        mypath = args.my_path
        if mypath is None:
            raise ValueError('Please appoint your own path!')
        if not os.path.exists(mypath):
            raise ValueError('Not a valid path of your own path!')
        main(path, mypath, args)