# Copyright (c) "Neo4j"
# Neo4j Sweden AB [https://neo4j.com]
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from __future__ import annotations

from contextlib import contextmanager
from time import monotonic

from . import _typing as t


if t.TYPE_CHECKING:
    from ._async.io import AsyncBolt
    from ._sync.io import Bolt


class Deadline:
    def __init__(self, timeout: float | None) -> None:
        if timeout is None or timeout == float("inf"):
            self._deadline = float("inf")
        else:
            self._deadline = monotonic() + timeout
        self._original_timeout = timeout

    @property
    def original_timeout(self) -> float | None:
        return self._original_timeout

    def expired(self) -> bool:
        return self.to_timeout() == 0

    def to_timeout(self) -> float | None:
        if self._deadline == float("inf"):
            return None
        timeout = self._deadline - monotonic()
        return max(0, timeout)

    def __eq__(self, other) -> bool:
        if isinstance(other, Deadline):
            return self._deadline == other._deadline
        return NotImplemented

    def __gt__(self, other) -> bool:
        if isinstance(other, Deadline):
            return self._deadline > other._deadline
        return NotImplemented

    def __ge__(self, other) -> bool:
        if isinstance(other, Deadline):
            return self._deadline >= other._deadline
        return NotImplemented

    def __lt__(self, other) -> bool:
        if isinstance(other, Deadline):
            return self._deadline < other._deadline
        return NotImplemented

    def __le__(self, other) -> bool:
        if isinstance(other, Deadline):
            return self._deadline <= other._deadline
        return NotImplemented

    @classmethod
    def from_timeout_or_deadline(
        cls, timeout: t.Self | float | None
    ) -> t.Self:
        if isinstance(timeout, (float, int)) or timeout is None:
            return cls(timeout)
        return timeout

    def __str__(self) -> str:
        return f"Deadline(timeout={self._original_timeout})"

    def __bool__(self) -> bool:
        """Whether a deadline is set (:data:`True`) or not (:data:`False`)."""
        return self._deadline != float("inf")


merge_deadlines = min


def merge_deadlines_and_timeouts(
    *deadline: Deadline | None,
) -> Deadline:
    deadlines = map(Deadline.from_timeout_or_deadline, deadline)
    return merge_deadlines(deadlines)


@contextmanager
def connection_deadline(
    connection: AsyncBolt | Bolt,
    deadline: Deadline | None,
) -> t.Generator[None]:
    with (
        connection_read_deadline(connection, deadline),
        connection_write_deadline(connection, deadline),
    ):
        yield


def connection_read_deadline(
    connection: AsyncBolt | Bolt,
    deadline: Deadline | None,
) -> t.AbstractContextManager[None]:
    return _connection_deadline_wrapper(
        deadline,
        connection.socket.get_read_deadline,
        connection.socket.set_read_deadline,
    )


def connection_write_deadline(
    connection: AsyncBolt | Bolt,
    deadline: Deadline | None,
) -> t.AbstractContextManager[None]:
    return _connection_deadline_wrapper(
        deadline,
        connection.socket.get_write_deadline,
        connection.socket.set_write_deadline,
    )


@contextmanager
def _connection_deadline_wrapper(
    deadline: Deadline | None,
    deadline_getter: t.Callable[[], Deadline],
    deadline_setter: t.Callable[[Deadline | None], None],
) -> t.Generator[None]:
    if deadline is None:
        # nothing to do here
        yield
        return
    original_deadline = deadline_getter()
    deadline = merge_deadlines(
        d for d in (deadline, original_deadline) if d is not None
    )
    deadline_setter(deadline)
    try:
        yield
    finally:
        deadline_setter(original_deadline)