import os
import re
import time
import typing

import flask
import prometheus_client
import prometheus_client.multiprocess

from oauthclientbridge import compat

if typing.TYPE_CHECKING:
    from typing import Text  # noqa: F401

registry = prometheus_client.CollectorRegistry()

if "prometheus_multiproc_dir" in os.environ:
    prometheus_client.multiprocess.MultiProcessCollector(registry)


TIME_BUCKETS = (
    0.0001,
    0.00055,
    0.001,
    0.0028,
    0.0046,
    0.0064,
    0.0082,
    0.01,
    0.028,
    0.046,
    0.064,
    0.082,
    0.1,
    0.4,
    0.7,
    1.0,
    4.0,
    7.0,
    10.0,
    float("inf"),
)

BYTE_BUCKETS = (
    8,
    22,
    36,
    50,
    64,
    176,
    288,
    400,
    512,
    1408,
    2304,
    3200,
    4096,
    11264,
    18432,
    25600,
    32768,
    90112,
    147456,
    204800,
    float("inf"),
)

RETRY_BUCKETS = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, float("inf"))

# Rest of these get populated lazily with http_%d as fallback.
HTTP_STATUS = {429: "http_too_many_requests"}

DBErrorCounter = prometheus_client.Counter(
    "oauth_database_error_total",
    "Database errors.",
    ["query", "error"],
    registry=registry,
)

DBLatencyHistorgram = prometheus_client.Histogram(
    "oauth_database_latency_seconds",
    "Database query latency.",
    ["query"],
    buckets=TIME_BUCKETS,
    registry=registry,
)

ServerErrorCounter = prometheus_client.Counter(
    "oauth_server_error_total",
    "OAuth errors returned to users.",
    ["endpoint", "status", "error"],
    registry=registry,
)

ServerLatencyHistogram = prometheus_client.Histogram(
    "oauth_server_latency_seconds",
    "Overall request latency.",
    ["endpoint", "status"],
    buckets=TIME_BUCKETS,
    registry=registry,
)

ServerRequestSizeHistogram = prometheus_client.Histogram(
    "oauth_server_request_bytes",
    "Overall request size.",
    ["endpoint", "status"],
    buckets=BYTE_BUCKETS,
    registry=registry,
)

ServerResponseSizeHistogram = prometheus_client.Histogram(
    "oauth_server_response_bytes",
    "Overall response size.",
    ["endpoint", "status"],
    buckets=BYTE_BUCKETS,
    registry=registry,
)

ClientErrorCounter = prometheus_client.Counter(
    "oauth_client_error_total",
    "OAuth errors from upstream provider.",
    ["endpoint", "status", "error"],
    registry=registry,
)

ClientRetryHistogram = prometheus_client.Histogram(
    "oauth_client_retries",
    "OAuth fetch retries.",
    ["endpoint", "status"],
    buckets=RETRY_BUCKETS,
    registry=registry,
)

ClientLatencyHistogram = prometheus_client.Histogram(
    "oauth_client_latency_seconds",
    "Overall request latency.",
    ["endpoint", "status"],
    buckets=TIME_BUCKETS,
    registry=registry,
)

ClientResponseSizeHistogram = prometheus_client.Histogram(
    "oauth_client_response_bytes",
    "Overall response size.",
    ["endpoint", "status"],
    buckets=BYTE_BUCKETS,
    registry=registry,
)


def status(code):  # type: (int) -> str
    if code not in HTTP_STATUS:
        text = compat.responses.get(code, str(code)).lower()
        HTTP_STATUS[code] = "http_%s" % re.sub(r"[ -]", "_", text)
    return HTTP_STATUS[code]


def endpoint():  # type: () -> Text
    return getattr(flask.request.url_rule, "endpoint", "notfound")


def before_request():  # type: () -> None
    flask.request._stats_latency_start_time = time.time()  # type: ignore


def after_request(response):
    request_latency = (
        time.time() - flask.request._stats_latency_start_time  # type: ignore
    )
    labels = {"endpoint": endpoint(), "status": status(response.status_code)}

    ServerLatencyHistogram.labels(**labels).observe(request_latency)
    if response.content_length is not None:
        ServerResponseSizeHistogram.labels(**labels).observe(response.content_length)
    if flask.request.content_length is not None:
        ServerRequestSizeHistogram.labels(**labels).observe(
            flask.request.content_length
        )
    return response


def export_metrics():  # type: () -> flask.Response
    text = prometheus_client.generate_latest(registry)
    return flask.Response(text, mimetype=prometheus_client.CONTENT_TYPE_LATEST)

@@ -28,22 +28,8 @@
 time
-%0Aimport typing
 %0A%0Aim
@@ -145,76 +145,8 @@
 at%0A%0A
-if typing.TYPE_CHECKING:%0A    from typing import Text  # noqa: F401%0A%0A
 regi
@@ -2876,29 +2876,18 @@
 tus(code
-)
 : 
- # type: (
 int) -%3E 
@@ -2889,16 +2889,17 @@
 ) -%3E str
+:
 %0A    if 
@@ -3100,29 +3100,16 @@
 nt()
-:  # type: () -%3E Text
+ -%3E str:
 %0A   
@@ -3198,29 +3198,16 @@
 st()
-:  # type: ()
  -%3E None
 %0A   
@@ -3202,16 +3202,17 @@
  -%3E None
+:
 %0A    fla
@@ -3276,24 +3276,78 @@
 e: ignore%0A%0A%0A
+# TODO: Figure our why I can't type annotate response%0A
 def after_re
@@ -3361,16 +3361,34 @@
 esponse)
+ -%3E flask.Response
 :%0A    re
@@ -3974,21 +3974,8 @@
 cs()
-:  # type: ()
  -%3E 
@@ -3988,16 +3988,17 @@
 Response
+:
 %0A    tex

Update stats typing