import matplotlib.pyplot as plt
from datetime import datetime
import configparser
import os
import platform

from pprint import pprint
from typing import List
import argparse

from databricks_cli.sdk.api_client import ApiClient
from databricks_cli.jobs.api import JobsApi
from databricks_cli.runs.api import RunsApi

from itertools import cycle

cycol = cycle('bgrcmk')


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--profile',
                        help='databricks profile',
                        default='delphi-dev')
    return parser.parse_args()


def get_workers_diffs(run_times: List[tuple]) -> List[tuple]:
    sorted_runs = sorted(run_times, key=lambda d: d[0])

    enriched_runs = []
    previos_dt_change = 0
    for run in sorted_runs:
        if previos_dt_change == 0:
            dt_change = run[0]
            enriched_runs.append((dt_change, run[1]))
            previos_dt_change = dt_change
        else:
            dt_additional = previos_dt_change + (
                (run[0] - previos_dt_change) * 0.99)
            dt_change = run[0]
            enriched_runs.append((dt_additional, 0))
            enriched_runs.append((dt_change, run[1]))
            previos_dt_change = dt_change

    # add values for correct graphical representation

    accumulated_workers = []
    previous_run_workers = 0
    for run in enriched_runs:
        dt_change = datetime.utcfromtimestamp(int(run[0] / 1000))
        accumulated_workers.append((dt_change, previous_run_workers + run[1]))
        previous_run_workers = previous_run_workers + run[1]

    return accumulated_workers


if __name__ == "__main__":
    args = parse_arguments()
    config = configparser.ConfigParser()
    # dbx config path
    if platform.system() == "Windows":
        config_path = os.getenv('HOMEDRIVE', '') + os.getenv(
            'HOMEPATH', '') + '\\.databrickscfg'
    else:
        config_path = os.getenv('HOME', '') + '/.databrickscfg'
    config.read(config_path)

    api_client = ApiClient(host=config[args.profile]['host'],
                           token=config[args.profile]['token'])

    jobs_api = JobsApi(api_client)
    runs_api = RunsApi(api_client)

    instance_type_spark_version_runs = {}

    for job in jobs_api.list_jobs()['jobs']:
        job_metadata = jobs_api.get_job(job_id=job['job_id'])['settings']

        if 'new_cluster' in job_metadata:
            driver_node_type = job_metadata['new_cluster']['driver_node_type_id'] \
                if 'driver_node_type_id' in job_metadata['new_cluster'] \
                else ""
        else:
            # skip jobs without new cluster spec
            continue

        runs = runs_api.list_runs(job_id=job['job_id'], active_only=False, completed_only=True, offset=0, limit=100)['runs'] \
            if 'runs' in runs_api.list_runs(job_id=job['job_id'], active_only=False, completed_only=True, offset=0, limit=100) \
            else []

        run_times = []
        for run in runs:
            if 'num_workers' in run['cluster_spec']['new_cluster']:
                workers = run['cluster_spec']['new_cluster']['num_workers']
            elif 'autoscale' in run['cluster_spec']['new_cluster']:
                workers = run['cluster_spec']['new_cluster']['autoscale'][
                    'max_workers']
            else:
                pprint(run['cluster_spec']['new_cluster'])
                raise
            #if workers == 0:
            #pprint(run)
            #raise
            #run_times.append({'workers' : workers, 'start_time': run['start_time'], 'end_time' : run['end_time']})
            run_times.append((run['start_time'], workers))
            run_times.append((run['end_time'], -workers))

        if 'node_type_id' in job_metadata['new_cluster']:
            key = f"{job_metadata['new_cluster']['node_type_id']}_{job_metadata['new_cluster']['spark_version']}"
        elif 'instance_pool_id' in job_metadata['new_cluster']:
            key = f"{job_metadata['new_cluster']['instance_pool_id']}_{job_metadata['new_cluster']['spark_version']}"
        else:
            key = f"no_data_{job_metadata['new_cluster']['spark_version']}"

        if len(run_times) > 0:
            if key in instance_type_spark_version_runs:
                #new_run_times = instance_type_spark_version_runs[key].extend(run_times)
                instance_type_spark_version_runs[key] += run_times
            else:
                instance_type_spark_version_runs[key] = run_times

        #instance_and_spark_types_for_jobs.append({
        #    'job_id' : job['job_id'],
        #    'name': job_metadata['name'],
        #    'node_type':job_metadata['new_cluster']['node_type_id'],
        #    'driver_node_type' : driver_node_type,
        #    'spark_version' : job_metadata['new_cluster']['spark_version'],
        #    'workers' : workers,
        #    'run_times' : run_times
        #})

    plots = []

    for id, run_times in instance_type_spark_version_runs.items():
        pprint(id)
        pprint(get_workers_diffs(run_times))
        # Plotting both the curves simultaneously
        plt.plot(*zip(*get_workers_diffs(run_times)),
                 color=next(cycol),
                 label=id)

    # Naming the x-axis, y-axis and the whole graph
    plt.xlabel("Time")
    plt.ylabel("Workers")
    # Adding legend, which helps us recognize the curve according to it's color
    plt.legend()
    # To load the display window
    plt.show()