"""Helpers shared by the security tests."""

from __future__ import annotations

import json
from collections.abc import Iterable
from dataclasses import dataclass, field
from typing import Any

from starlette.types import ASGIApp, Message, Scope


@dataclass
class AsgiResult:
    """What one raw ASGI call produced and how much of the request body it consumed."""

    status: int = 0
    body: bytes = b""
    headers: dict[str, str] = field(default_factory=dict)
    bytes_read: int = 0
    receive_calls: int = 0

    def json(self) -> Any:
        return json.loads(self.body)


async def asgi_request(
    app: ASGIApp,
    method: str,
    path: str,
    *,
    headers: dict[str, str] | None = None,
    chunks: Iterable[bytes] = (),
    host: str = "mcp.example.com",
    scheme: str = "https",
    client: tuple[str, int] = ("198.51.100.7", 4000),
) -> AsgiResult:
    """Drive ``app`` with a counting ``receive`` so tests can assert how much body was read."""
    result = AsgiResult()
    pending = list(chunks)
    raw_headers = [(b"host", host.encode())]
    for name, value in (headers or {}).items():
        raw_headers.append((name.lower().encode(), value.encode()))
    scope: Scope = {
        "type": "http",
        "asgi": {"version": "3.0"},
        "http_version": "1.1",
        "method": method,
        "scheme": scheme,
        "path": path,
        "raw_path": path.encode(),
        "query_string": b"",
        "root_path": "",
        "headers": raw_headers,
        "client": client,
        "server": (host, 443),
    }

    async def receive() -> Message:
        result.receive_calls += 1
        if not pending:
            return {"type": "http.disconnect"}
        chunk = pending.pop(0)
        result.bytes_read += len(chunk)
        return {"type": "http.request", "body": chunk, "more_body": bool(pending)}

    async def send(message: Message) -> None:
        if message["type"] == "http.response.start":
            result.status = message["status"]
            result.headers = {
                key.decode().lower(): value.decode() for key, value in message["headers"]
            }
        elif message["type"] == "http.response.body":
            result.body += message.get("body", b"")

    await app(scope, receive, send)
    return result