from tree_traverser import backend, CompressedTree
import datetime
import psutil
from tqdm import tqdm
from pathlib import Path
import json
from more_itertools import chunked
process = psutil.Process()

def massage_request(r):
    return {k : v if isinstance(v, list) else [v]
            for k, v in r.items()}


if __name__ == "__main__":

    config = """
---
type: remote
host: databridge-prod-catalogue1-ope.ewctest.link
port: 10000
engine: remote
store: remote
    """

    request = {
            "class": "d1",
            "dataset": "climate-dt",
            # "date": "19920420",
        }
    
    data_path = Path("data/compressed_tree_climate_dt.json")
    if not data_path.exists():
        compressed_tree = CompressedTree({})
    else:
        compressed_tree = CompressedTree.load(data_path)

    fdb = backend.PyFDB(fdb_config = config)

    visited_path = Path("data/visited_dates.json")
    if not visited_path.exists():
        visited_dates = set()
    else:
        with open(visited_path, "r") as f:
            visited_dates = set(json.load(f))

    today = datetime.datetime.today()
    start = datetime.datetime.strptime("19920420", "%Y%m%d")
    date_list = [start + datetime.timedelta(days=x) for x in range((today - start).days)]
    date_list = [d.strftime("%Y%m%d") for d in date_list if d not in visited_dates] 
    for dates in chunked(tqdm(date_list), 5):
        print(dates[0])
        print(f"Memory usage: {(process.memory_info().rss)/1e6:.1f} MB")

        r = request | dict(date = dates)
        tree = fdb.traverse_fdb(massage_request(r))

        compressed_tree.insert_tree(tree)
        compressed_tree.save(data_path)
        
        for date in dates:
            visited_dates.add(date)

        with open(visited_path, "w") as f:
            json.dump(list(visited_dates), f)
        
        # print(compressed_tree.reconstruct_compressed_ecmwf_style())
