import itertools
import json
import base64
import pathlib
import uuid
from collections import defaultdict

import cfgrib
import eccodes
import numpy as np
import xarray as xr

from .magician import Magician
from . import gridutils as gu

import logging

logger = logging.getLogger("gribscan")


def find_stream(f, needle, buffersize=1024 * 1024):
    keep_going = True
    while keep_going:
        start = f.tell()
        buf = f.read(buffersize)
        if len(buf) < buffersize:
            keep_going = False
        try:
            idx = buf.index(needle)
        except ValueError:
            f.seek(-len(needle), 1)
            continue
        else:
            pos = start + idx
            f.seek(pos)
            return pos


def detect_large_grib1_special_coding(f, part_size):
    """
    This is from eccodes src/grib_io.c /* Special coding */ (couldn't find it in the specs...)
    """
    if part_size & 0x800000:  # this is a large grib, hacks are coming...
        start = f.tell()
        data = f.read(part_size)
        f.seek(start)
        assert data[7] == 1, "large grib mode only exists in Grib 1"

        s0len = 8
        s1start = s0len
        s1len = int.from_bytes(data[s1start : s1start + 3], "big")
        flags = data[s1start + 7]
        has_s2 = bool(flags & (1 << 7))
        has_s3 = bool(flags & (1 << 6))

        s2start = s1start + s1len
        if has_s2:
            s2len = int.from_bytes(data[s2start : s2start + 3], "big")
        else:
            s2len = 0

        s3start = s2start + s2len
        if has_s3:
            s3len = int.from_bytes(data[s3start : s3start + 3], "big")
        else:
            s3len = 0

        s4start = s3start + s3len

        s4len = int.from_bytes(data[s4start : s4start + 3], "big")
        if s4len < 120:
            return (part_size & 0x7FFFFF) * 120 - s4len + 4
        else:
            return part_size

    else:  # normal grib
        return part_size


def _split_file(f, skip=0):
    """
    splits a gribfile into individual messages
    """
    if hasattr(f, "size"):
        size = f.size
    else:
        f.seek(0, 2)
        size = f.tell()
        f.seek(0)
    part = 0

    logger.debug(f"reading GRIB file with size {size}")

    while f.tell() < size:
        logger.debug(f"extract part {part + 1}")
        start = f.tell()
        indicator = f.peek(16)
        if indicator[:4] != b"GRIB":
            logger.info(f"non-consecutive messages, searching for part {part + 1}")
            if (start := find_stream(f, b"GRIB")) is None:
                # Handle non-GRIB data after the last GRIB message in a file
                return
            indicator = f.peek(16)
        if len(indicator) < 16:
            indicator = f.read(16)
            f.seek(start)
        if len(indicator) < 16:
            logger.info(f"couldn't peek or seek indicator, assuming end of file at {start}")
            return

        grib_edition = indicator[7]

        if grib_edition == 1:
            part_size = int.from_bytes(indicator[4:7], "big")
            part_size = detect_large_grib1_special_coding(f, part_size)
        elif grib_edition == 2:
            part_size = int.from_bytes(indicator[8:16], "big")
        else:
            raise ValueError(f"unknown grib edition: {grib_edition}")

        data = f.read(part_size)
        if data[-4:] != b"7777":
            logger.warning(f"part {part + 1} is broken")
            f.seek(start + 1)
        else:
            yield start, part_size, grib_edition, data

        part += 1
        if skip and part > skip:
            break


EXTRA_PARAMETERS = [
    "forecastTime",
    "indicatorOfUnitOfTimeRange",
    "lengthOfTimeRange",
    "indicatorOfUnitForTimeRange",
    "productDefinitionTemplateNumber",
    "N",
    "timeRangeIndicator",
    "P1",
    "P2",
    "numberIncludedInAverage",
]

production_template_numbers = {
    0: {"forcastTime": True, "timeRange": False},
    1: {"forcastTime": True, "timeRange": False},
    2: {"forcastTime": True, "timeRange": False},
    3: {"forcastTime": True, "timeRange": False},
    4: {"forcastTime": True, "timeRange": False},
    5: {"forcastTime": True, "timeRange": False},
    6: {"forcastTime": True, "timeRange": False},
    7: {"forcastTime": True, "timeRange": False},
    15: {"forcastTime": True, "timeRange": False},
    32: {"forcastTime": True, "timeRange": False},
    33: {"forcastTime": True, "timeRange": False},
    40: {"forcastTime": True, "timeRange": False},
    41: {"forcastTime": True, "timeRange": False},
    44: {"forcastTime": True, "timeRange": False},
    45: {"forcastTime": True, "timeRange": False},
    48: {"forcastTime": True, "timeRange": False},
    51: {"forcastTime": True, "timeRange": False},
    53: {"forcastTime": True, "timeRange": False},
    54: {"forcastTime": True, "timeRange": False},
    55: {"forcastTime": True, "timeRange": False},
    56: {"forcastTime": True, "timeRange": False},
    57: {"forcastTime": True, "timeRange": False},
    58: {"forcastTime": True, "timeRange": False},
    60: {"forcastTime": True, "timeRange": False},
    1000: {"forcastTime": True, "timeRange": False},
    1002: {"forcastTime": True, "timeRange": False},
    1100: {"forcastTime": True, "timeRange": False},
    40033: {"forcastTime": True, "timeRange": False},
    40455: {"forcastTime": True, "timeRange": False},
    40456: {"forcastTime": True, "timeRange": False},
    20: {"forcastTime": False, "timeRange": False},
    30: {"forcastTime": False, "timeRange": False},
    31: {"forcastTime": False, "timeRange": False},
    254: {"forcastTime": False, "timeRange": False},
    311: {"forcastTime": False, "timeRange": False},
    2000: {"forcastTime": False, "timeRange": False},
    8: {"forcastTime": True, "timeRange": True},
    9: {"forcastTime": True, "timeRange": True},
    10: {"forcastTime": True, "timeRange": True},
    11: {"forcastTime": True, "timeRange": True},
    12: {"forcastTime": True, "timeRange": True},
    13: {"forcastTime": True, "timeRange": True},
    14: {"forcastTime": True, "timeRange": True},
    34: {"forcastTime": True, "timeRange": True},
    42: {"forcastTime": True, "timeRange": True},
    43: {"forcastTime": True, "timeRange": True},
    46: {"forcastTime": True, "timeRange": True},
    47: {"forcastTime": True, "timeRange": True},
    61: {"forcastTime": True, "timeRange": True},
    67: {"forcastTime": True, "timeRange": True},
    68: {"forcastTime": True, "timeRange": True},
    91: {"forcastTime": True, "timeRange": True},
    1001: {"forcastTime": True, "timeRange": True},
    1101: {"forcastTime": True, "timeRange": True},
    10034: {"forcastTime": True, "timeRange": True},
}

# according to http://www.cosmo-model.org/content/consortium/generalMeetings/general2014/wg6-pompa/grib2/grib/pdtemplate_4.41.htm
time_range_units = {
    0: 60,  # np.timedelta64(1, "m"),
    1: 60 * 60,  # np.timedelta64(1, "h"),
    2: 24 * 60 * 60,  # np.timedelta64(1, "D"),
    # 3   Month
    # 4   Year
    # 5   Decade (10 years)
    # 6   Normal (30 years)
    # 7   Century (100 years)
    # 8-9 Reserved
    10: 3 * 60 * 60,  # np.timedelta64(3, "h"),
    11: 6 * 60 * 60,  # np.timedelta64(6, "h"),
    12: 12 * 60 * 60,  # np.timedelta64(12, "h"),
    13: 1,  # np.timedelta64(1, "s"),
    # 14-191  Reserved
    # 192-254 Reserved for local use
    # 255 Missing
}


def get_time_offset(gribmessage, lean_towards="end"):
    """Calculate time offset based on GRIB definition.

    See: https://codes.ecmwf.int/grib/format/grib1/ctable/5/
    """
    offset = 0  # np.timedelta64(0, "s")
    edition = int(gribmessage["editionNumber"])
    if edition == 1:
        timeRangeIndicator = int(gribmessage["timeRangeIndicator"])
        if timeRangeIndicator == 0:
            unit = time_range_units[
                int(gribmessage.get("indicatorOfUnitOfTimeRange", 255))
            ]
            offset += int(gribmessage["P1"]) * unit
        elif timeRangeIndicator == 1:
            pass
        elif timeRangeIndicator in [2, 3]:
            unit = time_range_units[
                int(gribmessage.get("indicatorOfUnitOfTimeRange", 255))
            ]
            if lean_towards == "start":
                offset += int(gribmessage["P1"]) * unit
            elif lean_towards == "end":
                offset += int(gribmessage["P2"]) * unit
        elif timeRangeIndicator == 4:
            unit = time_range_units[
                int(gribmessage.get("indicatorOfUnitOfTimeRange", 255))
            ]
            offset += int(gribmessage["P2"]) * unit
        elif timeRangeIndicator == 10:
            unit = time_range_units[
                int(gribmessage.get("indicatorOfUnitOfTimeRange", 255))
            ]
            offset += (int(gribmessage["P1"]) * 256 + int(gribmessage["P2"])) * unit
        elif timeRangeIndicator == 123:
            unit = time_range_units[
                int(gribmessage.get("indicatorOfUnitOfTimeRange", 255))
            ]
            if lean_towards == "end":
                N = int(gribmessage["numberIncludedInAverage"])
                offset += N * int(gribmessage["P2"]) * unit
        else:
            raise NotImplementedError(
                f"don't know how to handle timeRangeIndicator {timeRangeIndicator}"
            )
    else:
        try:
            options = production_template_numbers[
                int(gribmessage["productDefinitionTemplateNumber"])
            ]
        except KeyError:
            return offset
        if options["forcastTime"]:
            unit = time_range_units[
                int(gribmessage.get("indicatorOfUnitOfTimeRange", 255))
            ]
            offset += gribmessage.get("forecastTime", 0) * unit
        if options["timeRange"]:
            unit = time_range_units[
                int(gribmessage.get("indicatorOfUnitOfTimeRange", 255))
            ]
            length = gribmessage.get("lengthOfTimeRange", 0)
            if isinstance(length, np.ndarray):
                if lean_towards == "start":
                    length = int(length[0])
                elif lean_towards == "end":
                    length = int(length[-1])
            offset += length * unit
    return offset


def arrays_to_list(o):
    try:
        return o.tolist()
    except AttributeError:
        return o


def scan_gribfile(filelike, **kwargs):
    for offset, size, grib_edition, data in _split_file(filelike):
        mid = eccodes.codes_new_from_message(data)
        m = cfgrib.cfmessage.CfMessage(mid)
        t = eccodes.codes_get_native_type(m.codes_id, "values")
        s = eccodes.codes_get_size(m.codes_id, "values")

        global_attrs = {k: m[k] for k in cfgrib.dataset.GLOBAL_ATTRIBUTES_KEYS}
        for uuid_key in ["uuidOfHGrid", "uuidOfVGrid"]:
            try:
                global_attrs[uuid_key] = str(
                    uuid.UUID(eccodes.codes_get_string(mid, uuid_key))
                )
            except eccodes.KeyValueNotFoundError:
                pass

        idx = {
            "globals": global_attrs,
            "attrs": {
                k: m.get(k, None)
                for k in cfgrib.dataset.DATA_ATTRIBUTES_KEYS
                + cfgrib.dataset.EXTRA_DATA_ATTRIBUTES_KEYS
            },
            "parameter_code": {
                k: m.get(k, None)
                for k in ["discipline", "parameterCategory", "parameterNumber"]
            },
            "posix_time": m["time"] + get_time_offset(m),
            "domain": m["globalDomain"],
            "member": m.get("number", None),
            "time": f"{m['hour']:02d}{m['minute']:02d}",
            "date": f"{m['year']:04d}{m['month']:02d}{m['day']:02d}",
            "levtype": m.get("typeOfLevel", None),
            "level": m.get("level", None),
            "type": m.get("dataType", None),
            "referenceTime": m["time"],
            "step": m["step"],
            "_offset": offset,
            "_length": size,
            "array": {
                "dtype": np.dtype(t).str,
                "shape": [s],
            },
            "extra": {
                k: arrays_to_list(m.get(k, None))
                for k in (EXTRA_PARAMETERS + gu.params_for_gridType(m["gridType"]))
            },
            **kwargs,
        }

        if (param := m.get("shortName", "unknown")) != "unknown":
            idx["param"] = param
        else:
            idx["param"] = ".".join(map(str, idx["parameter_code"].values()))

        yield idx


def write_index(gribfile, idxfile=None, outdir=None, force=False):
    p = pathlib.Path(gribfile)
    if outdir is None:
        outdir = p.parent

    if idxfile is None:
        # Replace GRIB suffixes but append to all others (e.g. `fcst_phy2m.202410`).
        suffix_map = {s: ".index" for s in (".grib", ".grb", ".grib2", ".grb2")}
        idxfile = pathlib.Path(outdir) / (
            p.with_suffix(suffix_map.get(p.suffix, p.suffix + ".index")).name
        )

    if idxfile.exists() and not force:
        raise FileExistsError(f"Index file {idxfile} already exists!")

    # We need to use the gribfile (str) variable because Path() objects
    # collapse the "/./" notation used to denote subtrees.
    gen = scan_gribfile(open(p, "rb"), filename=gribfile)

    tempfile = idxfile.with_suffix(".index.partial")
    with open(tempfile, "w") as output_file:
        for record in gen:
            json.dump(record, output_file)
            output_file.write("\n")

    if force or not idxfile.exists():
        tempfile.rename(idxfile)
    else:
        logger.warning(f"Index file {idxfile} got created during runtime.")


def parse_index(indexfile, m2key, duplicate="replace"):
    index = {}
    with open(indexfile, "r") as f:
        for line in f:
            meta = json.loads(line)
            tinfo = m2key(meta)
            if tinfo in index:
                if duplicate == "replace":
                    index[tinfo] = meta
                elif duplicate == "keep":
                    continue
                elif duplicate == "error":
                    raise Exception(f"Duplicate message step: {tinfo}")
            else:
                index[tinfo] = meta
    return list(index.values())


def is_value(v):
    if v is None or v == "undef" or v == "unknown":
        return False
    else:
        return True


def inspect_grib_indices(messages, magician):
    coords_by_key = defaultdict(lambda: tuple(set() for _ in magician.dimkeys))
    size_by_key = defaultdict(set)
    attrs_by_key = {}
    extra_by_key = {}
    dtype_by_key = {}
    global_attrs = {}

    for msg in messages:
        varkey, coords = magician.m2key(msg)
        for existing, new in zip(coords_by_key[varkey], coords):
            existing.add(new)
        size_by_key[varkey].add(msg["array"]["shape"][0])
        attrs_by_key[varkey] = {k: v for k, v in msg["attrs"].items() if is_value(v)}
        extra_by_key[varkey] = {k: v for k, v in msg["extra"].items() if is_value(v)}
        dtype_by_key[varkey] = msg["array"]["dtype"]
        global_attrs = msg["globals"]

    for k, v in size_by_key.items():
        assert len(v) == 1, f"inconsistent shape of {k}"

    size_by_key = {k: list(v)[0] for k, v in size_by_key.items()}

    varinfo = {}
    for varkey, coords in coords_by_key.items():
        if all(len(c) == 1 for c in coords):
            dims = ()
            dim_id = ()
            shape = ()
        else:
            dims, dim_id, shape = map(
                tuple,
                zip(
                    *(
                        (dim, i, len(coords))
                        for i, (dim, coords) in enumerate(zip(magician.dimkeys, coords))
                        if len(coords) != 1
                    )
                ),
            )

        info = {
            "dims": dims,
            "shape": shape,
            "dim_id": dim_id,
            "coords": tuple(coords_by_key[varkey][i] for i in dim_id),
            "dtype": dtype_by_key[varkey],
            "attrs": attrs_by_key[varkey],
            "extra": extra_by_key[varkey],
        }

        computed_coords = gu.varinfo2coords(info)

        info = {
            **info,
            "data_dims": list(computed_coords.dims),
            "data_shape": [computed_coords.sizes[d] for d in computed_coords.dims],
        }

        varinfo[varkey] = {
            **info,
            **magician.variable_hook(varkey, info),
        }

    coords = defaultdict(set)
    for _, info in varinfo.items():
        for dim, cs in zip(info["dims"], info["coords"]):
            coords[dim] |= cs

    coords = {
        **{
            k: xr.DataArray(
                np.asarray(list(sorted(c))),
                dims=(k,),
                attrs=gu.default_attrs.get(k, {}),
            )
            for k, c in coords.items()
        },
        **computed_coords.variables,
        **magician.extra_coords(varinfo),
    }

    return global_attrs, coords, varinfo


def build_refs(messages, global_attrs, coords, varinfo, magician):
    coords_inv = {
        k: {v: i for i, v in enumerate(vs.values)} for k, vs in coords.items()
    }

    refs = {}
    for msg in messages:
        key, coord = magician.m2key(msg)
        info = varinfo[key]
        cs = [coord[d] for d in info["dim_id"]]
        chunk_id = ".".join(
            itertools.chain(
                map(str, [coords_inv[d][c] for d, c in zip(info["dims"], cs)]),
                ["0"] * len(info["data_dims"])
            )
        )
        refs[info["name"] + "/" + chunk_id] = [
            msg["filename"],
            msg["_offset"],
            msg["_length"],
        ]

    for varkey, info in varinfo.items():
        refs[info["name"] + "/.zattrs"] = json.dumps(
            {
                **info["attrs"],
                "_ARRAY_DIMENSIONS": list(info["dims"]) + list(info["data_dims"]),
            }
        )
        shape = [len(coords[dim]) for dim in info["dims"]] + list(info["data_shape"])
        chunks = [1 for _ in info["shape"]] + list(info["data_shape"])
        refs[info["name"] + "/.zarray"] = json.dumps(
            {
                "shape": shape,
                "chunks": chunks,
                "compressor": {"id": "gribscan.rawgrib"},
                "dtype": info["dtype"],
                "fill_value": info["attrs"].get("missingValue", 9999),
                "filters": None,
                "order": "C",
                "zarr_format": 2,
            }
        )

    for name, cs in coords.items():
        cs, array_meta, compressor = magician.coords_hook(name, cs)

        if compressor is None:
            compressor_id = None
            data = bytes(cs.values)
        else:
            compressor_id = compressor.get_config()
            data = bytes(compressor.encode(cs.values))

        refs[f"{name}/.zattrs"] = json.dumps({**cs.attrs, "_ARRAY_DIMENSIONS": cs.dims})
        refs[f"{name}/.zarray"] = json.dumps(
            {
                **{
                    "chunks": list(cs.shape),
                    "compressor": compressor_id,
                    "dtype": cs.dtype.str,
                    "fill_value": None,
                    "filters": None,
                    "order": "C",
                    "shape": list(cs.shape),
                    "zarr_format": 2,
                },
                **array_meta,
            }
        )
        refs[f"{name}/0"] = "base64:" + base64.b64encode(data).decode("ascii")

    refs[".zgroup"] = json.dumps({"zarr_format": 2})
    refs[".zattrs"] = json.dumps(magician.globals_hook(global_attrs))

    return refs


def is_zarr_key(key):
    return key.endswith((".zarray", ".zgroup", ".zattrs"))


def consolidate_metadata(refs):
    return json.dumps(
        {
            "zarr_consolidated_format": 1,
            "metadata": {
                key: json.loads(value)
                for key, value in refs.items()
                if is_zarr_key(key)
            },
        }
    )


def subtree(path, sep="/./"):
    """Return sub-tree of a given path.

    The start of a sub-tree is marked by a user-defined string (default '/./').

    Example:

    >>> subtree("/foo/bar/./baz/")
    "baz/"

    Notes:
        This funcion mimicks the behaviour of rsync in -R/--relative mode.
        https://askubuntu.com/a/552122
    """
    return path.split(sep)[-1]


def prepend_path(refs, prefix):
    """Prepend a path-prefix to all target filenames in a given reference filesystem.

    For absolute target paths, the existing target parents are overwritten.
    """
    return {
        k: [(pathlib.Path(prefix) / subtree(target[0])).as_posix()] + target[1:]
        if isinstance(target, list) else target
        for k, target in refs.items()
    }


def compress_extra_attributes(messages):
    it = iter(messages)
    mlast = next(it)
    yield mlast
    for m in it:
        if "extra" in m:
            lastextra = mlast["extra"]
            # if extra attribute in this message is large (i.e. a list or dict) and is the same as in previous message, replace it by a reference to the previous one
            extra = {
                k: lastextra[k]
                if k in lastextra and isinstance(v, (list, dict)) and lastextra[k] == v
                else v
                for k, v in m["extra"].items()
            }
            m = {**m, "extra": extra}
        yield m
        mlast = m


def grib_magic(filenames, magician=None, global_prefix=None):
    if magician is None:
        magician = Magician()

    messages = list(
        compress_extra_attributes(
            msg
            for filename in filenames
            for msg in parse_index(filename, magician.m2key)
        )
    )

    messages_by_dataset = defaultdict(list)
    for message in messages:
        messages_by_dataset[magician.m2dataset(message)].append(message)

    refs_by_dataset = {}
    for dataset, messages in messages_by_dataset.items():
        global_attrs, coords, varinfo = inspect_grib_indices(messages, magician)
        refs = build_refs(messages, global_attrs, coords, varinfo, magician)
        refs[".zmetadata"] = consolidate_metadata(refs)
        if global_prefix is None:
            refs_by_dataset[dataset] = refs
        else:
            refs_by_dataset[dataset] = prepend_path(refs, global_prefix)

    return refs_by_dataset
