import logging
from os import getenv
from typing import List, Union

import requests
from typeguard import typechecked

from loki_api import get_logger, log_it

GRAFANA_LOGS_URL = getenv(
    "GRAFANA_LOGS_URL", "https://logs-prod-us-central1.grafana.net"
)


@typechecked
class LokiClient:
    """Python client for interacting with the Grafana [Loki HTTP API](https://grafana.com/docs/loki/latest/api/).
    By default, this module will query https://logs-prod-us-central1.grafana.net,
    however, this can be overridden with the `GRAFANA_LOGS_URL` environment variable.

    The Loki API itself doesn't support authentication as noted [here](https://grafana.com/docs/loki/latest/operations/authentication/)
    However, this client does support HTTP Basic Auth and requires a username
    and password to initialize a new client.

    Parameters
    ----------
    username: str
        Your username when using basic auth

    password: str
        Your password when using basic auth

    Usage
    -----
    ```
    client = LokiClient("user", "pass")

    # Total requests made against a datasource in the past 30 days
    client.query_range(
        'sum(count_over_time({kind="usage_insights"} | json | datasourceName="TS - psi-prd - 693257622622" | eventName="data-request" [30d]))',
    )
    ```
    """

    def __init__(
        self, username: str, password: str, logger: logging.Logger = None
    ):
        self._session = requests.Session()
        self._session.auth = (username, password)
        self.logger = logger or get_logger()

    def query_range(
        self,
        query: str,
        start: Union[str, int, float] = None,
        end: Union[str, int, float] = None,
        step: float = None,
        interval: float = None,
        direction: str = None,
        limit: int = None,
    ) -> List[dict]:
        """
        Parameters
        ----------
        query: str
            The [LogQL](https://grafana.com/docs/loki/latest/logql/) query to perform

        start: Union[str, int, float]
            The start time for the query as a nanosecond Unix epoch or
            another supported format. Defaults to one hour ago.

        end: Union[str, int, float]
            The end time for the query as a nanosecond Unix epoch or another
            supported format. Defaults to now.

        step: float
            Query resolution step width in duration format or float number of seconds.
            duration refers to Prometheus duration strings of the form [0-9]+[smhdwy].
            For example, 5m refers to a duration of 5 minutes.
            Defaults to a dynamic value based on start and end.
            Only applies to query types which produce a matrix response.

        interval: float
            This parameter is experimental; see the explanation under Step versus interval.
            Only return entries at (or greater than) the specified interval,
            can be a duration format or float number of seconds.
            Only applies to queries which produce a stream response.

        direction: str
            Determines the sort order of logs
            Supported values are forward or backward. Defaults to backward.

        limit: int
            The max number of entries to return. It defaults to 100.
            Only applies to query types which produce a stream(log lines) response.
        """
        assert direction in [
            "backward",
            "forward",
            None,
        ], "direction must be 'forward' or 'backward'"

        return self._make_request(
            "loki/api/v1/query_range",
            query=query,
            start=start,
            end=end,
            step=step,
            interval=interval,
            direction=direction,
            limit=limit,
        )

    def query(
        self,
        query: str,
        time: Union[str, int, float] = None,
        direction: str = None,
        limit: int = None,
    ) -> List[dict]:
        """
        query: str
            The [LogQL](https://grafana.com/docs/loki/latest/logql/) query to perform

        time: Union[str, int, float]
            The evaluation time for the query as a nanosecond Unix epoch or
            another supported format. Defaults to now.

        direction: str
            Determines the sort order of logs
            Supported values are forward or backward. Defaults to backward.

        limit: int
            The max number of entries to return. It defaults to 100.
            Only applies to query types which produce a stream(log lines) response.
        """
        assert direction in [
            "backward",
            "forward",
            None,
        ], "direction must be 'forward' or 'backward'"

        return self._make_request(
            "loki/api/v1/query",
            query=query,
            time=time,
            direction=direction,
            limit=limit,
        )

    @log_it
    def _make_request(self, endpoint: str, **kwargs):
        params = {k: v for k, v in kwargs.items() if v}
        response = self._session.get(
            url=f"{GRAFANA_LOGS_URL}/{endpoint}",
            params=params,
        )
        try:
            response.raise_for_status()
        except requests.exceptions.HTTPError as ex:
            self.logger.error(
                f"[{response.status_code}] {response.reason} - {response.text}"
            )
            raise ex
        result = response.json()
        return result["data"]["result"]

    def __repr__(self):
        return f"{self.__class__.__name__}()"
