from collections import defaultdict
import copy
from distutils.version import LooseVersion
from functools import lru_cache
import io
import struct
import os
import operator
import re
import numbers

import numpy as np
import pandas as pd

import fsspec
from pandas.api.types import is_categorical_dtype

from fastparquet import __version__

PANDAS_VERSION = LooseVersion(pd.__version__)
created_by = f"fastparquet-python version {__version__} (build 0)"


class ParquetException(Exception):
    """Generic Exception related to unexpected data format when
     reading parquet file."""
    pass


def default_mkdirs(f):
    os.makedirs(f, exist_ok=True)


PATH_DATE_FMT = '%Y%m%d_%H%M%S.%f'


def path_string(o):
    if isinstance(o, pd.Timestamp):
        return o.isoformat()
    return str(o)


default_open = open


def default_remove(paths):
    for path in paths:
        try:
            os.unlink(path)
        except IOError:
            pass
    

def val_from_meta(x, meta):
    try:
        if meta['pandas_type'] == 'categorical':
            return x
        t = np.dtype(meta['numpy_type'])
        if t == "bool":
            return x in [True, "true", "True", 't', "T", 1, "1"]
        return np.dtype(t).type(x)
    except ValueError:
        if meta['numpy_type'] == 'datetime64[ns]':
            return pd.to_datetime(x, format=PATH_DATE_FMT)
        else:
            raise


def val_to_num(x, meta=None):
    """Parse a string as a number, date or timedelta if possible"""
    if meta:
        return val_from_meta(x, meta)
    return _val_to_num(x)


@lru_cache(1000)
def _val_to_num(x):
    if isinstance(x, numbers.Real):
        return x
    if x in ['now', 'NOW', 'TODAY', '']:
        return x
    if type(x) == str and x.lower() == 'nan':
        return x
    if x == "True":
        return True
    if x == "False":
        return False
    try:
        return int(x, base=10)
    except:
        pass
    try:
        return float(x)
    except:
        pass
    try:
        return pd.Timestamp(x)
    except:
        pass
    try:
        # TODO: determine the valid usecases for this, then try to limit the set
        #  ofstrings which may get inadvertently converted to timedeltas
        return pd.Timedelta(x)
    except:
        return x


def ensure_bytes(s):
    return s.encode('utf-8') if isinstance(s, str) else s


def check_column_names(columns, *args):
    """Ensure that parameters listing column names have corresponding columns"""
    for arg in args:
        if isinstance(arg, (tuple, list)):
            missing = set(arg) - set(columns)
            if missing:
                raise ValueError("Following columns were requested but are "
                                 "not available: %s.\n"
                                 "All requested columns: %s\n"
                                 "Available columns: %s"
                                 "" % (missing, arg, columns))


def reset_row_idx(data: pd.DataFrame) -> pd.DataFrame:
    """Reset row (multi-)index as column(s) of the DataFrame.

    Multi-index are stored in columns, one per index level.

    Parameters
    ----------
    data : dataframe

    Returns
    -------
    dataframe
    """
    if isinstance(data.index, pd.MultiIndex):
        for name, cats, codes in zip(data.index.names, data.index.levels,
                                     data.index.codes):
            data = data.assign(**{name: pd.Categorical.from_codes(codes,
                                                                  cats)})
        data.reset_index(drop=True)
    else:
        data = data.reset_index()
    return data


def metadata_from_many(file_list, verify_schema=False, open_with=default_open,
                       root=False, fs=None):
    """
    Given list of parquet files, make a FileMetaData that points to them

    Parameters
    ----------
    file_list: list of paths of parquet files
    verify_schema: bool (False)
        Whether to assert that the schemas in each file are identical
    open_with: function
        Use this to open each path.
    root: str
        Top of the dataset's directory tree, for cases where it can't be
        automatically inferred.
    fs: fsspsec.AbstractFileSystem
        Used in preference to open_with, if given

    Returns
    -------
    basepath: the root path that other paths are relative to
    fmd: metadata thrift structure
    """
    from fastparquet import api

    legacy = True
    if all(isinstance(pf, api.ParquetFile) for pf in file_list):
        pfs = file_list
        file_list = [pf.fn for pf in pfs]
    elif all(not isinstance(pf, api.ParquetFile) for pf in file_list):

        if verify_schema or fs is None or len(file_list) < 3:
            pfs = [api.ParquetFile(fn, open_with=open_with) for fn in file_list]
        else:
            # activate new code path here
            f0 = file_list[0]
            pf0 = api.ParquetFile(f0, open_with=open_with)
            if pf0.file_scheme not in ['empty', 'simple']:
                # set of directories, revert
                pfs = [pf0] + [api.ParquetFile(fn, open_with=open_with) for fn in file_list[1:]]
            else:
                # permits concurrent fetch of footers; needs fsspec >= 2021.6
                size = int(1.4 * pf0._head_size)
                pieces = fs.cat(file_list[1:], start=-size)
                sizes = {path: int.from_bytes(piece[-8:-4], "little") + 8 for
                         path, piece in pieces.items()}
                not_bigenough = [path for path, s in sizes.items() if s > size]
                if not_bigenough:
                    new_pieces = fs.cat(not_bigenough, start=-max(sizes.values()))
                    pieces.update(new_pieces)
                pieces = {k: _get_fmd(v) for k, v in pieces.items()}
                pieces = [(fn, pieces[fn]) for fn in file_list[1:]]  # recover ordering
                legacy = False
    else:
        raise ValueError("Merge requires all PaquetFile instances or none")
    basepath, file_list = analyse_paths(file_list, root=root)

    if legacy:
        # legacy code path
        if verify_schema:
            for pf in pfs[1:]:
                if pf._schema != pfs[0]._schema:
                    raise ValueError('Incompatible schemas')

        fmd = copy.copy(pfs[0].fmd)  # we inherit "created by" field
        rgs = []

        for pf, fn in zip(pfs, file_list):
            if pf.file_scheme not in ['simple', 'empty']:
                for rg in pf.row_groups:
                    rg = copy.copy(rg)
                    rg.columns = [copy.copy(c) for c in rg.columns]
                    for chunk in rg.columns:
                        chunk.file_path = '/'.join(
                            [fn, chunk.file_path if isinstance(chunk.file_path, str) else chunk.file_path.decode()]
                        )
                    rgs.append(rg)

            else:
                for rg in pf.row_groups:
                    rg = copy.copy(rg)
                    rg.columns = [copy.copy(c) for c in rg.columns]
                    for chunk in rg.columns:
                        chunk.file_path = fn
                    rgs.append(rg)

        fmd.row_groups = rgs
        fmd.num_rows = sum(rg.num_rows for rg in fmd.row_groups)
        return basepath, fmd

    for rg in pf0.fmd.row_groups:
        # chunks of first file, which would have file_path=None
        rg.columns[0].file_path = f0[len(basepath):].lstrip("/")

    rgs0 = pf0.fmd.row_groups
    for k, v in pieces:
        # Set file paths on other files
        rgs = v.row_groups or []
        for rg in rgs:
            rg.columns[0].file_path = k[len(basepath):].lstrip("/")
        rgs0.extend(rgs)
    pf0.fmd.row_groups = rgs0
    pf0.fmd.num_rows = sum(rg.num_rows for rg in pf0.fmd.row_groups)
    return basepath, pf0.fmd


def _get_fmd(inbytes):
    from .cencoding import from_buffer

    f = io.BytesIO(inbytes)
    f.seek(-8, 2)
    head_size = struct.unpack('<i', f.read(4))[0]
    f.seek(-(head_size + 8), 2)
    data = f.read(head_size)
    return from_buffer(data, "FileMetaData")


# simple cache to avoid re compile every time
seps = {}


def ex_from_sep(sep):
    """Generate regex for category folder matching"""
    if sep not in seps:
        if sep in r'\^$.|?*+()[]':
            s = re.compile(r"([a-zA-Z_0-9]+)=([^\\{}]+)".format(sep))
        else:
            s = re.compile("([a-zA-Z_0-9]+)=([^{}]+)".format(sep))
        seps[sep] = s
    return seps[sep]


def analyse_paths(file_list, root=False):
    """Consolidate list of file-paths into  parquet relative paths"""
    path_parts_list = [join_path(fn).split('/') for fn in file_list]
    if root is False:
        basepath = path_parts_list[0][:-1]
        for i, path_parts in enumerate(path_parts_list):
            j = len(path_parts) - 1
            for k, (base_part, path_part) in enumerate(
                    zip(basepath, path_parts)):
                if base_part != path_part:
                    j = k
                    break
            basepath = basepath[:j]
        l = len(basepath)

    else:
        basepath = join_path(root).split('/')
        l = len(basepath)
        assert all(p[:l] == basepath for p in path_parts_list
                   ), "All paths must begin with the given root"
    out_list = []
    for path_parts in path_parts_list:
        out_list.append('/'.join(path_parts[l:]))  # use '/'.join() instead of join_path to be consistent with split('/')

    return '/'.join(basepath), out_list  # use '/'.join() instead of join_path to be consistent with split('/')


def infer_dtype(column):
    try:
        return pd.api.types.infer_dtype(column, skipna=False)
    except AttributeError:
        return pd.lib.infer_dtype(column)


def groupby_types(iterable):
    groups = defaultdict(list)
    for x in iterable:
        groups[type(x)].append(x)
    return groups


def get_column_metadata(column, name):
    """Produce pandas column metadata block"""
    # from pyarrow.pandas_compat
    # https://github.com/apache/arrow/blob/master/python/pyarrow/pandas_compat.py
    inferred_dtype = infer_dtype(column)
    dtype = column.dtype
    if str(dtype) == "bool":
        # pandas accidentally calls this "boolean"
        inferred_dtype = "bool"

    if is_categorical_dtype(dtype):
        extra_metadata = {
            'num_categories': len(column.cat.categories),
            'ordered': column.cat.ordered,
        }
        dtype = column.cat.codes.dtype
    elif hasattr(dtype, 'tz'):
        try:
            stz = str(dtype.tz)
            if "UTC" in stz and ":" in stz:
                extra_metadata = {'timezone': stz.strip("UTC")}
            elif len(str(stz)) == 3:  # like "UTC", "CET", ...
                extra_metadata = {'timezone': str(stz)}
            elif getattr(dtype.tz, "zone", False):
                extra_metadata = {'timezone': dtype.tz.zone}
            elif "pytz" not in stz:
                pd.Series([pd.to_datetime('now', utc=True)]).dt.tz_localize(stz)
                extra_metadata = {'timezone': stz}
            elif "Offset" in stz:
                extra_metadata = {'timezone': f"{dtype.tz._minutes // 60:+03}:00"}
            else:
                raise KeyError
        except Exception as e:
            raise ValueError("Time-zone information could not be serialised: "
                             "%s, please use another" % str(dtype.tz)) from e
    else:
        extra_metadata = None

    if isinstance(name, tuple):
        name = str(name)
    elif not isinstance(name, str):
        raise TypeError(
            'Column name must be a string. Got column {} of type {}'.format(
                name, type(name).__name__
            )
        )

    return {
        'name': name,
        'field_name': name,
        'pandas_type': {
            'string': 'unicode',
            'datetime64': (
                'datetimetz' if hasattr(dtype, 'tz')
                else 'datetime'
            ),
            'integer': str(dtype),
            'floating': str(dtype),
        }.get(inferred_dtype, inferred_dtype),
        'numpy_type': get_numpy_type(dtype),
        'metadata': extra_metadata,
    }


def get_numpy_type(dtype):
    if is_categorical_dtype(dtype):
        return 'category'
    elif "Int" in str(dtype):
        return str(dtype).lower()
    elif str(dtype) == "boolean":
        return "bool"
    elif str(dtype) == "string":
        return "object"
    else:
        return str(dtype)


def get_file_scheme(paths):
    """For the given row groups, figure out if the partitioning scheme

    Parameters
    ----------
    paths: list of str
        normally from row_group.columns[0].file_path

    Returns
    -------
    'empty': no rgs at all
    'simple': all rgs in a single file
    'flat': multiple files in one directory
    'hive': directories are all `key=value`; all files are at the same
        directory depth
    'drill': assume directory names are labels, and field names are of the
        form dir0, dir1; all files are at the same directory depth
    'other': none of the above, assume no partitioning
    """
    if not paths:
        return 'empty'
    if set(paths) == {None}:
        return 'simple'
    if None in paths:
        return 'other'
    parts = [p.split('/') for p in paths]
    lens = [len(p) for p in parts]
    if len(set(lens)) > 1:
        return 'other'
    if set(lens) == {1}:
        return 'flat'
    matches = all(all("=" in p[1:-1] for p in part[:-1]) for part in parts)
    return "hive" if matches else "drill"


def join_path(*path):
    return "/".join([p.replace("\\", "/").rstrip("/") for p in path if p])


_json_decoder = [None]


def json_decoder():
    import importlib
    if _json_decoder[0] is None:
        for lib in ['orjson', 'ujson', 'rapidjson', 'json']:
            try:
                package = importlib.import_module(lib)
                _json_decoder[0] = package.loads
                break
            except (ImportError, AttributeError):
                pass
    return _json_decoder[0]


def _strip_path_tail(paths) -> set:
    return {path.rsplit("/", 1)[0] if "/" in path else "" for path in paths}


ops = {
    "==": operator.eq,
    "=": operator.eq,
    "!=": operator.ne,
    ">": operator.gt,
    ">=": operator.ge,
    "<": operator.lt,
    "<=": operator.le
}


def norm_col_name(name, is_index:bool=None):
    if isinstance(name, tuple):
        if is_index:
            return name[0]
        else:
            return str(name)
    return name


def get_fs(fn, open_with, mkdirs):
    fs = None
    if "FastParquetImpl.write.<locals>.<lambda>" in str(open_with):
        import inspect
        so = inspect.getclosurevars(open_with).nonlocals["storage_options"] or {}
        fs, fn = fsspec.core.url_to_fs(fn, **so)
        open_with = fs.open
        mkdirs = mkdirs or (lambda d: fs.mkdirs(d, exist_ok=True))
    return fs, fn, open_with, mkdirs

