
import json
import pickle
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

class Tools:
    @staticmethod
    def load_parquet(file_path):
        print("[Loading] Loading local file: " + str(file_path))
        load_path = file_path
        read_pd = pd.read_parquet(load_path)
        read_pd = read_pd.to_json(orient='records')
        read_pd = json.loads(read_pd)
        print("[Loaded]")
        return read_pd

    @staticmethod
    def write_parquet(path, jsonl_obj):
        assert False, "Not implemented"
        return None


    @staticmethod
    def load_jsonl(file_path):
        json_objects = []
        with open(file_path, 'r', encoding='utf8') as f:
            for line in f:
                json_objects.append(json.loads(line.strip()))
        return json_objects

    @staticmethod
    def write_jsonl(path, content):
        with open(path, 'w') as f:
            for line in content:
                f.write(json.dumps(line) + "\n")

    @staticmethod
    def load_tasks(task_path):
        result = dict()
        lines = Tools.load_jsonl(task_path)
        for line in lines:
            result[line['task_id']] = line
        return result

    @staticmethod
    def dump_pickle(path, content):
        with open(path, 'wb') as f:
            pickle.dump(content, f)

    @staticmethod
    def load_pickle(path):
        with open(path, 'rb') as f:
            return pickle.load(f)

    @staticmethod
    def write_file(path, content):
        with open(path, 'w', encoding='utf8') as f:
            f.write(content)
