from collections import namedtuple
from glob import glob
import enum

from b2luigi.core.utils import product_dict, fill_kwargs_with_lists

import b2luigi

from parse import parse


class DataMode(enum.Enum):
    raw = "raw"
    mdst = "mdst"
    cdst = "cdst"
    skimmed_raw = "skimmed_raw"


class DataTask(b2luigi.ExternalTask):
    data_mode = b2luigi.EnumParameter(enum=DataMode)
    experiment_number = b2luigi.IntParameter()
    run_number = b2luigi.IntParameter()
    prefix = b2luigi.Parameter()
    file_name = b2luigi.Parameter()


class RawDataTask(DataTask):
    data_mode = DataMode.raw

    def output(self):
        yield {"raw_output.root": b2luigi.LocalTarget(_build_data_path(self))}


class DstDataTask(DataTask):
    release = b2luigi.Parameter()
    prod = b2luigi.IntParameter()
    database = b2luigi.IntParameter()

    def output(self):
        yield {"full_output.root": b2luigi.LocalTarget(_build_data_path(self))}


class SkimmedRawDataTask(DstDataTask):
    data_mode = DataMode.skimmed_raw

    def output(self):
        yield {"raw_output.root": b2luigi.LocalTarget(_build_data_path(self))}


class MdstDataTask(DstDataTask):
    data_mode = DataMode.mdst

    def output(self):
        yield {"output.root": b2luigi.LocalTarget(_build_data_path(self))}


class CdstDataTask(DstDataTask):
    data_mode = DataMode.cdst

    def output(self):
        yield {"output.root": b2luigi.LocalTarget(_build_data_path(self))}


requires_raw_data = b2luigi.requires(RawDataTask)
requires_skimmed_raw_data = b2luigi.requires(SkimmedRawDataTask)
requires_mdst_data = b2luigi.requires(MdstDataTask)
requires_cdst_data = b2luigi.requires(CdstDataTask)


def _get_dir_structure(data_mode):
    if data_mode == DataMode.mdst:
        return b2luigi.get_setting("mdst_dir_structure",
                                   "/hsm/belle2/bdata/Data/release-{p.release}/DB{p.database:08d}/prod{p.prod:08d}/" +
                                   "e{p.experiment_number:04d}/4S/r{p.run_number:05d}/all/mdst/sub00/" +
                                   "mdst.{p.prefix}.{p.experiment_number:04d}.{p.run_number:05d}.{p.file_name}.root")
    if data_mode == DataMode.cdst:
        return b2luigi.get_setting("cdst_dir_structure",
                                   "/hsm/belle2/bdata/Data/release-{p.release}/DB{p.database:08d}/prod{p.prod:08d}/" +
                                   "e{p.experiment_number:04d}/4S/r{p.run_number:05d}/all/cdst/sub00/" +
                                   "cdst.{p.prefix}.{p.experiment_number:04d}.{p.run_number:05d}.{p.file_name}.root")
    if data_mode == DataMode.skimmed_raw:
        return b2luigi.get_setting("skimmed_raw_dir_structure",
                                   "/hsm/belle2/bdata/Data/release-{p.release}/DB{p.database:08d}/prod{p.prod:08d}/" +
                                   "e{p.experiment_number:04d}/4S/r{p.run_number:05d}/all/raw/sub00/" +
                                   "raw.{p.prefix}.{p.file_name}.{p.experiment_number:04d}.{p.run_number:05d}.root")
    if data_mode == DataMode.raw:
        return b2luigi.get_setting("raw_dir_structure",
                                   "/ghi/fs01/belle2/bdata/Data/Raw/e{p.experiment_number:04d}/r{p.run_number:05d}/sub00/" +
                                   "{p.prefix}.{p.experiment_number:04d}.{p.run_number:05d}.{p.file_name}.root")


def _build_data_path(parameters):
    mode = parameters.data_mode

    dir_structure = _get_dir_structure(mode)
    return dir_structure.format(p=parameters)


def _parse_data_path(data_mode, data_path):
    extracted_kwargs = None

    dir_structure = _get_dir_structure(data_mode)
    extracted_kwargs = parse(dir_structure, data_path)

    if not extracted_kwargs:
        raise ValueError("The path does not fit the given format")
    extracted_kwargs = {key.replace("p.", ""): value for key, value in extracted_kwargs.named.items()}
    return extracted_kwargs


def _get_data_kwargs(data_mode, experiment_number, run_number, prefix=None, file_name=None, **other_kwargs):
    if file_name is None:
        file_name = "*"

    if prefix is None:
        prefix = "*"

    all_kwargs = fill_kwargs_with_lists(data_mode=data_mode, experiment_number=experiment_number, run_number=run_number,
                                        prefix=prefix,
                                        file_name=file_name, **other_kwargs)
    for kwargs in product_dict(**all_kwargs):
        # The build_data_path wants an object, so lets convert the dict to a named tuple
        kwargs = namedtuple('GenericDict', kwargs.keys())(**kwargs)
        for data_file in glob(_build_data_path(kwargs)):
            yield _parse_data_path(data_mode, data_file)


def clone_on_mdst(self, task_class, experiment_number, run_number, release, prod, database, prefix=None, file_name=None,
                  **additional_kwargs):
    # TODO: make database not needed
    for kwargs in _get_data_kwargs(data_mode=DataMode.mdst, experiment_number=experiment_number, run_number=run_number,
                                   release=release,
                                   prod=prod, database=database, prefix=prefix, file_name=file_name):
        yield self.clone(task_class, **kwargs, **additional_kwargs)


def clone_on_cdst(self, task_class, experiment_number, run_number, release, prod, database, prefix=None, file_name=None,
                  **additional_kwargs):
    # TODO: make database not needed
    for kwargs in _get_data_kwargs(data_mode=DataMode.cdst, experiment_number=experiment_number, run_number=run_number,
                                   release=release,
                                   prod=prod, database=database, prefix=prefix, file_name=file_name):
        yield self.clone(task_class, **kwargs, **additional_kwargs)


def clone_on_skimmed_raw(self, task_class, experiment_number, run_number, release, prod, database, prefix=None, file_name=None,
                         **additional_kwargs):
    # TODO: make database not needed
    for kwargs in _get_data_kwargs(data_mode=DataMode.skimmed_raw, experiment_number=experiment_number, run_number=run_number,
                                   release=release,
                                   prod=prod, database=database, prefix=prefix, file_name=file_name):
        yield self.clone(task_class, **kwargs, **additional_kwargs)


def clone_on_raw(self, task_class, experiment_number, run_number, prefix=None, file_name=None, **additional_kwargs):
    for kwargs in _get_data_kwargs(data_mode=DataMode.raw, experiment_number=experiment_number, run_number=run_number,
                                   prefix=prefix,
                                   file_name=file_name):
        yield self.clone(task_class, **kwargs, **additional_kwargs)
