# Copyright (C) 2018-2020  The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information

"""Elastic Search backend

"""

import datetime  # noqa
import logging

from copy import deepcopy
from typing import Any, Dict

from elasticsearch import helpers

from swh.core import utils

logger = logging.getLogger(__name__)


DEFAULT_CONFIG = {
    "elasticsearch": {
        "cls": "local",
        "args": {
            "index_name_prefix": "swh-tasks",
            "storage_nodes": ["localhost:9200"],
            "client_options": {
                "sniff_on_start": False,
                "sniff_on_connection_fail": True,
                "http_compress": False,
                "sniffer_timeout": 60,
            },
        },
    }
}


def get_elasticsearch(cls: str, args: Dict[str, Any] = {}):
    """Instantiate an elastic search instance

    """
    if cls == "local":
        from elasticsearch import Elasticsearch
    elif cls == "memory":
        from .elasticsearch_memory import MemoryElasticsearch as Elasticsearch  # type: ignore  # noqa
    else:
        raise ValueError("Unknown elasticsearch class `%s`" % cls)

    return Elasticsearch(**args)


class ElasticSearchBackend:
    """ElasticSearch backend to index tasks

    This uses an elasticsearch client to actually discuss with the
    elasticsearch instance.

    """

    def __init__(self, **config):
        self.config = deepcopy(DEFAULT_CONFIG)
        self.config.update(config)
        es_conf = self.config["elasticsearch"]
        args = deepcopy(es_conf["args"])
        self.index_name_prefix = args.pop("index_name_prefix")
        self.storage = get_elasticsearch(
            cls=es_conf["cls"],
            args={
                "hosts": args.get("storage_nodes", []),
                **args.get("client_options", {}),
            },
        )
        # document's index type (cf. /data/elastic-template.json)
        self.doc_type = "task"

    def initialize(self):
        self.storage.indices.put_mapping(
            index=f"{self.index_name_prefix}-*",
            doc_type=self.doc_type,
            # to allow type definition below
            include_type_name=True,
            # to allow install mapping even if no index yet
            allow_no_indices=True,
            body={
                "properties": {
                    "task_id": {"type": "double"},
                    "task_policy": {"type": "text"},
                    "task_status": {"type": "text"},
                    "task_run_id": {"type": "double"},
                    "arguments": {
                        "type": "object",
                        "properties": {
                            "args": {"type": "nested", "dynamic": False},
                            "kwargs": {"type": "text"},
                        },
                    },
                    "type": {"type": "text"},
                    "backend_id": {"type": "text"},
                    "metadata": {"type": "object", "enabled": False},
                    "scheduled": {
                        "type": "date",
                        "format": "yyyy-MM-dd HH:mm:ss||yyyy-MM-dd||strict_date_optional_time||epoch_millis",  # noqa
                    },
                    "started": {
                        "type": "date",
                        "format": "yyyy-MM-dd HH:mm:ss||yyyy-MM-dd||strict_date_optional_time||epoch_millis",  # noqa
                    },
                    "ended": {
                        "type": "date",
                        "format": "yyyy-MM-dd HH:mm:ss||yyyy-MM-dd||strict_date_optional_time||epoch_millis",  # noqa
                    },
                }
            },
        )

        self.storage.indices.put_settings(
            index=f"{self.index_name_prefix}-*",
            allow_no_indices=True,
            body={
                "index": {
                    "codec": "best_compression",
                    "refresh_interval": "1s",
                    "number_of_shards": 1,
                }
            },
        )

    def create(self, index_name) -> None:
        """Create and initialize index_name with mapping for all indices
           matching `swh-tasks-` pattern

        """
        assert index_name.startswith(self.index_name_prefix)
        self.storage.indices.create(index_name)

    def compute_index_name(self, year, month):
        """Given a year, month, compute the index's name.

        """
        return "%s-%s-%s" % (self.index_name_prefix, year, "%02d" % month)

    def mget(self, index_name, doc_ids, chunk_size=500, source=True):
        """Retrieve document's full content according to their ids as per
           source's setup.

           The `source` allows to retrieve only what's interesting, e.g:
           - source=True ; gives back the original indexed data
           - source=False ; returns without the original _source field
           - source=['task_id'] ; returns only task_id in the _source field

        Args:
            index_name (str): Name of the concerned index.
            doc_ids (generator): Generator of ids to retrieve
            chunk_size (int): Number of documents chunk to send for retrieval
            source (bool/[str]): Source of information to return

        Yields:
            document indexed as per source's setup

        """
        if isinstance(source, list):
            source = {"_source": ",".join(source)}
        else:
            source = {"_source": str(source).lower()}

        for ids in utils.grouper(doc_ids, n=1000):
            res = self.storage.mget(
                body={"ids": list(ids)},
                index=index_name,
                doc_type=self.doc_type,
                params=source,
            )
            if not res:
                logger.error("Error during retrieval of data, skipping!")
                continue

            for doc in res["docs"]:
                found = doc.get("found")
                if not found:
                    msg = "Doc id %s not found, not indexed yet" % doc["_id"]
                    logger.warning(msg)
                    continue
                yield doc["_source"]

    def _streaming_bulk(self, index_name, doc_stream, chunk_size=500):
        """Bulk index data and returns the successful indexed data's
           identifier.

        Args:
            index_name (str): Name of the concerned index.
            doc_stream (generator): Generator of documents to index
            chunk_size (int): Number of documents chunk to send for indexation

        Yields:
            document id indexed

        """
        actions = (
            {
                "_index": index_name,
                "_op_type": "index",
                "_type": self.doc_type,
                "_source": data,
            }
            for data in doc_stream
        )
        for ok, result in helpers.streaming_bulk(
            client=self.storage,
            actions=actions,
            chunk_size=chunk_size,
            raise_on_error=False,
            raise_on_exception=False,
        ):
            if not ok:
                logger.error("Error during %s indexation. Skipping.", result)
                continue
            yield result["index"]["_id"]

    def is_index_opened(self, index_name: str) -> bool:
        """Determine if an index is opened or not

        """
        try:
            self.storage.indices.stats(index_name)
            return True
        except Exception:
            # fails when indice is closed (no other api call found)
            return False

    def streaming_bulk(self, index_name, doc_stream, chunk_size=500, source=True):
        """Bulk index data and returns the successful indexed data as per
           source's setup.

           the `source` permits to retrieve only what's of interest to
           us, e.g:

           - source=True ; gives back the original indexed data
           - source=False ; returns without the original _source field
           - source=['task_id'] ; returns only task_id in the _source field

           Note that:
           - if the index is closed, it will be opened
           - if the index does not exist, it will be created and opened

           This keeps the index opened for performance reasons.

        Args:
            index_name (str): Name of the concerned index.
            doc_stream (generator): Document generator to index
            chunk_size (int): Number of documents chunk to send
            source (bool, [str]): the information to return

        """
        # index must exist
        if not self.storage.indices.exists(index_name):
            self.create(index_name)
        # index must be opened
        if not self.is_index_opened(index_name):
            self.storage.indices.open(index_name)

        indexed_ids = self._streaming_bulk(
            index_name, doc_stream, chunk_size=chunk_size
        )
        yield from self.mget(
            index_name, indexed_ids, chunk_size=chunk_size, source=source
        )
