import os
from typing import Any, Callable, Dict, List, Optional, Set, Union

import pandas as pd
from pyteomics.mgf import MGFBase
from pyteomics.mzid import MzIdentML
from pyteomics.mzml import MzML

from mmproteo.utils import log, visualization
from mmproteo.utils.formats.mgf import read_mgf
from mmproteo.utils.formats.mz import read_mzid, read_mzml


def iter_entries(iterator: Union[MGFBase, MzIdentML, MzML], logger: log.Logger = log.DEFAULT_LOGGER) \
        -> List[Dict[str, Any]]:
    logger.debug("Iterator type = " + str(type(iterator)))
    entries = list(iterator)
    logger.debug("Length: %d" % len(entries))
    if len(entries) > 0:
        logger.debug("Example:\n" + visualization.pretty_print_json(entries[0]))
    return entries


def read_parquet(filename: str, logger: log.Logger = log.DEFAULT_LOGGER) -> pd.DataFrame:
    return pd.read_parquet(filename)


_FILE_READING_CONFIG: Dict[str, Callable[[str, log.Logger], pd.DataFrame]] = {
    "mgf": read_mgf,
    "mzid": read_mzid,
    "mzml": read_mzml,
    "parquet": read_parquet,
}


def get_readable_file_extensions() -> Set[str]:
    return set(_FILE_READING_CONFIG.keys())


def read(filename: str, filename_col: Optional[str] = "%s_filename", logger: log.Logger = log.DEFAULT_LOGGER) \
        -> pd.DataFrame:
    _, ext = separate_extension(filename=filename, extensions=get_readable_file_extensions())

    if len(ext) > 0:
        logger.debug("Started reading %s file '%s'" % (ext, filename))
        df = _FILE_READING_CONFIG[ext](filename, logger)
        logger.info("Finished reading %s file '%s'" % (ext, filename))
    else:
        raise NotImplementedError
    if filename_col is not None:
        if "%s" in filename_col:
            col = filename_col % ext
        else:
            col = filename_col
        if col in df.columns:
            logger.warning("'%s' already exists in DataFrame columns. Please choose a different column name.")
        else:
            df[col] = filename.split(os.sep)[-1]
    return df


def separate_extension(filename: str, extensions: Set[str]) -> (str, str):
    lower_filename = filename.lower()
    longest_extension = ""
    for ext in extensions:
        if lower_filename.endswith("." + ext) and len(longest_extension) < len(ext):
            longest_extension = ext

    if len(longest_extension) > 0:
        base_filename = filename[:len(filename) - len(longest_extension) - len(".")]
        return base_filename, longest_extension
    else:
        return filename, ""


