import urllib.parse
from collections.abc import Iterator
from typing import Any

from starlette.datastructures import QueryParams
from starlette.types import ASGIApp, Receive, Scope, Send


class QueryStringFlatteningMiddleware:
    def __init__(self, app: ASGIApp, skip_keys: list[str] | None = None) -> None:
        self.app = app
        self.skip_keys = skip_keys or []

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] not in {"http", "websocket"}:
            await self.app(scope, receive, send)
            return

        if query_string := scope.get("query_string", ""):
            flattened_params = list(self._yield_flattened_params(query_string))
            scope["query_string"] = urllib.parse.urlencode(
                flattened_params, doseq=True
            ).encode("utf-8")

        await self.app(scope, receive, send)

    def _yield_flattened_params(self, query_string: str) -> Iterator[tuple[str, Any]]:
        for key, value in QueryParams(query_string).multi_items():
            if key in self.skip_keys:
                yield key, value.strip()
            else:
                yield from [
                    (key, entry) for entry in value.split(",") if entry.strip() != ""
                ]