from __future__ import annotations

import sys
from typing import TextIO

from .profiler import RequestProfiler
from .render import render_trace


class WaterfallMiddleware:
    def __init__(
        self,
        app,
        *,
        stream: TextIO | None = None,
        width: int = 60,
        hotspot_count: int = 5,
        max_nodes: int = 5000,
    ) -> None:
        self.app = app
        self.stream = stream if stream is not None else sys.stdout
        self.width = width
        self.hotspot_count = hotspot_count
        self.max_nodes = max_nodes

    async def __call__(self, scope, receive, send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        method = scope.get("method", "HTTP")
        path = scope.get("path", "/")
        status_code = 500
        profiler = RequestProfiler(method, path, max_nodes=self.max_nodes)

        async def send_wrapper(message) -> None:
            nonlocal status_code
            if message["type"] == "http.response.start":
                status_code = message["status"]
            await send(message)

        profiler.start()
        try:
            await self.app(scope, receive, send_wrapper)
        finally:
            trace = profiler.stop(status_code=status_code)
            self.stream.write(
                render_trace(
                    trace,
                    width=self.width,
                    hotspot_count=self.hotspot_count,
                )
            )
            self.stream.write("\n")
            self.stream.flush()
