# Copyright 2016 Criteo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Adapter between BigGraphite and Graphite."""
from __future__ import absolute_import  # Otherwise graphite is this module.

import threading
import time
from datetime import datetime

from graphite import carbonlink
from graphite import intervals
from graphite import node
from graphite import readers
from graphite.logger import log
from graphite.render import hashing

from biggraphite import accessor as bg_accessor
from biggraphite import accessor_cache as bg_accessor_cache
from biggraphite import glob_utils
from biggraphite import graphite_utils
from biggraphite import metric as bg_metric

_CONFIG_NAME = "biggraphite"

try:
    from graphite.readers import utils as readers_utils

    BaseReader = readers_utils.BaseReader
except ImportError:
    BaseReader = object
try:
    from graphite.finders import utils as finders_utils

    BaseFinder = finders_utils.BaseFinder
except ImportError:
    BaseFinder = object
try:
    FetchInProgress = readers.FetchInProgress
except AttributeError:
    FetchInProgress = None


class Error(Exception):
    """Base class for all exceptions from this module."""


_DISABLED = "DISABLED"


class Reader(BaseReader):
    """As per the Graphite API, fetches points for and metadata for a given metric."""

    __slots__ = (
        "_accessor",
        "_metadata_cache",
        "_carbonlink",
        "_metric",
        "_metric_name",
    )

    def __init__(self, accessor, metadata_cache, carbonlink, metric_name, metric=None):
        """Create a new reader."""
        assert accessor
        assert metadata_cache

        self._accessor = accessor
        self._metadata_cache = metadata_cache
        self._carbonlink = carbonlink
        self._metric = metric
        self._metric_name = metric_name

    def __get_cached_datapoints(self, stage):
        cached_datapoints = []
        try:
            # TODO: maybe this works with non stage0 now, need to check.
            if stage.stage0 and self._carbonlink:
                cached_datapoints = self._carbonlink.query(self._metric_name)
        except Exception:
            log.exception("Failed CarbonLink query '%s'" % self._metric_name)
        return cached_datapoints

    def __get_metadata(self):
        self.__refresh_metric()
        if self._metric is not None:
            return self._metric.metadata
        else:
            return bg_metric.MetricMetadata.create()

    def __get_time_info(self, start_time, end_time, now, shift=False):
        """Constrain the provided range in an aligned interval within retention."""
        retention = (
            self.__get_metadata().retention
            or bg_metric.MetricMetadata._DEFAULT_RETENTION
        )
        return retention.align_time_window(start_time, end_time, now, shift=shift)

    def __refresh_metric(self):
        """Set self._metric."""
        # TODO: This can be blocking, so we need to do one of:
        # - forward metadata to the constructor (possibly by adding mget
        #   to the metadata cache and the accessor)
        # - make fetch_async really async
        # - use a thread pool in `fetch_multi`
        if self._metric is None:
            self._metric = self._metadata_cache.get_metric(self._metric_name)

    def _merge_cached_points(
        self,
        stage,
        start,
        step,
        aggregation_method,
        points,
        cached_datapoints,
        raw_step=None,
    ):
        """Merge points from carbonlink with points from database."""
        if isinstance(cached_datapoints, dict):
            cached_datapoints = list(cached_datapoints.items())

        return readers.utils.merge_with_cache(
            cached_datapoints,
            start,
            step,
            points,
            func=aggregation_method,
            raw_step=raw_step,
        )

    def fetch_async(self, start_time, end_time, now=None, requestContext=None):
        """Fetch point for a given interval as per the Graphite API.

        Args:
          start_time: Timestamp to fetch points from, will constrained by retention policy.
          end_time: Timestamp to fetch points until, will constrained by retention policy.
          now: Current timestamp as a float, defaults to time.time(), for tests.

        Returns:
          A callable that returns a tuple made of (rounded start time,
          rounded end time, stage precision), points
          Points is a list for which missing points are set to None.
        """
        fetch_start = time.time()
        log.rendering(
            "fetch(%s, %d, %d) - start" % (self._metric_name, start_time, end_time)
        )

        self.__refresh_metric()
        if now is None:
            now = time.time()

        metadata = self.__get_metadata()
        start_time, end_time, stage = self.__get_time_info(start_time, end_time, now)
        start_step = stage.step(start_time)
        points_num = stage.step(end_time) - start_step
        step = stage.precision
        aggregation_method = metadata.aggregator.carbon_name
        raw_step = metadata.retention.stage0.precision

        if not self._metric:
            # The metric doesn't exist, let's fail gracefully.
            ts_and_points = []
        else:
            # This returns a generator which we can iterate on later.
            ts_and_points = self._accessor.fetch_points(
                self._metric, start_time, end_time, stage
            )

        def read_points():
            read_start = time.time()

            cached_datapoints = self.__get_cached_datapoints(stage)

            # TODO: Consider wrapping an array (using NaN for None) for
            # speed&memory efficiency
            points = [None] * points_num
            for ts, point in ts_and_points:
                index = stage.step(ts) - start_step
                points[index] = point

            if cached_datapoints:
                points = self._merge_cached_points(
                    stage,
                    start_step,
                    step,
                    aggregation_method,
                    points,
                    cached_datapoints,
                    raw_step=raw_step,
                )

            now = time.time()
            log.rendering(
                "fetch(%s, %d, %d) - %d points - read: %f secs - total: %f secs"
                % (
                    self._metric_name,
                    start_time,
                    end_time,
                    len(points),
                    now - read_start,
                    now - fetch_start,
                )
            )
            return (start_time, end_time, stage.precision), points

        log.rendering(
            "fetch(%s, %d, %d) - started" % (self._metric_name, start_time, end_time)
        )

        return read_points

    def fetch(self, start_time, end_time, now=None, requestContext=None):
        """Fetch point for a given interval as per the Graphite API.

        Args:
          start_time: Timestamp to fetch points from, will constrained by retention policy.
          end_time: Timestamp to fetch points until, will constrained by retention policy.
          now: Current timestamp as a float, defaults to time.time(), for tests.

        Returns:
          A tuple made of (rounded start time, rounded end time, stage precision), points
          Points is a list for which missing points are set to None.
        """
        results = self.fetch_async(start_time, end_time, now=now, requestContext=None)
        if FetchInProgress:
            return FetchInProgress(results)