import json
from typing import Callable

from fastapi import Request
from prometheus_client import Counter, Gauge
from prometheus_fastapi_instrumentator import Instrumentator, metrics
from prometheus_fastapi_instrumentator.metrics import Info
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import Message

SUBSYSTEM = "model"
NAMESPACE = "turing_api"

# Define Prometheus metrics
instrumentator = Instrumentator(
    should_group_status_codes=False,
    should_ignore_untemplated=True,
    should_respect_env_var=False,
    should_instrument_requests_inprogress=True,
    excluded_handlers=["/metrics"],
    inprogress_name="fastapi_inprogress",
    inprogress_labels=True
)

## Define custom metric for tracking requested languages
def http_requested_languages_total(
        metric_name: str = "Total HTTP requested languages",
        metric_description: str = "Total number of HTTP requests per programming language", 
        metric_namespace: str = NAMESPACE,
        metric_subsystem: str = SUBSYSTEM  ) -> Callable[[Info],None]:
    METRIC = Counter(
        metric_name,
        metric_description,
        namespace=metric_namespace,
        subsystem=metric_subsystem,
        labelnames=["language"]
    )
    async def instrumentation(info: Info) -> None:
        try:
            if info.modified_handler != "/predict":
                return
            lang = info.request.query_params.get("language")
        except Exception:
            print("Failed to get language from request")
            lang = "other"

        METRIC.labels(language=lang).inc()

    return instrumentation


## Define custom metrics for tracking code comments in requests
http_request_code_comments_total = Counter (
        "Total HTTP request code comments",
        "Total number of comments in HTTP requests", 
        namespace=NAMESPACE,
        subsystem=SUBSYSTEM,
        labelnames=["language"]
    )

## Define custom metrics for tracking characters in code comments
http_request_comment_characters_total = Counter(
        "Total HTTP request code comment characters",
        "Total number of characters in the HTTP requests", 
        namespace=NAMESPACE,
        subsystem=SUBSYSTEM,
        labelnames=["endpoint","language"]
    )
   

## Define custom metric for tracking maximum characters per comment
http_request_maximum_characters_per_comment = Gauge(
        "Maximum characters per comment",
        "Maximum number of characters in a single comment from HTTP requests", 
        namespace=NAMESPACE,
        subsystem=SUBSYSTEM,
        labelnames=["language"]
    )


## Middleware to extract and record metrics from request body
class PrometheusBodyMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if request.url.path != "/predict":
            return await call_next(request)
        
        body_bytes = await request.body()
        query_params = request.query_params
        try:
            if body_bytes:
                language = query_params.get("language", "unknown")
                body_json = json.loads(body_bytes)
                print(f"Request body JSON: {body_json}")
                texts = body_json.get("texts")
                if texts:
                    total_characters = sum(len(example) for example in texts if example)
                    max_characters = max((len(example) for example in texts if example), default=0)
                    http_request_comment_characters_total.labels(endpoint="/predict", language=language).inc(total_characters)
                    http_request_maximum_characters_per_comment.labels(language=language).set(max_characters)
                    http_request_code_comments_total.labels(language=language).inc(len(texts))
        except (json.JSONDecodeError, UnicodeDecodeError):
            pass

        async def receive() -> Message:
            return {"type": "http.request", "body": body_bytes}
        
        request._receive = receive

        response = await call_next(request)
        return response

## Register metrics with the instrumentator
instrumentator.add(
    metrics.request_size(
        should_include_handler=True,
        should_include_method=False,
        should_include_status=True,
        metric_namespace=NAMESPACE,
        metric_subsystem=SUBSYSTEM,
    )
).add(
    metrics.response_size(
        should_include_handler=True,
        should_include_method=False,
        should_include_status=True,
        metric_namespace=NAMESPACE,
        metric_subsystem=SUBSYSTEM
    )
).add(
    http_requested_languages_total()
).add(
    metrics.requests(
        should_include_handler=True,
        should_include_method=True,
        should_include_status=True,
        metric_namespace=NAMESPACE,
        metric_subsystem=SUBSYSTEM
    )
).add(
    metrics.latency(
        should_include_handler=True,        
        should_include_method=False,
        should_include_status=True,
        metric_namespace=NAMESPACE,
        metric_subsystem=SUBSYSTEM
    )
)