from __future__ import annotations as _annotations

from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import cast

from acp import connect_to_agent
from acp.client.connection import ClientSideConnection
from acp.interfaces import Client
from websockets.asyncio.client import ClientConnection, connect

from ._auth import bearer_headers
from ._config import TransportOptions
from ._stream_adapter import WebSocketStreamBridge, open_websocket_stream_bridge

__all__ = ("RemoteClientConnection", "connect_remote_agent")

_HeaderPairs = Mapping[str, str] | Sequence[tuple[str, str]]


@dataclass(slots=True)
class RemoteClientConnection:
    connection: ClientSideConnection
    websocket: ClientConnection
    streams: WebSocketStreamBridge

    async def close(self) -> None:
        await self.connection.close()
        await self.streams.close()


async def connect_remote_agent(
    client: Client,
    url: str,
    *,
    options: TransportOptions | None = None,
    headers: _HeaderPairs | None = None,
    bearer_token: str | None = None,
) -> RemoteClientConnection:
    resolved_options = options or TransportOptions()
    resolved_headers = _merge_headers(headers, bearer_headers(bearer_token))
    websocket = await connect(
        url,
        additional_headers=resolved_headers,
        compression=resolved_options.compression,
        open_timeout=resolved_options.open_timeout,
        ping_interval=resolved_options.ping_interval,
        ping_timeout=resolved_options.ping_timeout,
        close_timeout=resolved_options.close_timeout,
        max_size=resolved_options.max_size,
        max_queue=resolved_options.max_queue,
    )
    streams = await open_websocket_stream_bridge(
        websocket,
        reader_limit=resolved_options.reader_limit,
    )
    connection = connect_to_agent(client, streams.writer, streams.reader)
    return RemoteClientConnection(connection=connection, websocket=websocket, streams=streams)


def _merge_headers(
    base: _HeaderPairs | None,
    extra: dict[str, str] | None,
) -> _HeaderPairs | None:
    if not extra:
        return base
    if base is None:
        return extra
    if isinstance(base, Mapping):
        base_mapping = cast(Mapping[str, str], base)
        merged_items: list[tuple[str, str]] = [*base_mapping.items(), *extra.items()]
        return merged_items
    base_pairs = cast(Sequence[tuple[str, str]], base)
    merged_pairs: list[tuple[str, str]] = [*base_pairs, *extra.items()]
    return merged_pairs
