"""
RoutingPrismaWrapper: routes Prisma reads to a read-replica client and writes
to a writer client. Used when DATABASE_URL_READ_REPLICA is configured;
otherwise PrismaClient uses the writer-only PrismaWrapper directly.
"""

import os
from collections.abc import Callable
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Final

from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.prisma_client import PrismaWrapper

if TYPE_CHECKING:
    from prisma.types import HttpConfig

# Per-model action methods that read from the database. These are routed to
# the read replica when one is configured.
_MODEL_READ_METHODS: Final = frozenset(
    {
        "find_first",
        "find_first_or_raise",
        "find_many",
        "find_unique",
        "find_unique_or_raise",
        "count",
        "group_by",
        "query_first",
        "query_raw",
    }
)

# Top-level Prisma client methods that read from the database.
_TOP_LEVEL_READ_METHODS: Final = frozenset({"query_first", "query_raw"})


class _RoutedActions:
    """Per-model accessor that sends reads to the reader and writes to the writer.

    `should_use_reader` is consulted on every read dispatch so a mid-call flip
    of the routing wrapper's reader-availability flag (e.g. after the reader
    fails a recreate) is observed without re-fetching the actions accessor.
    """

    __slots__ = ("_reader_actions", "_should_use_reader", "_writer_actions")

    def __init__(
        self,
        writer_actions: object,
        reader_actions: object,
        should_use_reader: Callable[[], bool],
    ):
        self._writer_actions = writer_actions
        self._reader_actions = reader_actions
        self._should_use_reader = should_use_reader

    def __getattr__(self, name: str) -> object:
        if name in _MODEL_READ_METHODS and self._should_use_reader():
            return getattr(self._reader_actions, name)
        return getattr(self._writer_actions, name)


class WriterPinnedClient:
    """PrismaClient-shaped view whose `.db` resolves to the writer while it is available.

    Read-after-write paths (e.g. the model reconcile a /model/new triggers to
    verify its own just-committed row) must not read through a lagging read
    replica: the row is not replayed there yet, so the reconcile concludes the
    write is missing and fails the request even though it is durable (#38556).

    While the writer is degraded (`writer_unavailable`), the pin yields to the
    routed wrapper so reconcile reads keep working from the replica: a proxy
    that starts during a primary outage must still load DB-backed models, and
    no read-after-write hazard exists then because writes are failing anyway.
    """

    __slots__ = ("db",)

    def __init__(self, db: "PrismaWrapper | RoutingPrismaWrapper") -> None:
        self.db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) and not db.writer_unavailable else db


class RoutingPrismaWrapper:
    """
    Routes Prisma operations between a writer and a reader Prisma client.

    Reads (find_*, count, group_by, query_raw, query_first) go to the reader;
    everything else (writes, transactions, raw execute) goes to the writer.
    Lifecycle methods (connect, disconnect, IAM token refresh) act on both
    clients so callers do not need to know about the split. When
    IAM_TOKEN_DB_AUTH is enabled, both writer and reader refresh their tokens
    independently on their own ~12-minute cadence.

    Reader degradation: a reader-side failure (failed connect, failed
    recreate) is non-fatal — the wrapper sets `_reader_unavailable=True`, logs
    a warning, and routes subsequent reads to the writer. The next successful
    `connect()` or `recreate_prisma_client()` clears the flag. This keeps the
    proxy serving traffic during transient reader outages instead of failing
    startup or returning errors for read-heavy endpoints.

    Writer degradation: a writer-side `connect()` failure while the reader
    connects is likewise non-fatal — the wrapper sets
    `_writer_unavailable=True`, logs a warning, and keeps serving reads from
    the reader (key lookups, DB-stored model loads) so a proxy that starts
    during a primary outage still serves inference from the replica. Writes
    fail at call time until the writer recovers; the PrismaClient DB health
    watchdog polls `writer_unavailable` and drives the writer reconnect,
    which clears the flag via `recreate_prisma_client`. Only when BOTH sides
    fail to connect does `connect()` raise (full DB outage — the existing
    `allow_requests_on_db_unavailable` startup handling applies).
    """

    def __init__(self, writer: PrismaWrapper, reader: PrismaWrapper):
        self._writer = writer
        self._reader = reader
        # When True, reads fall back to the writer. Flipped on by reader
        # connect/recreate failures and flipped off on the next reader recovery.
        self._reader_unavailable: bool = False
        self._writer_unavailable: bool = False

    @property
    def writer(self) -> PrismaWrapper:
        return self._writer

    @property
    def reader(self) -> PrismaWrapper:
        return self._reader

    @property
    def read_target(self) -> PrismaWrapper:
        """The wrapper `_TOP_LEVEL_READ_METHODS` dispatch to right now.

        Callers that need to reason about the engine a read actually ran on
        (e.g. recovering from prepared statements that went stale on it) must
        consult this rather than `writer`, and `__getattr__` routes through it
        so the two cannot drift apart.
        """
        return self._writer if self._reader_unavailable else self._reader

    @property
    def reader_unavailable(self) -> bool:
        return self._reader_unavailable

    @property
    def writer_unavailable(self) -> bool:
        return self._writer_unavailable

    def mark_writer_recovered(self) -> None:
        """Clear the degraded-writer flag after an external health probe proved
        the writer reachable. Needed when recovery happens without
        `recreate_prisma_client` (e.g. an IAM token refresh already recreated
        the writer engine), which is otherwise the only runtime path that
        clears the flag — without this, the watchdog would keep firing
        reconnect attempts against an already-healthy writer."""
        self._writer_unavailable = False

    def _should_use_reader(self) -> bool:
        return not self._reader_unavailable

    @staticmethod
    async def _try_connect(client: PrismaWrapper, timeout: int | timedelta | None = None) -> Exception | None:
        if client.is_connected() is True:
            return None
        try:
            await client.connect(timeout)
            return None
        except Exception as e:
            return e

    async def connect(self, timeout: int | timedelta | None = None) -> None:
        writer_error: Final = await self._try_connect(self._writer, timeout)
        if writer_error is None:
            self._writer_unavailable = False
            verbose_proxy_logger.info("[writer] DB connected")
        reader_error: Final = await self._try_connect(self._reader, timeout)
        if reader_error is None:
            self._reader_unavailable = False
            verbose_proxy_logger.info("[reader] DB connected")
        if writer_error is None and reader_error is None:
            return
        if writer_error is not None and reader_error is not None:
            raise writer_error
        if reader_error is not None:
            # Degrade gracefully: the proxy keeps serving traffic with reads
            # routed to the writer until the reader endpoint is reachable.
            # Aborting startup here would tie proxy availability to an
            # opt-in, best-effort reader endpoint.
            self._reader_unavailable = True
            verbose_proxy_logger.warning(
                "Failed to connect to read replica DB: %s. "
                "Falling back to the writer for reads until the reader is reachable.",
                reader_error,
            )
            return
        self._writer_unavailable = True
        verbose_proxy_logger.warning(
            "Failed to connect to primary (writer) DB: %s. "
            "Serving reads from the read replica; writes will fail until the writer recovers.",
            writer_error,
        )

    async def disconnect(self, timeout: float | timedelta | None = None) -> None:
        first_error: BaseException | None = None
        for client in (self._writer, self._reader):
            try:
                await client.disconnect(timeout)
            except Exception as e:
                if first_error is None:
                    first_error = e
                verbose_proxy_logger.warning("Error disconnecting Prisma client: %s", e)
        if first_error is not None:
            raise first_error

    def is_connected(self) -> bool:
        # Reflects writer health only. The reader is best-effort; its
        # availability is tracked via `_reader_unavailable` and a degraded
        # reader must NOT cause a writer reconnect (would loop indefinitely
        # since recreate_prisma_client only fixes writer-side problems).
        return bool(self._writer.is_connected())

    async def start_token_refresh_task(self) -> None:
        await self._writer.start_token_refresh_task()
        await self._reader.start_token_refresh_task()

    async def stop_token_refresh_task(self) -> None:
        await self._writer.stop_token_refresh_task()
        await self._reader.stop_token_refresh_task()

    async def recreate_prisma_client(
        self,
        new_db_url: str,
        http_client: "HttpConfig | None" = None,
        *,
        expected_generation: int | None = None,
    ) -> bool:
        """Recreate both writer and reader Prisma clients.

        The writer reconnect path in PrismaClient calls
        `self.db.recreate_prisma_client(...)`. Without this method, a DB-wide
        connectivity event would only re-create the writer; the reader engine
        would stay broken and every routed read would fail. We always recreate
        the writer first (its URL is the one passed in), then best-effort
        recreate the reader. A reader failure flips `_reader_unavailable=True`
        so reads transparently fall through to the writer.

        `expected_generation` is forwarded to the writer's optimistic-lock
        guard. If the writer recreate is skipped (another path already replaced
        the engine — issue #29176), we skip the reader too rather than churning
        it needlessly, and return ``False``.
        """
        writer_recreated: Final = await self._writer.recreate_prisma_client(
            new_db_url,
            http_client=http_client,
            expected_generation=expected_generation,
        )
        if not writer_recreated:
            return False
        self._writer_unavailable = False
        try:
            await self._recreate_reader(http_client=http_client)
            self._reader_unavailable = False
        except Exception as e:
            self._reader_unavailable = True
            verbose_proxy_logger.warning(
                "Failed to recreate reader Prisma client: %s. "
                "Reads will fall back to the writer until the reader recovers.",
                e,
            )
        return True

    async def _recreate_reader(self, http_client: "HttpConfig | None" = None) -> None:
        """Resolve the reader URL and recreate its Prisma client.

        Token-authenticated readers regenerate their token (host/port/user came
        from the parsed reader URL at construction time). Password-authenticated
        readers reuse the URL stored in `DATABASE_URL_READ_REPLICA`.
        """
        if self._reader.iam_token_db_auth:
            new_reader_url: Final = self._reader.get_rds_iam_token()
            if not new_reader_url:
                raise RuntimeError(f"Failed to generate fresh {self._reader.token_label} for read replica")
            await self._reader.recreate_prisma_client(new_reader_url, http_client=http_client)
            return
        reader_url: Final = os.getenv("DATABASE_URL_READ_REPLICA", "")
        if not reader_url:
            raise RuntimeError("DATABASE_URL_READ_REPLICA not set; cannot recreate read replica client")
        await self._reader.recreate_prisma_client(reader_url, http_client=http_client)

    def __getattr__(self, name: str) -> Any:
        if name in _TOP_LEVEL_READ_METHODS:
            return getattr(self.read_target, name)
        writer_attr: Final[object] = getattr(self._writer, name)
        # Per-model action accessors are non-callable instances that expose
        # both `find_many` and `create`. Methods like execute_raw / batch_ /
        # tx are callables and stay on the writer untouched.
        if not callable(writer_attr) and hasattr(writer_attr, "find_many") and hasattr(writer_attr, "create"):
            try:
                reader_attr: Final[object] = getattr(self._reader, name)
            except AttributeError:
                return writer_attr
            return _RoutedActions(writer_attr, reader_attr, self._should_use_reader)
        return writer_attr
