import asyncio
import asyncio.protocols
import hmac
from codecs import CodecInfo
from collections.abc import Callable, Iterable, Iterator
from hashlib import md5, sha256
from typing import Any, ClassVar, Generic, NewType, Text, TypeVar, overload
from typing_extensions import Final, Literal, TypeAlias, final

import asyncpg.pgproto.pgproto

from ..connect_utils import _ConnectionParameters
from ..pgproto.pgproto import WriteBuffer
from ..types import Attribute, Type

_T = TypeVar("_T")
_Record = TypeVar("_Record", bound=Record)
_OtherRecord = TypeVar("_OtherRecord", bound=Record)
_PreparedStatementState = TypeVar(
    "_PreparedStatementState", bound=PreparedStatementState[Any]
)

_NoTimeoutType = NewType("_NoTimeoutType", object)
_TimeoutType: TypeAlias = float | None | _NoTimeoutType

BUILTIN_TYPE_NAME_MAP: Final[dict[str, int]]
BUILTIN_TYPE_OID_MAP: Final[dict[int, str]]
NO_TIMEOUT: Final[_NoTimeoutType]

hashlib_md5 = md5  # noqa: Y026

@final
class ConnectionSettings(asyncpg.pgproto.pgproto.CodecContext):
    __pyx_vtable__: Any = ...
    def __init__(self, conn_key: Any) -> None: ...
    def add_python_codec(
        self,
        typeoid: int,
        typename: str,
        typeschema: str,
        typekind: str,
        encoder: Callable[[Any], Any],
        decoder: Callable[[Any], Any],
        format: Any,
    ) -> Any: ...
    def clear_type_cache(self) -> None: ...
    def get_data_codec(
        self, oid: int, format: Any = ..., ignore_custom_codec: bool = ...
    ) -> Any: ...
    def get_text_codec(self) -> CodecInfo: ...
    def register_data_types(self, types: Iterable[Any]) -> None: ...
    def remove_python_codec(
        self, typeoid: int, typename: str, typeschema: str
    ) -> None: ...
    def set_builtin_type_codec(
        self,
        typeoid: int,
        typename: str,
        typeschema: str,
        typekind: str,
        alias_to: str,
        format: Any = ...,
    ) -> Any: ...
    def __getattr__(self, name: str) -> Any: ...
    def __reduce__(self) -> Any: ...

@final
class PreparedStatementState(Generic[_Record]):
    closed: bool = ...
    name: str = ...
    query: str = ...
    refs: int = ...
    record_class: type[_Record] = ...
    ignore_custom_codec: bool = ...
    __pyx_vtable__: Any = ...
    def __init__(
        self,
        name: str,
        query: str,
        protocol: BaseProtocol[Any],
        record_class: type[_Record],
        ignore_custom_codec: bool,
    ) -> None: ...
    def _get_parameters(self) -> tuple[Type, ...]: ...
    def _get_attributes(self) -> tuple[Attribute, ...]: ...
    def _init_types(self) -> set[int]: ...
    def _init_codecs(self) -> None: ...
    def attach(self) -> None: ...
    def detach(self) -> None: ...
    def mark_closed(self) -> None: ...
    def __reduce__(self) -> Any: ...

class CoreProtocol:
    backend_pid: Any = ...
    backend_secret: Any = ...
    __pyx_vtable__: Any = ...
    def __init__(self, con_params: _ConnectionParameters) -> None: ...
    def is_in_transaction(self) -> bool: ...
    def __reduce__(self) -> Any: ...

class BaseProtocol(CoreProtocol, Generic[_Record]):
    queries_count: Any = ...
    is_ssl: bool = ...
    __pyx_vtable__: Any = ...
    def __init__(
        self,
        addr: Any,
        connected_fut: Any,
        con_params: _ConnectionParameters,
        record_class: type[_Record],
        loop: Any,
    ) -> None: ...
    def set_connection(self, connection: Any) -> None: ...
    def get_server_pid(self, *args: Any, **kwargs: Any) -> int: ...
    def get_settings(self, *args: Any, **kwargs: Any) -> ConnectionSettings: ...
    def get_record_class(self) -> type[_Record]: ...
    def abort(self) -> None: ...
    async def bind(
        self,
        state: PreparedStatementState[_OtherRecord],
        args: Any,
        portal_name: str,
        timeout: _TimeoutType,
    ) -> Any: ...
    @overload
    async def bind_execute(
        self,
        state: PreparedStatementState[_OtherRecord],
        args: Any,
        portal_name: str,
        limit: int,
        return_extra: Literal[False],
        timeout: _TimeoutType,
    ) -> list[_OtherRecord]: ...
    @overload
    async def bind_execute(
        self,
        state: PreparedStatementState[_OtherRecord],
        args: Any,
        portal_name: str,
        limit: int,
        return_extra: Literal[True],
        timeout: _TimeoutType,
    ) -> tuple[list[_OtherRecord], bytes, bool]: ...
    @overload
    async def bind_execute(  # pyright: ignore
        self,
        state: PreparedStatementState[_OtherRecord],
        args: Any,
        portal_name: str,
        limit: int,
        return_extra: bool,
        timeout: _TimeoutType,
    ) -> list[_OtherRecord] | tuple[list[_OtherRecord], bytes, bool]: ...
    async def bind_execute_many(
        self,
        state: PreparedStatementState[_OtherRecord],
        args: Any,
        portal_name: str,
        timeout: _TimeoutType,
    ) -> None: ...
    async def close(self, timeout: _TimeoutType) -> None: ...
    def _get_timeout(self, timeout: _TimeoutType) -> float | None: ...
    def _is_cancelling(self) -> bool: ...
    async def _wait_for_cancellation(self) -> None: ...
    async def close_statement(
        self, state: PreparedStatementState[_OtherRecord], timeout: _TimeoutType
    ) -> Any: ...
    async def copy_in(self, *args: Any, **kwargs: Any) -> str: ...
    async def copy_out(self, *args: Any, **kwargs: Any) -> str: ...
    async def execute(self, *args: Any, **kwargs: Any) -> Any: ...
    def is_closed(self, *args: Any, **kwargs: Any) -> Any: ...
    def is_connected(self, *args: Any, **kwargs: Any) -> Any: ...
    def data_received(self, data: Any) -> None: ...
    def connection_made(self, transport: Any) -> None: ...
    def connection_lost(self, exc: Exception | None) -> None: ...
    def pause_writing(self, *args: Any, **kwargs: Any) -> Any: ...
    @overload
    async def prepare(
        self,
        stmt_name: str,
        query: str,
        timeout: float | None = ...,
        *,
        state: _PreparedStatementState,
        ignore_custom_codec: bool = ...,
        record_class: None,
    ) -> _PreparedStatementState: ...
    @overload
    async def prepare(
        self,
        stmt_name: str,
        query: str,
        timeout: float | None = ...,
        *,
        state: None = ...,
        ignore_custom_codec: bool = ...,
        record_class: type[_OtherRecord],
    ) -> PreparedStatementState[_OtherRecord]: ...
    async def query(self, *args: Any, **kwargs: Any) -> str: ...
    def resume_writing(self, *args: Any, **kwargs: Any) -> Any: ...
    def __reduce__(self) -> Any: ...

@final
class Codec:
    __pyx_vtable__: Any = ...
    def __init__(
        self,
        name: str,
        schema: str,
        kind: str,
        type: int,
        format: int,
        xformat: int,
        c_encoder: Any,
    ) -> None: ...
    def __reduce__(self) -> Any: ...

class DataCodecConfig:
    __pyx_vtable__: Any = ...
    def __init__(self, cache_key: Any) -> None: ...
    def add_python_codec(
        self,
        typeoid: int,
        typename: str,
        typeschema: str,
        typekind: str,
        encoder: Callable[[ConnectionSettings, WriteBuffer, object], object],
        decoder: Callable[..., object],
        format: Any,
        xformat: Any,
    ) -> Any: ...
    def add_types(self, types: Iterable[Any]) -> Any: ...
    def clear_type_cache(self) -> None: ...
    def declare_fallback_codec(self, oid: int, name: str, schema: str) -> Codec: ...
    def remove_python_codec(
        self, typeoid: int, typename: str, typeschema: str
    ) -> Any: ...
    def set_builtin_type_codec(
        self,
        typeoid: int,
        typename: str,
        typeschema: str,
        typekind: str,
        alias_to: str,
        format: Any = ...,
    ) -> Any: ...
    def __reduce__(self) -> Any: ...

class Protocol(BaseProtocol[_Record], asyncio.protocols.Protocol): ...

class Record:
    @overload
    def get(self, key: str) -> Any | None: ...
    @overload
    def get(self, key: str, default: _T) -> Any | _T: ...
    def items(self) -> Iterator[tuple[str, Any]]: ...
    def keys(self) -> Iterator[str]: ...
    def values(self) -> Iterator[Any]: ...
    @overload
    def __getitem__(self, index: str) -> Any: ...
    @overload
    def __getitem__(self, index: int) -> Any: ...
    @overload
    def __getitem__(self, index: slice) -> tuple[Any, ...]: ...
    def __iter__(self) -> Iterator[Any]: ...
    def __contains__(self, x: object) -> bool: ...
    def __len__(self) -> int: ...

class Timer:
    def __init__(self, budget: float | None) -> None: ...
    def __enter__(self) -> None: ...
    def __exit__(self, et: object, e: object, tb: object) -> None: ...
    def get_remaining_budget(self) -> float: ...
    def has_budget_greater_than(self, amount: float) -> bool: ...

@final
class SCRAMAuthentication:
    AUTHENTICATION_METHODS: ClassVar[list[str]]
    DEFAULT_CLIENT_NONCE_BYTES: ClassVar[int]
    DIGEST = sha256  # noqa: Y026
    REQUIREMENTS_CLIENT_FINAL_MESSAGE: ClassVar[list[str]]
    REQUIREMENTS_CLIENT_PROOF: ClassVar[list[str]]
    SASLPREP_PROHIBITED: ClassVar[tuple[Callable[[Text], bool], ...]]
    authentication_method: bytes
    authorization_message: bytes | None
    client_channel_binding: bytes
    client_first_message_bare: bytes | None
    client_nonce: bytes | None
    client_proof: bytes | None
    password_salt: bytes | None
    password_iterations: int
    server_first_message: bytes | None
    server_key: hmac.HMAC | None
    server_nonce: bytes | None
