import time
from typing import Dict

import elastic_poc
import snowflake_poc

functions = {
    'elastic': {
        'playlist_positions_over_time':
            elastic_poc.playlist_positions_over_time,
        'current_placements_in_playlist':
            elastic_poc.current_placements_in_playlist
    },
    'snowflake': {
        'playlist_positions_over_time':
            snowflake_poc.playlist_positions_over_time,
        'current_placements_in_playlist':
            snowflake_poc.current_placements_in_playlist
    }
}


def average(times):
    return sum(times) / len(times)


def time_function(store: str, operation: str, args: Dict, client):
    fn = functions[store][operation]
    t_start = time.time()
    fn(args, client)
    t_end = time.time()
    return t_end - t_start


def time_n_calls(store: str, operation: str, args: Dict, client, n=10):
    fn = functions[store][operation]
    print(f'Getting average of {n} calls of {store} {operation}')
    times = []
    for _ in range(n):
        t_start = time.time()
        fn(args, client)
        t_end = time.time()
        times.append(t_end - t_start)
    return average(times)


def time_n_calls_v2(function, args, n=10):
    times = []
    for _ in range(n):
        t_start = time.time()
        function(args)
        t_end = time.time()
        times.append(t_end - t_start)
    return average(times)