import urllib.parse
from typing import Any, Iterator, List, Optional, Tuple

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

DEFAULT_SKIP_KEYS = ["orderBy", "order_by"]


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

    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)
        return

    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
            else:
                yield from [
                    (key, entry) for entry in value.split(",") if entry.strip() != ""
                ]