import time
import os
import pickle
from collections import defaultdict


# os.makedirs(".wandb_cache", exist_ok=True)


class WandbCache:
    def __init__(self, cache_dir=".wandb_cache"):
        self.cache_dir = cache_dir
        os.makedirs(cache_dir, exist_ok=True)

    def get_cache_path(self, run_id, keys):
        # Sort keys to ensure consistent filenames
        key_str = "_".join(sorted(keys))
        filename = f"{run_id}_{key_str}.pkl"
        filename = filename.replace("/", "__")
        return os.path.join(self.cache_dir, filename)

    def get_history(self, run, keys):
        cache_path = self.get_cache_path(run.id, keys)

        # Try to load from cache first
        if os.path.exists(cache_path):
            try:
                with open(cache_path, 'rb') as f:
                    cached_data = pickle.load(f)
                    # print(f"Loading history from cache for run {run.id}")
                    return cached_data
            except:
                # print(f"Failed to load cache for run {run.id}, fetching from wandb")
                pass

        # If not in cache or cache invalid, fetch from wandb
        print(f"Fetching history from wandb for run {run.id}")
        history = run.history(keys=keys)

        # Only save to cache if run is finished
        if run.state == "finished":
            try:
                with open(cache_path, 'wb') as f:
                    pickle.dump(history, f)
                print(f"Saved history to cache for run {run.id}")
            except Exception as e:
                print(f"Failed to save cache for run {run.id}: {e}")

        return history

# Initialize cache
# wandb_cache = WandbCache()

# Use cache to get history
# df = wandb_cache.get_history(run, [x_metric] + y_metrics)