import asyncio
import logging
import socket
from collections.abc import AsyncIterator, Sequence
from contextlib import AbstractAsyncContextManager
from typing import Any, Final

import hiredis

from ai.backend.logging import BraceStyleAdapter

from .types import RedisTarget, aobject

__all__ = (
    "RedisClient",
    "RedisConnection",
)


_keepalive_options: dict[int, int] = {}


# macOS does not support several TCP_ options
# so check if socket package includes TCP options before adding it
if (_TCP_KEEPIDLE := getattr(socket, "TCP_KEEPIDLE", None)) is not None:
    _keepalive_options[_TCP_KEEPIDLE] = 20

if (_TCP_KEEPINTVL := getattr(socket, "TCP_KEEPINTVL", None)) is not None:
    _keepalive_options[_TCP_KEEPINTVL] = 5

if (_TCP_KEEPCNT := getattr(socket, "TCP_KEEPCNT", None)) is not None:
    _keepalive_options[_TCP_KEEPCNT] = 3


class Ellipsis:
    pass


class RedisError(RuntimeError):
    pass


ellipsis: Final[Ellipsis] = Ellipsis()
log = BraceStyleAdapter(logging.getLogger(__spec__.name))

BUF_SIZE: Final[int] = 1 * 1024 * 1024


class RedisClient(aobject):
    reader: asyncio.StreamReader
    writer: asyncio.StreamWriter

    verbose: bool

    _prev_buf: bytes

    def __init__(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        verbose: bool = False,
    ) -> None:
        self.reader = reader
        self.writer = writer

        self.verbose = verbose

    async def execute(
        self,
        command: Sequence[str | int | float | bytes | memoryview],
        *,
        command_timeout: float | None = None,
    ) -> Any:
        return (
            await self._send(
                [command],
                command_timeout=command_timeout,
            )
        )[0]

    async def subscribe_reader(
        self,
    ) -> AsyncIterator[Any]:
        hiredis_reader = hiredis.Reader(notEnoughData=ellipsis)
        while True:
            data = await self.reader.read(BUF_SIZE)
            if not data:
                break
            hiredis_reader.feed(data)
            msg = hiredis_reader.gets()
            if msg == ellipsis and self.reader.at_eof():
                raise EOFError
            if msg is None:
                continue
            yield msg

    async def pipeline(
        self,
        commands: Sequence[Sequence[str | int | float | bytes | memoryview]],
        *,
        command_timeout: float | None = None,
        return_exception: bool = False,
    ) -> Any:
        return await self._send(
            commands,
            command_timeout=command_timeout,
            return_exception=return_exception,
        )

    async def _send(
        self,
        commands: Sequence[Sequence[str | int | float | bytes | memoryview]],
        *,
        command_timeout: float | None = None,
        return_exception: bool = False,
    ) -> list[Any]:
        """
        Executes a function that issues Redis commands or returns a pipeline/transaction of commands,
        with automatic retries upon temporary connection failures.

        Note that when retried, the given function may be executed *multiple* times, so the caller
        should take care of side-effects of it.
        """

        if command_timeout:
            _timeout_secs = command_timeout
        else:
            _timeout_secs = 10.0

        while True:
            try:
                hiredis_reader = hiredis.Reader(notEnoughData=ellipsis)
                _blobs = b""
                for command in commands:
                    request_blob = hiredis.pack_command(tuple(command))  # type: ignore[arg-type]
                    self.writer.write(request_blob)
                    _blobs += request_blob
                async with asyncio.timeout(_timeout_secs):
                    await self.writer.drain()

                if self.verbose:
                    log.debug("requests: {}", commands)
                    log.debug("raw request: {}", _blobs)

                results = []
                first_result = ellipsis

                _buf = b""

                met_unexpected_eof = False

                try:
                    while first_result == ellipsis:
                        async with asyncio.timeout(_timeout_secs):
                            buf = await self.reader.read(n=BUF_SIZE)
                        _buf += buf
                        hiredis_reader.feed(buf)
                        first_result = hiredis_reader.gets()
                        if first_result == ellipsis and self.reader.at_eof():
                            met_unexpected_eof = True
                            break
                except hiredis.ProtocolError:
                    raise

                if met_unexpected_eof:
                    log.warning("Met unexpected EOF")
                    raise EOFError

                if self.verbose:
                    log.debug("raw response: {}", _buf)

                if isinstance(first_result, hiredis.HiredisError):
                    if not return_exception:
                        raise first_result
                results.append(first_result)

                while hiredis_reader.has_data():
                    next_result = hiredis_reader.gets()
                    if isinstance(next_result, hiredis.HiredisError):
                        if not return_exception:
                            raise next_result
                    results.append(next_result)

                try:
                    if len(results) != len(commands):
                        log.warning("requests: {}", commands)
                        log.warning("responses: {}", results)
                        log.warning("raw request: {}", _blobs)
                        log.warning("raw response: {}", _buf)
                        log.warning("previous raw response: {}", self._prev_buf)
                        if _buf.startswith(self._prev_buf):
                            log.warning(
                                "new response contains data of previous response, without it the response should formed like:"
                            )
                            log.warning("{}", _buf[len(self._prev_buf) :])
                        raise RedisError("Response count does not match with number of requests!")
                finally:
                    self._prev_buf = _buf

                if self.verbose:
                    for request, response in zip(commands, results, strict=True):
                        log.debug("{} -> {}", request, response)

                return results
            except asyncio.CancelledError:
                raise
            finally:
                await asyncio.sleep(0)


class RedisConnection(AbstractAsyncContextManager[RedisClient]):
    _redis_target: RedisTarget
    _db: int

    _socket_timeout: float | None
    _socket_connect_timeout: float | None
    _keepalive_options: dict[int, int]

    def __init__(
        self,
        redis_target: RedisTarget,
        *,
        db: int = 0,
        socket_timeout: float | None = 5.0,
        socket_connect_timeout: float | None = 2.0,
        keepalive_options: dict[int, int] = _keepalive_options,
    ) -> None:
        self._redis_target = redis_target
        self._db = db

        self._socket_timeout = socket_timeout
        self._socket_connect_timeout = socket_connect_timeout
        self._keepalive_options = keepalive_options

    async def connect(self) -> RedisClient:
        if self._redis_target.sentinel:
            raise RuntimeError("Redis with sentinel not supported for this library")

        redis_url = self._redis_target.addr
        if redis_url is None:
            raise ValueError("Redis URL is not configured")

        host = str(redis_url[0])
        port = redis_url[1]
        password = self._redis_target.password

        async with asyncio.timeout(self._socket_connect_timeout):
            reader, writer = await asyncio.open_connection(host, port)
            sock: socket.socket = writer.get_extra_info("socket")
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True)

            for opt, val in self._keepalive_options.items():
                sock.setsockopt(socket.IPPROTO_TCP, opt, val)

        self.writer = writer
        self.reader = reader

        client = RedisClient(reader, writer)

        if password:
            await client.execute(
                ["AUTH", password],
                command_timeout=self._socket_timeout,
            )

        await client.execute(
            ["HELLO", "3"],
            command_timeout=self._socket_timeout,
        )
        await client.execute(
            ["SELECT", self._db],
            command_timeout=self._socket_timeout,
        )

        return client

    async def __aenter__(self) -> RedisClient:
        return await self.connect()

    async def disconnect(self) -> None:
        if self.writer:
            try:
                self.writer.close()
            except RuntimeError as e:
                if str(e) == "Event loop is closed":
                    pass  # there's no more room we can do anything
                raise e

    async def __aexit__(self, *exc_info: Any) -> None:
        return await self.disconnect()
