import re
from collections import namedtuple
from collections.abc import Mapping
from enum import Enum
from urllib.parse import unquote_plus

import allel
import dask.array as da
import numpy as np
import pandas
import xarray as xr
from fsspec.core import url_to_fs
from fsspec.mapping import FSMap

DIM_VARIANT = "variants"
DIM_ALLELE = "alleles"
DIM_SAMPLE = "samples"
DIM_PLOIDY = "ploidy"


def gff3_parse_attributes(attributes_string):
    """Parse a string of GFF3 attributes ('key=value' pairs delimited by ';')
    and return a dictionary."""

    attributes = dict()
    fields = attributes_string.split(";")
    for f in fields:
        if "=" in f:
            key, value = f.split("=")
            key = unquote_plus(key).strip()
            value = unquote_plus(value.strip())
            attributes[key] = value
        elif len(f) > 0:
            # not strictly kosher, treat as a flag
            attributes[unquote_plus(f).strip()] = True
    return attributes


gff3_cols = (
    "contig",
    "source",
    "type",
    "start",
    "end",
    "score",
    "strand",
    "phase",
    "attributes",
)


def read_gff3(buf, compression="gzip"):

    # read as dataframe
    df = pandas.read_csv(
        buf,
        sep="\t",
        comment="#",
        names=gff3_cols,
        na_values=["", "."],
        compression=compression,
    )

    # parse attributes
    df["attributes"] = df["attributes"].apply(gff3_parse_attributes)

    return df


def unpack_gff3_attributes(df, attributes):

    df = df.copy()

    if attributes == "*":
        # discover all attribute keys
        attributes = set()
        for a in df["attributes"]:
            attributes.update(a.keys())
        attributes = sorted(attributes)

    # unpack attributes into columns
    for key in attributes:
        df[key] = df["attributes"].apply(lambda v: v.get(key, np.nan))
    del df["attributes"]

    return df


# zarr compatibility, version 2.11.0 introduced the BaseStore class
# see also https://github.com/malariagen/malariagen-data-python/issues/129

try:
    # zarr >= 2.11.0
    from zarr.storage import KVStore

    class SafeStore(KVStore):
        def __getitem__(self, key):
            try:
                return self._mutable_mapping[key]
            except KeyError as e:
                # raise a different error to ensure zarr propagates the exception, rather than filling
                raise FileNotFoundError(e)

        def __contains__(self, key):
            return key in self._mutable_mapping


except ImportError:
    # zarr < 2.11.0

    class SafeStore(Mapping):
        def __init__(self, store):
            self.store = store

        def __getitem__(self, key):
            try:
                return self.store[key]
            except KeyError as e:
                # raise a different error to ensure zarr propagates the exception, rather than filling
                raise FileNotFoundError(e)

        def __contains__(self, key):
            return key in self.store

        def __iter__(self):
            return iter(self.store)

        def __len__(self):
            return len(self.store)


class SiteClass(Enum):
    UPSTREAM = 1
    DOWNSTREAM = 2
    UTR5 = 3
    UTR3 = 4
    CDS_FIRST = 5
    CDS_MID = 6
    CDS_LAST = 7
    INTRON_FIRST = 8
    INTRON_MID = 9
    INTRON_LAST = 10


def da_from_zarr(z, inline_array, chunks="auto"):
    """Utility function for turning a zarr array into a dask array.

    N.B., dask does have it's own from_zarr() function but we roll
    our own here to get a little more control.

    """
    if chunks == "native" or z.dtype == object:
        # N.B., dask does not support "auto" chunks for arrays with object dtype
        chunks = z.chunks
    kwargs = dict(chunks=chunks, fancy=False, lock=False, inline_array=inline_array)
    try:
        d = da.from_array(z, **kwargs)
    except TypeError:
        # only later versions of dask support inline_array argument
        del kwargs["inline_array"]
        d = da.from_array(z, **kwargs)
    return d


def dask_compress_dataset(ds, indexer, dim):
    """Temporary workaround for memory issues when attempting to
    index an xarray dataset with a Boolean array.

    See also: https://github.com/pydata/xarray/issues/5054

    Parameters
    ----------
    ds : xarray.Dataset
    indexer : str
    dim : str

    Returns
    -------
    xarray.Dataset

    """
    if isinstance(indexer, str):
        indexer = ds[indexer].data

    # sanity checks
    assert isinstance(indexer, da.Array)
    assert indexer.ndim == 1
    assert indexer.dtype == bool
    assert indexer.shape[0] == ds.dims[dim]

    coords = dict()
    for k in ds.coords:
        a = ds[k]
        v = _dask_compress_dataarray(a, indexer, dim)
        coords[k] = (a.dims, v)

    data_vars = dict()
    for k in ds.data_vars:
        a = ds[k]
        v = _dask_compress_dataarray(a, indexer, dim)
        data_vars[k] = (a.dims, v)

    attrs = ds.attrs.copy()

    return xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs)


def _dask_compress_dataarray(a, indexer, dim):
    try:
        # find the axis for the given dimension
        axis = a.dims.index(dim)

    except ValueError:
        # array doesn't have the given dimension, return as-is
        v = a.data

    else:
        # apply the indexing operation
        v = da_compress(indexer, a.data, axis)

    return v


def da_compress(indexer, data, axis):
    """Wrapper for dask.array.compress() which computes chunk sizes faster."""

    # sanity checks
    assert isinstance(data, da.Array)
    assert isinstance(indexer, da.Array)
    assert isinstance(axis, int)
    assert indexer.shape[0] == data.shape[axis]

    # useful variables
    old_chunks = data.chunks
    axis_old_chunks = old_chunks[axis]
    axis_n_chunks = len(axis_old_chunks)

    # ensure indexer and data are chunked in the same way
    indexer = indexer.rechunk((axis_old_chunks,))

    # apply the indexing operation
    v = da.compress(indexer, data, axis=axis)

    # need to compute chunks sizes in order to know dimension sizes;
    # would normally do v.compute_chunk_sizes() but that is slow for
    # multidimensional arrays, so hack something more efficient

    axis_new_chunks = tuple(
        indexer.map_blocks(
            lambda b: np.sum(b, keepdims=True),
            chunks=((1,) * axis_n_chunks,),
        ).compute()
    )
    new_chunks = tuple(
        [axis_new_chunks if i == axis else c for i, c in enumerate(old_chunks)]
    )
    v._chunks = new_chunks

    return v


def init_filesystem(url, **kwargs):
    """Initialise an fsspec filesystem from a given base URL and parameters."""

    # special case Google Cloud Storage, use anonymous access, avoids a delay
    if url.startswith("gs://") or url.startswith("gcs://"):
        kwargs["token"] = "anon"
    elif "gs://" in url:
        # chained URL
        kwargs["gs"] = dict(token="anon")
    elif "gcs://" in url:
        # chained URL
        kwargs["gcs"] = dict(token="anon")

    # process the url using fsspec
    fs, path = url_to_fs(url, **kwargs)

    # path compatibility, fsspec/gcsfs behaviour varies between version
    while path.endswith("/"):
        path = path[:-1]

    return fs, path


def init_zarr_store(fs, path):
    """Initialise a zarr store (mapping) from an fsspec filesystem."""

    return SafeStore(FSMap(fs=fs, root=path, check=False, create=False))


Region = namedtuple("Region", ["contig", "start", "end"])


def _handle_region_coords(resource, region):

    region_pattern_match = re.search(r"([a-zA-Z0-9]+)\:(.+)\-(.+)", region)
    if region_pattern_match:
        # parse region string that contains genomic coordinates
        region_split = region_pattern_match.groups()

        contig = region_split[0]
        start = int(region_split[1].replace(",", ""))
        end = int(region_split[2].replace(",", ""))

        if contig not in resource.contigs:
            raise ValueError(f"Contig {contig} does not exist in the dataset.")
        elif (
            start < 0
            or end <= start
            or end > resource.genome_sequence(region=contig).shape[0]
        ):
            raise ValueError("Provided genomic coordinates are not valid.")

        return Region(contig, start, end)

    else:
        return None


def _handle_region_feature(resource, region):
    gene_annotation = resource.geneset(attributes=["ID"])
    results = gene_annotation.query(f"ID == '{region}'")
    if not results.empty:
        # region is a feature ID
        feature = results.squeeze()
        return Region(feature.contig, feature.start, feature.end)
    else:
        return None


def resolve_region(resource, region):
    """Parse the provided region and return `Region(contig, start, end)`.
    Supports contig names, gene names and genomic coordinates"""

    if isinstance(region, Region):
        # region is already Region tuple, nothing to do
        return region

    if isinstance(region, (list, tuple)):
        # multiple regions, normalise to list and resolve components
        return [resolve_region(resource, r) for r in region]

    # check type, fail early if bad
    if not isinstance(region, str):
        raise TypeError("The region parameter must be a string or Region object.")

    # check if region is a chromosome arm
    if region in resource.contigs:
        return Region(region, None, None)

    # check if region is a region string providing coordinates
    region_from_coords = _handle_region_coords(resource, region)
    if region_from_coords is not None:
        return region_from_coords

    # check if region is a gene annotation feature ID
    region_from_feature = _handle_region_feature(resource, region)
    if region_from_feature is not None:
        return region_from_feature

    raise ValueError(
        f"Region {region!r} is not a valid contig, region string or feature ID."
    )


def locate_region(region, pos):
    """Get array slice and a parsed genomic region.

    Parameters
    ----------
    region : Region
        Region to locate.
    pos : array-like
        Positions to be searched.

    Returns
    -------
    loc_region : slice

    """
    pos = allel.SortedIndex(pos)
    loc_region = pos.locate_range(region.start, region.end)
    return loc_region


def xarray_concat(datasets, **kwargs):
    if len(datasets) == 0:
        return datasets[0]
    else:
        return xr.concat(datasets, **kwargs)
