from __future__ import annotations

import io
import uuid

from fsspec.core import OpenFile, get_fs_token_paths, open_files
from fsspec.utils import read_block
from fsspec.utils import tokenize as fs_tokenize

from dask.highlevelgraph import HighLevelGraph

MAGIC = b"Obj\x01"
SYNC_SIZE = 16


def read_long(fo):
    """variable-length, zig-zag encoding."""
    c = fo.read(1)
    b = ord(c)
    n = b & 0x7F
    shift = 7
    while (b & 0x80) != 0:
        b = ord(fo.read(1))
        n |= (b & 0x7F) << shift
        shift += 7
    return (n >> 1) ^ -(n & 1)


def read_bytes(fo):
    """a long followed by that many bytes of data."""
    size = read_long(fo)
    return fo.read(size)


def read_header(fo):
    """Extract an avro file's header

    fo: file-like
        This should be in bytes mode, e.g., io.BytesIO

    Returns dict representing the header

    Parameters
    ----------
    fo: file-like
    """
    assert fo.read(len(MAGIC)) == MAGIC, "Magic avro bytes missing"
    meta = {}
    out = {"meta": meta}
    while True:
        n_keys = read_long(fo)
        if n_keys == 0:
            break
        for _ in range(n_keys):
            # ignore dtype mapping for bag version
            read_bytes(fo)  # schema keys
            read_bytes(fo)  # schema values
    out["sync"] = fo.read(SYNC_SIZE)
    out["header_size"] = fo.tell()
    fo.seek(0)
    out["head_bytes"] = fo.read(out["header_size"])
    return out


def open_head(fs, path, compression):
    """Open a file just to read its head and size"""
    with OpenFile(fs, path, compression=compression) as f:
        head = read_header(f)
    size = fs.info(path)["size"]
    return head, size


def read_avro(urlpath, blocksize=100000000, storage_options=None, compression=None):
    """Read set of avro files

    Use this with arbitrary nested avro schemas. Please refer to the
    fastavro documentation for its capabilities:
    https://github.com/fastavro/fastavro

    Parameters
    ----------
    urlpath: string or list
        Absolute or relative filepath, URL (may include protocols like
        ``s3://``), or globstring pointing to data.
    blocksize: int or None
        Size of chunks in bytes. If None, there will be no chunking and each
        file will become one partition.
    storage_options: dict or None
        passed to backend file-system
    compression: str or None
        Compression format of the targe(s), like 'gzip'. Should only be used
        with blocksize=None.
    """
    from dask import compute, delayed
    from dask.bag import from_delayed
    from dask.utils import import_required

    import_required(
        "fastavro", "fastavro is a required dependency for using bag.read_avro()."
    )

    storage_options = storage_options or {}
    if blocksize is not None:
        fs, fs_token, paths = get_fs_token_paths(
            urlpath, mode="rb", storage_options=storage_options
        )
        dhead = delayed(open_head)
        out = compute(*[dhead(fs, path, compression) for path in paths])
        heads, sizes = zip(*out)
        dread = delayed(read_chunk)

        offsets = []
        lengths = []
        for size in sizes:
            off = list(range(0, size, blocksize))
            length = [blocksize] * len(off)
            offsets.append(off)
            lengths.append(length)

        out = []
        for path, offset, length, head in zip(paths, offsets, lengths, heads):
            delimiter = head["sync"]
            f = OpenFile(fs, path, compression=compression)
            token = fs_tokenize(
                fs_token, delimiter, path, fs.ukey(path), compression, offset
            )
            keys = [f"read-avro-{o}-{token}" for o in offset]
            values = [
                dread(f, o, l, head, dask_key_name=key)
                for o, key, l in zip(offset, keys, length)
            ]
            out.extend(values)

        return from_delayed(out)
    else:
        files = open_files(urlpath, compression=compression, **storage_options)
        dread = delayed(read_file)
        chunks = [dread(fo) for fo in files]
        return from_delayed(chunks)


def read_chunk(fobj, off, l, head):
    """Get rows from raw bytes block"""
    import fastavro

    if hasattr(fastavro, "iter_avro"):
        reader = fastavro.iter_avro
    else:
        reader = fastavro.reader

    with fobj as f:
        chunk = read_block(f, off, l, head["sync"])
    head_bytes = head["head_bytes"]
    if not chunk.startswith(MAGIC):
        chunk = head_bytes + chunk
    i = io.BytesIO(chunk)
    return list(reader(i))


def read_file(fo):
    """Get rows from file-like"""
    import fastavro

    if hasattr(fastavro, "iter_avro"):
        reader = fastavro.iter_avro
    else:
        reader = fastavro.reader

    with fo as f:
        return list(reader(f))


def to_avro(
    b,
    filename,
    schema,
    name_function=None,
    storage_options=None,
    codec="null",
    sync_interval=16000,
    metadata=None,
    compute=True,
    **kwargs,
):
    """Write bag to set of avro files

    The schema is a complex dictionary describing the data, see
    https://avro.apache.org/docs/1.8.2/gettingstartedpython.html#Defining+a+schema
    and https://fastavro.readthedocs.io/en/latest/writer.html .
    Its structure is as follows::

        {'name': 'Test',
         'namespace': 'Test',
         'doc': 'Descriptive text',
         'type': 'record',
         'fields': [
            {'name': 'a', 'type': 'int'},
         ]}

    where the "name" field is required, but "namespace" and "doc" are optional
    descriptors; "type" must always be "record". The list of fields should
    have an entry for every key of the input records, and the types are
    like the primitive, complex or logical types of the Avro spec
    ( https://avro.apache.org/docs/1.8.2/spec.html ).

    Results in one avro file per input partition.

    Parameters
    ----------
    b: dask.bag.Bag
    filename: list of str or str
        Filenames to write to. If a list, number must match the number of
        partitions. If a string, must include a glob character "*", which will
        be expanded using name_function
    schema: dict
        Avro schema dictionary, see above
    name_function: None or callable
        Expands integers into strings, see
        ``dask.bytes.utils.build_name_function``
    storage_options: None or dict
        Extra key/value options to pass to the backend file-system
    codec: 'null', 'deflate', or 'snappy'
        Compression algorithm
    sync_interval: int
        Number of records to include in each block within a file
    metadata: None or dict
        Included in the file header
    compute: bool
        If True, files are written immediately, and function blocks. If False,
        returns delayed objects, which can be computed by the user where
        convenient.
    kwargs: passed to compute(), if compute=True

    Examples
    --------
    >>> import dask.bag as db
    >>> b = db.from_sequence([{'name': 'Alice', 'value': 100},
    ...                       {'name': 'Bob', 'value': 200}])
    >>> schema = {'name': 'People', 'doc': "Set of people's scores",
    ...           'type': 'record',
    ...           'fields': [
    ...               {'name': 'name', 'type': 'string'},
    ...               {'name': 'value', 'type': 'int'}]}
    >>> b.to_avro('my-data.*.avro', schema)  # doctest: +SKIP
    ['my-data.0.avro', 'my-data.1.avro']
    """
    # TODO infer schema from first partition of data
    from dask.utils import import_required

    import_required(
        "fastavro", "fastavro is a required dependency for using bag.to_avro()."
    )
    _verify_schema(schema)

    storage_options = storage_options or {}
    files = open_files(
        filename,
        "wb",
        name_function=name_function,
        num=b.npartitions,
        **storage_options,
    )
    name = f"to-avro-{uuid.uuid4().hex}"
    dsk = {
        (name, i): (
            _write_avro_part,
            (b.name, i),
            f,
            schema,
            codec,
            sync_interval,
            metadata,
        )
        for i, f in enumerate(files)
    }
    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[b])
    out = type(b)(graph, name, b.npartitions)
    if compute:
        out.compute(**kwargs)
        return [f.path for f in files]
    else:
        return out.to_delayed()


def _verify_schema(s):
    assert isinstance(s, dict), "Schema must be dictionary"
    for field in ["name", "type", "fields"]:
        assert field in s, f"Schema missing '{field}' field"
    assert s["type"] == "record", "Schema must be of type 'record'"
    assert isinstance(s["fields"], list), "Fields entry must be a list"
    for f in s["fields"]:
        assert "name" in f and "type" in f, f"Field spec incomplete: {f}"


def _write_avro_part(part, f, schema, codec, sync_interval, metadata):
    """Create single avro file from list of dictionaries"""
    import fastavro

    with f as f:
        fastavro.writer(f, schema, part, codec, sync_interval, metadata)
