import datetime as dt
import urllib.parse as urlparse

from collections import defaultdict
from concurrent import futures

import pandas as pd
import panel as pn
import param  # type: ignore
import requests

from ..util import parse_timedelta
from .base import Source, cached


class PrometheusSource(Source):
    """
    `PrometheusSource` allows querying Prometheus PromQL endpoints.

    The `PrometheusSource` is configured to return timeseries about CPU,
    memory and network usage as well as restarts for a list of
    Kubernetes pods specified by ID.
    """

    ae5_source = param.Parameter(doc="""
      An AE5Source instance to use for querying.""")

    cache_per_query = param.Boolean(default=False, doc="""
        Whether to query the whole dataset or individual queries.""")

    ids = param.List(default=[], doc="""
      List of pod IDs to query.""")

    metrics = param.List(
        default=['memory_usage', 'cpu_usage', 'restarts',
                 'network_receive_bytes', 'network_transmit_bytes'],
        doc="Names of metric queries to execute")

    promql_api = param.String(doc="""
      Name of the AE5 deployment exposing the Prometheus API""")

    period = param.String(default='3h', doc="""
      Period to query over specified as a string. Supports:

        - Week:   '1w'
        - Day:    '1d'
        - Hour:   '1h'
        - Minute: '1m'
        - Second: '1s'
    """)

    samples = param.Integer(default=200, doc="""
      Number of samples in the selected period to query. May be
      overridden by explicit step value.""")

    shared = param.Boolean(default=False, readonly=True, doc="""
      PrometheusSource cannot be shared because it has per-user state.""")

    step = param.String(doc="""
        Step value to use in PromQL query_range query.""")

    source_type = 'prometheus'

    _memory_usage_query = """sum by(container_name)
    (container_memory_usage_bytes{job="kubelet",
    cluster="", namespace="default", pod_name=POD_NAME,
    container_name=~"app", container_name!="POD"})"""

    _network_receive_bytes_query = """sort_desc(sum by (pod_name)
    (rate(container_network_receive_bytes_total{job="kubelet", cluster="",
    namespace="default", pod_name=POD_NAME}[1m])))"""

    _network_transmit_bytes_query = """sort_desc(sum by (pod_name)
    (rate(container_network_transmit_bytes_total{job="kubelet", cluster="",
    namespace="default", pod_name=POD_NAME}[1m])))"""

    _cpu_usage_query = """sum by (container_name)
    (rate(container_cpu_usage_seconds_total{job="kubelet", cluster="",
     namespace="default", image!="", pod_name=POD_NAME,
     container_name=~"app", container_name!="POD"}[1m]))"""

    _restarts_query = """max by (container)
     (kube_pod_container_status_restarts_total{job="kube-state-metrics",
    cluster="", namespace="default", pod=POD_NAME,
    container=~"app"})"""

    _metrics = {
        'memory_usage': {
            'query': _memory_usage_query,
            'schema': {"type": "number"}
        },
        'network_receive_bytes': {
            'query': _network_receive_bytes_query,
            'schema': {"type": "number"}
        },
        'network_transmit_bytes': {
            'query': _network_transmit_bytes_query,
            'schema': {"type": "number"}
        },
        'cpu_usage': {
            'query': _cpu_usage_query,
            'schema': {"type": "number"}
        },
        'restarts': {
            'query': _restarts_query,
            'schema': {"type": "number"}
        }
    }

    @property
    def panel(self):
        return pn.Param(self.param, parameters=['period', 'samples', 'step'],
                        sizing_mode='stretch_width', show_name=False)

    def _format_timestamps(self):
        end = dt.datetime.utcnow()
        end_formatted = end.isoformat("T") + "Z"
        period = parse_timedelta(self.period)
        if period is None:
            raise ValueError(f"Could not parse period '{self.period}'. "
                             "Must specify weeks ('1w'), days ('1d'), "
                             "hours ('1h'), minutes ('1m'), or "
                             "seconds ('1s').")
        start = end - period
        start_formatted = start.isoformat("T") + "Z"
        return start_formatted, end_formatted

    def _step_value(self):
        if self.step:
            return self.step
        return int((parse_timedelta(self.period)/self.samples).total_seconds())

    def _url_query_parameters(self, pod_id, query):
        """
        Uses regular expression to map ae5-tools pod_id to full id.
        """
        start_timestamp, end_timestamp = self._format_timestamps()
        regexp = f'anaconda-app-{pod_id}-.*'
        query = query.replace("pod_name=POD_NAME", f"pod_name=~'{regexp}'")
        query = query.replace("pod=POD_NAME", f"pod=~'{regexp}'")
        query_params = {
            'query': query, 'start': start_timestamp,
            'end': end_timestamp, 'step': self._step_value()
        }
        return urlparse.urlencode(query_params)

    def _get_query_url(self, metric, pod_id):
        "Return the full query URL"
        query_template = self._metrics[metric]['query']
        query_params = self._url_query_parameters(pod_id, query_template)
        if self.ae5_source:
            ae5 = self.ae5_source._session
            return f'https://{ae5._k8s_endpoint}.{ae5.hostname}/promql/query_range?{query_params}'
        else:
            return f'{self.promql_api}/query_range?{query_params}'

    def _get_query_json(self, query_url):
        "Function called in parallel to fetch JSON in ThreadPoolExecutor"
        if self.ae5_source:
            response = self.ae5_source._session.session.get(query_url, verify=False)
        else:
            response = requests.get(query_url, verify=False)
        data = response.json()
        if len(data) == 0:
            return None
        return data

    def _json_to_df(self, metric, values):
        "Convert JSON response to pandas DataFrame"
        df = pd.DataFrame(values, columns=['timestamp', metric])
        df[metric] = df[metric].astype(float)
        df['timestamp'] = pd.to_datetime(df['timestamp'], unit='s')
        return df.set_index('timestamp')

    def _fetch_data(self, pod_ids):
        "Returns fetched JSON in dictionary indexed by pod_id then metric name"
        if not pod_ids:
            return {}
        # ToDo: Remove need to slice
        triples = [
            (pod_id, metric, self._get_query_url(metric, pod_id[3:]))
            for pod_id in pod_ids for metric in self.metrics
        ]
        fetched_json = defaultdict(dict)
        with futures.ThreadPoolExecutor(len(triples)) as executor:
            tasks = {
                executor.submit(self._get_query_json, t[2]): t
                for t in triples
            }
            for future in futures.as_completed(tasks):
                (pod_id, metric, query_url) = tasks[future]
                try:
                    fetched_json[pod_id][metric] = future.result()
                except Exception as e:
                    fetched_json[pod_id][metric] = []
                    self.param.warning(
                        f"Could not fetch {metric} for pod {pod_id}. "
                        f"Query used: {query_url}, errored with {type(e)}({e})."
                    )
        return fetched_json

    def _make_query(self):
        json_data = self._fetch_data(self.ids)
        dfs = []
        for pod_id, pod_data in json_data.items():
            df = None
            for metric in self.metrics:
                data_df = self._json_to_df(metric, pod_data[metric])
                if df is None:
                    df = data_df
                else:
                    df = pd.merge(df, data_df, on='timestamp', how='outer')
            df = df.reset_index()
            df.insert(0, 'id', pod_id)
            dfs.append(df)
        if dfs:
            return pd.concat(dfs)
        else:
            return pd.DataFrame(columns=list(self.get_schema('timeseries')))

    def get_schema(self, table=None):
        dt_start, dt_end = self._format_timestamps()
        schema = {
            "id": {
                "type": "string",
                "enum": self.ids
            },
            "timestamp" : {
                "type": "string",
                "format": "date-time",
                "formatMinimum": dt_start,
                "formatMaximum": dt_end
            }
        }
        for k, mdef in self._metrics.items():
            schema[k] = mdef['schema']
        return {"timeseries": schema} if table is None else schema

    @cached
    def get(self, table, **query):
        if table not in ('timeseries',):
            raise ValueError(f"PrometheusSource has no '{table}' table, "
                             "it currently only has a 'timeseries' table.")
        return self._make_query()

    def get_tables(self):
        return list(self._metrics)
