# (C) Copyright 2020 ECMWF.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.
#

import datetime
import fnmatch
import logging
import os

from tqdm import tqdm

LOG = logging.getLogger(__name__)


def post_process_valid_date(field, h):
    date = h.get("validityDate")
    time = h.get("validityTime")
    field["datetime"] = datetime.datetime(
        date // 10000,
        date % 10000 // 100,
        date % 100,
        time // 100,
        time % 100,
    )
    # Note that we do not create a script here:
    # There is no ".isoformat()". Because the sql database
    # takes care of this conversion back and forth.
    return field


def post_process_parameter_level(field, h):
    param = field.get("param", None)
    if param is None:
        field["param_level"] = None
        return field

    level = field.get("levelist", None)
    if level is None:
        field["param_level"] = param
        return field

    field["param_level"] = f"{param}_{level}"
    return field


def post_process_statistics(field, h):
    values = h.get("values")
    field["mean"] = values.mean()
    field["std"] = values.std()
    field["min"] = values.min()
    field["max"] = values.max()
    field["shape"] = ",".join([str(_) for _ in values.shape])
    return field


def _index_grib_file(
    path,
    with_statistics=False,
    with_valid_date=True,
    with_parameter_level=True,
):
    import eccodes

    from earthkit.data.readers.grib.codes import GribCodesHandle

    post_process_mars = []
    if with_valid_date:
        post_process_mars.append(post_process_valid_date)
    if with_parameter_level:
        post_process_mars.append(post_process_parameter_level)
    if with_statistics:
        post_process_mars.append(post_process_statistics)

    def parse_field(h):
        field = h.as_namespace("mars")

        if post_process_mars:
            for f in post_process_mars:
                field = f(field, h)

        field["_path"] = path
        field["_offset"] = h.get_long("offset")
        field["_length"] = h.get_long("totalLength")
        field["_param_id"] = h.get_string("paramId")
        field["md5_grid_section"] = h.get("md5GridSection")

        # eccodes.codes_get_string(h, "number") returns "0"
        # when "number" is not in the iterator
        # remove? field["number"] = h.get("number")

        return field

    size = os.path.getsize(path)
    pbar = tqdm(
        desc=f"Parsing {path}",
        total=size,
        unit_scale=True,
        unit_divisor=1024,
        unit="B",
        leave=False,
        # position=TQDM_POSITION,
        dynamic_ncols=True,
    )

    with open(path, "rb") as f:
        old_position = f.tell()
        h = eccodes.codes_grib_new_from_file(f)

        while h:
            yield parse_field(GribCodesHandle(h, path, offset=old_position))

            position = f.tell()
            pbar.update(position - old_position)
            old_position = position
            h = eccodes.codes_grib_new_from_file(f)

    pbar.close()


def _index_url(url):
    import earthkit.data

    path = earthkit.data.from_source("url", url).path
    # TODO: should use download_and_cache
    # path = download_and_cache(url)
    for entry in _index_grib_file(path):
        del entry["_path"]
        yield entry


def _index_path(path):
    if os.path.isdir(path):
        for root, _, files in os.walk(path):
            for name in files:
                yield from _index_grib_file(os.path.join(root, name))
    else:
        yield from _index_grib_file(path)


class PathParserIterator:
    """Delays parsing the directory for the list of files
    until the iterator is actually used (calling __iter__)
    """

    def __init__(
        self,
        path,
        relative_paths,
        ignore=None,
        followlinks=True,
        verbose=False,
        with_statistics=True,
    ):
        if ignore is None:
            ignore = []
        self.ignore = ignore
        self.path = path
        self.relative_paths = relative_paths
        self.followlinks = followlinks
        self.verbose = verbose
        self.with_statistics = with_statistics

        self._tasks = None

    def __iter__(self):
        for path in tqdm(self.tasks, dynamic_ncols=True):
            for entry in self.process_one_task(path):
                yield entry

    @property
    def tasks(self):
        if self._tasks is not None:
            return self._tasks

        LOG.debug(f"Parsing files in path={self.path}")
        assert os.path.exists(self.path), f"{self.path} does not exist"
        # assert os.path.isdir(self.path), f"{self.path} is not a directory"

        def _ignore(path):
            for ignore in self.ignore:
                if fnmatch.fnmatch(os.path.basename(path), ignore):
                    return True
            return False

        tasks = []

        if not os.path.isdir(self.path):
            tasks.append(self.path)
        else:
            for root, _, files in os.walk(self.path, followlinks=self.followlinks):
                for name in files:
                    path = os.path.join(root, name)
                    if path in self.ignore:
                        continue
                    tasks.append(path)

        if tasks:
            if self.verbose:
                print(f"Found {len(tasks)} files to index.")
        else:
            LOG.error(f"Could not find any files to index in path={self.path}")

        self._tasks = tasks

        return self.tasks


class GribIndexingPathParserIterator(PathParserIterator):
    def process_one_task(self, path):
        LOG.debug(f"Parsing file {path}")

        if self.relative_paths is True:
            _path = os.path.relpath(path, self.path)
        elif self.relative_paths is False:
            _path = os.path.abspath(path)
        elif self.relative_paths is None:
            _path = path
        else:
            assert False, self.relative_paths

        try:
            # We could use reader(self, path) but this will create a json
            # grib-index auxiliary file in the cache.
            # Indexing 1M grib files lead to 1M in cache.
            #
            # We would need either to refactor the grib reader.

            from earthkit.data.readers.grib.parsing import _index_grib_file

            for field in _index_grib_file(path, with_statistics=self.with_statistics):
                field["_path"] = _path
                yield field
        except PermissionError as e:
            LOG.error(f"Could not read {path}: {e}")
            return
        except Exception as e:
            print(f"(grib-parsing) Ignoring {path}, {e}")
            LOG.exception(f"(grib-parsing) Ignoring {path}, {e}")
            return
