"""
SQLAlchemy ORM 自动租户过滤集成

在不改动业务代码的前提下，自动为包含 tenant_id 字段的模型追加租户过滤。
"""

from collections.abc import Callable
from typing import Any, Type

try:
    from sqlalchemy import event
    from sqlalchemy.engine import Engine
    from sqlalchemy.orm import Session, sessionmaker
    from sqlalchemy.orm.decl_api import DeclarativeMeta
    from sqlalchemy.sql import Select
except ImportError:
    event = None
    Engine = None
    Session = None
    sessionmaker = None
    DeclarativeMeta = None
    Select = None

from ai4love_tools.tenant.context import get_tenant_id
from ai4love_tools.tenant.exceptions import TenantIdMissingError

# 模型元数据缓存：记录哪些模型包含 tenant_id 字段
_model_tenant_cache: dict[Type, bool] = {}


def _has_tenant_column(model_class: Type) -> bool:
    """
    检查模型类是否包含租户字段

    Args:
        model_class: SQLAlchemy 模型类

    Returns:
        如果模型包含 tenant_id 字段则返回 True
    """
    if model_class in _model_tenant_cache:
        return _model_tenant_cache[model_class]

    if not hasattr(model_class, "__table__"):
        _model_tenant_cache[model_class] = False
        return False

    table = model_class.__table__
    has_tenant = "tenant_id" in table.columns

    _model_tenant_cache[model_class] = has_tenant
    return has_tenant


def _should_apply_tenant_filter(
    model_class: Type | None, enabled_models: set[Type] | None = None
) -> bool:
    """
    判断是否应该应用租户过滤

    Args:
        model_class: 模型类
        enabled_models: 启用租户过滤的模型白名单，如果为 None 则对所有包含 tenant_id 的模型启用

    Returns:
        是否应该应用租户过滤
    """
    if not model_class:
        return False

    if enabled_models is not None:
        return model_class in enabled_models

    return _has_tenant_column(model_class)


def _add_tenant_filter_to_select(
    select_stmt: Select, model_class: Type, tenant_id: str | int
) -> Select:
    """
    为 SELECT 语句添加租户过滤条件

    Args:
        select_stmt: SQLAlchemy Select 语句
        model_class: 模型类
        tenant_id: 租户ID

    Returns:
        添加了租户条件的 Select 语句
    """
    if not _should_apply_tenant_filter(model_class):
        return select_stmt

    # 检查是否已有租户条件
    for criterion in select_stmt.whereclause.children if select_stmt.whereclause else []:
        if hasattr(criterion, "left") and hasattr(criterion.left, "key"):
            if criterion.left.key == "tenant_id":
                return select_stmt

    # 添加租户条件
    tenant_column = getattr(model_class, "tenant_id", None)
    if tenant_column is None:
        return select_stmt

    return select_stmt.where(tenant_column == tenant_id)


def _before_flush_add_tenant_id(session: Session, flush_context: Any, instances: Any) -> None:
    """
    在 flush 前自动为 INSERT 操作补充 tenant_id

    Args:
        session: SQLAlchemy Session
        flush_context: flush 上下文
        instances: 实例列表
    """
    tenant_id = get_tenant_id()
    if not tenant_id:
        return

    for instance in session.new:
        if not hasattr(instance, "__class__"):
            continue

        model_class = instance.__class__
        if not _should_apply_tenant_filter(model_class):
            continue

        if hasattr(instance, "tenant_id") and getattr(instance, "tenant_id", None) is None:
            setattr(instance, "tenant_id", tenant_id)


def enable_sqlalchemy_tenant_isolation(
    engine_or_session_factory: Engine | sessionmaker | Type[Session],
    *,
    tenant_column: str = "tenant_id",
    require_tenant: bool = False,
    enabled_models: set[Type] | None = None,
) -> None:
    """
    启用 SQLAlchemy ORM 自动租户过滤

    在应用启动阶段调用一次即可全局生效。

    Args:
        engine_or_session_factory: SQLAlchemy Engine、sessionmaker 或 Session 类
        tenant_column: 租户字段名，默认 "tenant_id"
        require_tenant: 是否强制要求上下文中有租户ID，默认 False
        enabled_models: 启用租户过滤的模型白名单，如果为 None 则对所有包含 tenant_id 的模型启用

    Example:
        from sqlalchemy import create_engine
        from ai4love_tools.tenant.sqlalchemy_integration import enable_sqlalchemy_tenant_isolation

        engine = create_engine("sqlite:///db.sqlite")
        enable_sqlalchemy_tenant_isolation(engine)
    """
    if event is None or Session is None:
        raise ImportError("SQLAlchemy 未安装，无法启用租户隔离功能")

    # 注册 before_flush 事件（用于 INSERT）
    # 检查是否已经注册过，避免重复注册
    try:
        listeners = event.contains(Session, "before_flush", _before_flush_add_tenant_id)
        if not listeners:
            event.listen(Session, "before_flush", _before_flush_add_tenant_id)
    except (AttributeError, TypeError):
        # 如果 event.contains 不可用，直接注册
        event.listen(Session, "before_flush", _before_flush_add_tenant_id)

    # 对于 SQLAlchemy 2.0+，使用 do_orm_execute 钩子
    try:
        from sqlalchemy.orm import ORMExecuteState

        @event.listens_for(Session, "do_orm_execute", once=False)
        def _do_orm_execute(state: ORMExecuteState) -> None:
            """
            在执行 ORM 查询时自动添加租户过滤
            """
            tenant_id = get_tenant_id()

            if require_tenant and not tenant_id:
                raise TenantIdMissingError("执行 ORM 查询需要租户ID，但上下文中未设置")

            if not tenant_id:
                return

            if not isinstance(state.statement, Select):
                return

            # 获取查询涉及的模型
            # 对于简单的单表查询，可以从 statement 中提取模型
            # 对于复杂的多表查询，可能需要更复杂的逻辑
            for entity in state.statement.column_descriptions:
                if "entity" in entity:
                    model_class = entity["entity"]
                    if _should_apply_tenant_filter(model_class, enabled_models):
                        state.statement = _add_tenant_filter_to_select(
                            state.statement, model_class, tenant_id
                        )
                        break

    except ImportError:
        # SQLAlchemy 1.x 版本，使用其他方式
        pass
