import typing as t

from sqlalchemy import ClauseElement, Connection, Engine
from sqlalchemy import exc as sa_exc
from sqlalchemy.orm import Session as SQLAlchemySession

from ash_dal.database.utils import is_read_query

USE_SLAVE_HINT = "use_slave"
USE_MASTER_HINT = "use_master"


class Session(SQLAlchemySession):
    def get_bind(
        self,
        mapper: t.Any | None = None,  # pyright: ignore [reportMissingParameterType]
        *,
        clause: ClauseElement | None = None,
        bind: Engine | Connection | None = None,
        _sa_skip_events: bool | None = None,
        _sa_skip_for_implicit_returning: bool = False,
        **kw: t.Any,
    ) -> Engine | Connection:
        """
        Resolve the engine used for the current statement.

        1. Try SQLAlchemy default bind resolution first. This preserves SQLAlchemy-native binding behavior (explicit
           binds, mapper/table binds, and internal SQLAlchemy routing) whenever those are configured. However, normal
           behavior for ash-dal sessions is to not have a direct bind, so we fall back to master/slave routing.
        2. If a query in this session previously routed to master, subsequent reads will also route to master to
           preserve read-after-write consistency.
        3. Next we check for execution hints:
           - ``use_slave=True`` routes to the slave for this statement
           - ``use_master=True`` routes to the master for this statement
           If both hints are enabled, a ``ValueError`` is raised.
        4. Otherwise we route to the slave for reads and the master for writes.
        """
        try:
            return super().get_bind(
                mapper=mapper,
                clause=clause,
                bind=bind,
                _sa_skip_events=_sa_skip_events,
                _sa_skip_for_implicit_returning=_sa_skip_for_implicit_returning,
                **kw,
            )
        except sa_exc.UnboundExecutionError:
            # Route reads to replica by default, but preserve read-your-writes
            # consistency by sticking to master after a write in this session.
            master_engine = self.info.get("master")
            slave_engine = self.info.get("slave")
            execution_options = getattr(clause, "_execution_options", {}) if clause is not None else {}
            use_slave = bool(execution_options.get(USE_SLAVE_HINT))
            use_master = bool(execution_options.get(USE_MASTER_HINT))

            if use_slave and use_master:
                msg = "Execution hints use_slave and use_master cannot both be enabled."
                raise ValueError(msg) from None

            is_select_clause = is_read_query(clause) or use_slave
            if clause is not None and not is_select_clause and master_engine:
                self.info["force_master_reads"] = True
                return master_engine
            if is_select_clause and slave_engine and not self.info.get("force_master_reads") and not use_master:
                return slave_engine
            if master_engine:
                return master_engine
            raise
