#!/usr/bin/env python
# (C) Copyright 2021 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 itertools
import json
import logging
import os
from collections import defaultdict

import requests
from multiurl import download, robust

from .date import (
    canonical_time,
    end_step,
    expand_date,
    expand_list,
    expand_time,
    full_date,
)
from .urls import URLS

LOG = logging.getLogger(__name__)

HOURLY_PATTERN = (
    "{_url}/{_yyyymmdd}/{_H}z/{resol}/{_stream}/"
    "{_yyyymmddHHMMSS}-{step}h-{_stream}-{type}.{_extension}"
)

MONTHLY_PATTERN = (
    "{_url}/{_yyyymmdd}/{_H}z/{resol}/{_stream}/"
    "{_yyyymmddHHMMSS}-{fcmonth}m-{_stream}-{type}.{_extension}"
)

PATTERNS = {"mmsa": MONTHLY_PATTERN}
EXTENSIONS = {"tf": "bufr"}

ONCE = set()


def warning_once(*args, did_you_mean=None):

    if repr(args) in ONCE:
        return

    LOG.warning(*args)

    ONCE.add(repr(args))

    if did_you_mean:
        (words, vocabulary) = did_you_mean

        def levenshtein(a, b):
            if len(a) == 0:
                return len(b)

            if len(b) == 0:
                return len(a)

            if a[0].lower() == b[0].lower():
                return levenshtein(a[1:], b[1:])

            return 1 + min(
                [
                    levenshtein(a[1:], b[1:]),
                    levenshtein(a[1:], b),
                    levenshtein(a, b[1:]),
                ]
            )

        if not isinstance(words, (list, tuple)):
            words = [words]

        for word in words:
            distance, best = min((levenshtein(word, w), w) for w in vocabulary)
            if distance < min(len(word), len(best)):
                LOG.warning(
                    "Did you mean %r instead of %r?",
                    best,
                    word,
                )


class Result:
    def __init__(self, urls, target, dates, for_urls, for_index):
        self.urls = urls
        self.target = target
        if len(dates) == 1:
            self.datetime = dates[0]
        else:
            self.datetime = dates
        self.for_urls = for_urls
        self.for_index = for_index


class Client:
    def __init__(
        self,
        source="ecmwf",
        beta=True,
        preserve_request_order=False,
        infer_stream_keyword=True,
        debug=False,
    ):
        self._url = None
        self.source = source
        self.beta = beta
        self.preserve_request_order = preserve_request_order
        self.infer_stream_keyword = infer_stream_keyword

        if debug:
            logging.basicConfig(level=logging.DEBUG)

    @property
    def url(self):

        if self._url is None:

            if self.source.startswith("http://") or self.source.startswith("https://"):
                self._url = self.source
            else:
                if self.source not in URLS:
                    warning_once(
                        "Unknown source %r. Known sources are %r",
                        self.source,
                        list(URLS.keys()),
                        did_you_mean=(self.source, list(URLS.keys())),
                    )
                    raise ValueError("Unknown source %r" % (self.source,))
                self._url = URLS[self.source]

        return self._url

    def retrieve(self, request=None, target=None, **kwargs):
        result = self._get_urls(request, target=target, use_index=True, **kwargs)
        result.size = download(result.urls, target=result.target)
        return result

    def download(self, request=None, target=None, **kwargs):
        result = self._get_urls(request, target=target, use_index=False, **kwargs)
        result.size = download(result.urls, target=result.target)
        return result

    def latest(self, request=None, **kwargs):
        if request is None:
            params = dict(**kwargs)
        else:
            params = dict(**request)

        if "time" not in params:
            delta = datetime.timedelta(hours=6)
        else:
            delta = datetime.timedelta(days=1)

        date = full_date(0, params.get("time", 18))

        stop = date - datetime.timedelta(days=1, hours=6)

        while date > stop:
            result = self._get_urls(
                request=None,
                use_index=False,
                date=date,
                **params,
            )
            codes = [robust(requests.head)(url).status_code for url in result.urls]
            if len(codes) > 0 and all(c == 200 for c in codes):
                return date
            date -= delta

        raise ValueError("Cannot etablish latest date for %r" % (result.for_urls,))

    def _get_urls(self, request=None, use_index=None, target=None, **kwargs):

        assert use_index in (True, False)
        if request is None:
            params = dict(**kwargs)
        else:
            params = dict(**request)

        if "date" not in params:
            params["date"] = self.latest(params)

        if target is None:
            target = params.pop("target", None)

        for_urls, for_index = self.prepare_request(params)

        for_urls["_url"] = [self.url]

        seen = set()
        data_urls = []

        dates = set()

        for args in (
            dict(zip(for_urls.keys(), x)) for x in itertools.product(*for_urls.values())
        ):
            pattern = PATTERNS.get(args["stream"], HOURLY_PATTERN)
            date = full_date(args.pop("date", None), args.pop("time", None))
            dates.add(date)
            args["_yyyymmdd"] = date.strftime("%Y%m%d")
            args["_H"] = date.strftime("%H")
            args["_yyyymmddHHMMSS"] = date.strftime("%Y%m%d%H%M%S")
            args["_extension"] = EXTENSIONS.get(args["type"], "grib2")
            args["_stream"] = self.patch_stream(args)

            url = pattern.format(**args)
            if url not in seen:
                data_urls.append(url)
                seen.add(url)

        if for_index and use_index:
            data_urls = self.get_parts(data_urls, for_index)

        return Result(
            urls=data_urls,
            target=target,
            dates=sorted(dates),
            for_urls=for_urls,
            for_index=for_index,
        )

    def get_parts(self, data_urls, for_index):

        count = len(for_index)
        result = []
        line = None

        possible_values = defaultdict(set)

        for url in data_urls:
            base, _ = os.path.splitext(url)
            index_url = f"{base}.index"
            r = robust(requests.get)(index_url)
            r.raise_for_status()

            parts = []
            for line in r.iter_lines():
                line = json.loads(line)
                matches = []
                for i, (name, values) in enumerate(for_index.items()):
                    idx = line.get(name)
                    if idx is not None:
                        possible_values[name].add(idx)
                    if idx in values:
                        if self.preserve_request_order:
                            for j, v in enumerate(values):
                                if v == idx:
                                    matches.append((i, j))
                        else:
                            matches.append(line["_offset"])

                if len(matches) == count:
                    parts.append((tuple(matches), (line["_offset"], line["_length"])))

            if parts:
                result.append((url, tuple(p[1] for p in sorted(parts))))

        for name, values in for_index.items():
            diff = set(values).difference(possible_values[name])
            for d in diff:
                warning_once(
                    "No index entries for %s=%s",
                    name,
                    d,
                    did_you_mean=(d, possible_values[name]),
                )

        if not result:
            raise ValueError("Cannot find index entries matching %r" % (for_index,))

        return result

    def user_to_index(self, key, value, request, for_index):
        FOR_INDEX = {
            ("type", "ef"): ["cf", "pf"],
        }

        return FOR_INDEX.get((key, value), value)

    def user_to_url(self, key, value, request, for_urls):
        FOR_URL = {
            ("type", "cf"): "ef",
            ("type", "pf"): "ef",
            ("type", "em"): "ep",
            ("type", "es"): "ep",
            ("type", "fcmean"): "fc",
            ("stream", "mmsa"): "mmsf",
        }

        if key == "step" and for_urls["type"] == ["ep"]:
            if end_step(value) <= 240:
                return "240"
            else:
                return "360"

        return FOR_URL.get((key, value), value)

    def prepare_request(self, request=None, **kwargs):
        DEFAULTS_FC = dict(
            resol="0p4-beta" if self.beta else "0p4",
            type="fc",
            stream="oper",
            step=0,
            fcmonth=1,
        )

        DEFAULTS_EF = dict(
            resol="0p4-beta" if self.beta else "0p4",
            type=["cf", "pf"],
            stream="enfo",
            step=0,
            fcmonth=1,
        )

        DEFAULTS = {
            "enfo": DEFAULTS_EF,
            "waef": DEFAULTS_EF,
        }

        URL_COMPONENTS = (
            "date",
            "time",
            "resol",
            "stream",
            "type",  # Must be before 'step' in that list
            "step",
            "fcmonth",
        )

        INDEX_COMPONENTS = (
            "param",
            "type",
            "step",
            "fcmonth",
            "number",
            "levelist",
            "levtype",
        )

        CANONICAL = {
            "time": lambda x: str(canonical_time(x)),
            # "param": lambda x: str(x).lower(),
            # "type": lambda x: str(x).lower(),
            # "stream": lambda x: str(x).lower(),
        }

        EXPAND_LIST = {
            "date": expand_date,
            "time": expand_time,
        }

        OTHER_STEP = {"mmsa": "step"}

        POSPROCESSING = {
            "area",
            "grid",
            "rotation",
            "frame",
            "bitmap",
            "gaussian",
            "accuracy",
            "format",
        }

        POSSIBLE_VALUES = {
            "type": ["tf", "fc", "fcmean", "cf", "pf", "em", "ep", "es"],
            "stream": ["oper", "wave", "scda", "scwv", "enfo", "waef", "mmsa"],
        }

        if request is None:
            params = dict(**kwargs)
        else:
            params = dict(**request)

        defaults = DEFAULTS.get(params.get("stream"), DEFAULTS_FC)
        for key, value in defaults.items():
            params.setdefault(key, value)

        params.pop("target", None)

        params.pop(OTHER_STEP.get(params["stream"], "fcmonth"), None)

        postproc = POSPROCESSING.intersection(set(params.keys()))
        if postproc:
            warning_once("MARS post-processing keywords %r not supported", postproc)

        for_urls = defaultdict(list)
        for_index = defaultdict(list)
        ignored = set()

        def sorter(kv):
            a = kv[0]
            if a in URL_COMPONENTS:
                return URL_COMPONENTS.index(a)
            if a in INDEX_COMPONENTS:
                return INDEX_COMPONENTS.index(a) + len(URL_COMPONENTS)
            return len(URL_COMPONENTS) + len(INDEX_COMPONENTS)

        for k, v in sorted(params.items(), key=sorter):
            if isinstance(v, str):
                v = v.split("/")

            if not isinstance(v, (list, tuple)):
                v = [v]

            v = EXPAND_LIST.get(k, expand_list)(v)

            # Return canonical forms
            v = [CANONICAL.get(k, str)(x) for x in v]

            if k.startswith("_"):
                continue

            if k in POSSIBLE_VALUES:
                possible_values = POSSIBLE_VALUES[k]
                for value in v:
                    if value not in possible_values:
                        warning_once(
                            "Unknown value %r for keyword %r",
                            value,
                            k,
                            did_you_mean=(value, possible_values),
                        )

            if k in INDEX_COMPONENTS:
                for values in [self.user_to_index(k, x, params, for_index) for x in v]:
                    if not isinstance(values, (list, tuple)):
                        values = [values]
                    for value in values:
                        if value not in for_index[k]:
                            for_index[k].append(value)

            if k in URL_COMPONENTS:
                for values in [self.user_to_url(k, x, params, for_urls) for x in v]:
                    if not isinstance(values, (list, tuple)):
                        values = [values]
                    for value in values:
                        if value not in for_urls[k]:
                            for_urls[k].append(value)

            if (
                k not in URL_COMPONENTS
                and k not in INDEX_COMPONENTS
                and k not in POSPROCESSING
            ):
                ignored.add(k)

        if ignored:
            warning_once(
                "The following keywords %r are ignored",
                ignored,
                did_you_mean=(list(ignored), URL_COMPONENTS + INDEX_COMPONENTS),
            )

        if params.get("type") == "tf":
            for_index.clear()

        return (dict(**for_urls), dict(**for_index))

    def patch_stream(self, args):
        URL_STREAM_MAPPING = {
            ("oper", "06"): "scda",
            ("oper", "18"): "scda",
            ("wave", "06"): "scwv",
            ("wave", "18"): "scwv",
            #
            ("oper", "ef"): "enfo",
            ("wave", "ef"): "waef",
            ("oper", "ep"): "enfo",
            ("wave", "ep"): "waef",
            ("scda", "ef"): "enfo",
            ("scwv", "ef"): "waef",
            ("scda", "ep"): "enfo",
            ("scwv", "ep"): "waef",
        }
        stream, time, type = args["stream"], args["_H"], args["type"]

        if not self.infer_stream_keyword:
            return stream

        stream = URL_STREAM_MAPPING.get((stream, time), stream)
        stream = URL_STREAM_MAPPING.get((stream, type), stream)

        return stream
