from os import path
from threading import Event, Thread
from time import sleep

from nabu.utils import partition_dict
from silx.io.url import DataUrl
from tqdm import tqdm


def format_dataurl_to_scan(scan_url):
    if isinstance(scan_url, str):
        # assuming format like silx:///path/to/data?entry/path/in/h5
        fname, h5path = scan_url.split("?")
    elif isinstance(scan_url, DataUrl):
        fname = scan_url.file_path()
        h5path = scan_url.data_path()
    scan_name = path.basename(fname)
    scan_num = h5path.split("/")[0]
    res = "%s:%s" % (scan_name, scan_num)
    return res


def check_progress(scans_and_partial_output_files):
    res = dict.fromkeys(scans_and_partial_output_files.keys(), 0)
    for scan, partial_output_files in scans_and_partial_output_files.items():
        output_exists = [path.isfile(fname) for fname in partial_output_files]
        res[scan] = sum(output_exists)
    return res


def do_in_thread(func, *args, **kwargs):
    t = Thread(target=func, args=args, kwargs=kwargs)
    t.start()
    return t


class ScansIntegrationMonitor(Thread):
    """Thread class with a stop() method. The thread itself has to check
    regularly for the stopped() condition.
    """

    def __init__(self, period, scans_to_partial_output_files):
        self.period = period
        self.scans_to_partial_output_files = scans_to_partial_output_files
        self.n_scans = len(self.scans_to_partial_output_files)
        self._stop_event = Event()
        super().__init__(target=self._monitor)

    def _monitor(self):
        self.scan_status = dict.fromkeys(self.scans_to_partial_output_files.keys(), 0)
        self.scan_finished = dict.fromkeys(
            self.scans_to_partial_output_files.keys(),
            False,
        )
        while not (all(list(self.scan_finished.values()))) and not (
            self._stop_event.is_set()
        ):
            for _i, scan_url_txt in enumerate(self.scans_to_partial_output_files):
                partial_out_files = self.scans_to_partial_output_files[scan_url_txt]
                output_exists = [path.isfile(fname) for fname in partial_out_files]
                n_processed = sum(output_exists)
                self.scan_status[scan_url_txt] = n_processed
                if n_processed == len(partial_out_files):
                    self.scan_finished[scan_url_txt] = True
            sleep(self.period)

    def stop(self):
        self._stop_event.set()

    def stopped(self):
        return self._stop_event.is_set()


def monitor_scans(
    scans_info,
    period=1,
    n_threads=1,
    callback=None,
    callback_args=None,
    callback_kwargs=None,
):
    """
    Monitor the integration progress.

    Parameters
    ----------
    scans_info: dict
        Dictionary whose structure is explained in ProcessingConfig.build_tasks_distribution()
    period: int, optional
        Update period in seconds. Each thread will periodically check the output files,
        so using a small period and/or a large number of threads might hammer the filesystem.
    n_threads: int, optional
        Number of threads for monitoring
    callback: callable, optional
        Callback function to be called each time a scan processing is complete.
        This is a function where the first two parameters are:
          - A HDF5Dataset object: data structure holding information on the scan
          - a list of strings: the path to partial output files
    callback_args: list, optional
    callback_kwargs: dict, optional
    """
    # Configure callback, if needed
    callback_args = callback_args or []
    callback_kwargs = callback_kwargs or {}

    scans_to_partial_output_files = {
        url: scans_info[url]["partial_output_files"] for url in scans_info
    }
    scan_finished = dict.fromkeys(scans_to_partial_output_files.keys(), False)

    # Split scans among threads
    scans_to_partial_output_files_parts = partition_dict(
        scans_to_partial_output_files,
        n_threads,
    )

    # Spawn monitoring threads and initialize progressbars
    threads = []
    progress_bars = {}
    monitor = {}
    for i, scans_to_partial_output_files_part in enumerate(
        scans_to_partial_output_files_parts,
    ):
        t = ScansIntegrationMonitor(period, scans_to_partial_output_files_part)
        threads.append(t)
        for j, scan_url_txt in enumerate(scans_to_partial_output_files_part):
            progress_bar = tqdm(
                total=len(scans_to_partial_output_files_part[scan_url_txt]),
                position=i * t.n_scans + j,
                desc=format_dataurl_to_scan(scan_url_txt),
            )
            progress_bars[scan_url_txt] = progress_bar
            monitor[scan_url_txt] = t

    # Start monitoring threads
    for t in threads:
        t.start()

    # Main thread updates the progress bars
    try:
        while not (all(scan_finished.values())):
            for scan_url_txt in scans_to_partial_output_files:
                thread = monitor[scan_url_txt]
                progress_bar = progress_bars[scan_url_txt]
                progress_bar.update(thread.scan_status[scan_url_txt] - progress_bar.n)
                if thread.scan_finished[scan_url_txt]:
                    scan_finished[scan_url_txt] = True
            sleep(0.5)

    except KeyboardInterrupt:
        print("Stopping monitoring")

    finally:
        for t in threads:
            t.stop()
