import uuid
import hydra
import numpy as np
from tqdm import tqdm
from sgcrl.utils.imports import get_class, instantiate_class
from sgcrl.data.episodes_readers.base_episodes_readers import EpisodesReader
from sgcrl.data.dbs.on_disk import PytorchOnDiskEpisodesDB

def store_in_episode_db(
    directory: str,
    episodes_reader: EpisodesReader,
):
    episodes_db = PytorchOnDiskEpisodesDB(directory=directory)

    # Store in db
    for episode in tqdm(episodes_reader, desc='writing episodes in db'):
        episodes_db.write("MyApp", str(uuid.UUID(int=np.random.randint(1e7), version=4)), "action", episode)

def store_image_out_episode_db(
    directory: str,
    episodes_reader: EpisodesReader,
):
    # episodes_db = PytorchOnDiskEpisodesDB(directory=directory)

    for episode in tqdm(episodes_reader, desc='writing episodes in db'):
        print(episode.keys())
        raise SystemExit(0)

if __name__ == "__main__":
    @hydra.main(version_base=None, config_path="./yamls",config_name="visual_antmaze_large")
    def main(cfg):
        
        directory = f'./episodes_dbs/{cfg.dataset_name}/'

        # store
        episodes_reader = instantiate_class(cfg.episodes_reader)
        store_image_out_episode_db(directory=directory,episodes_reader=episodes_reader)

        # verify
        # episodes_db = PytorchOnDiskEpisodesDB(directory=directory)
        # for id in tqdm(episodes_db.get_ids()):
        #     episode = episodes_db[id]
        #     observations = episode['observations']

    main()
