from datetime import datetime
from sqlalchemy import event, true
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.orm import Session, declarative_base, with_loader_criteria
from sqlalchemy.orm.attributes import flag_modified
from app.core.config import settings
from app.core.logging import get_logger
from typing import AsyncGenerator

from app.db.datascope import DataScopeHandle, SYSTEM_USER
from app.db.models.base_entity import BaseEntity

logger = get_logger(__name__)

# 创建数据库引擎
engine = create_async_engine(
    settings.database_url,
    echo=False,  # 关闭SQL调试日志以减少输出
    future=True,
    pool_pre_ping=True,
    pool_recycle=3600
)

# 创建会话工厂
AsyncSessionLocal = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False
)

@event.listens_for(Session, "do_orm_execute")
def _apply_data_scope(orm_execute_state):
    # only act on SELECT queries generated by ORM
    if not orm_execute_state.is_select:
        return

    allowed = DataScopeHandle.allowed_users()
    if not allowed or len(allowed) == 0:
        return

    # predicate builder: return None to skip for classes that opt-out
    def criteria_fn(cls):
        # skip if the mapped class explicitly disables data-scope
        if getattr(cls, "__ignore_data_scope__", False):
            return true()
        # some classes may not have created_by column; guard dynamically
        col = getattr(cls, "created_by", None)
        if col is None:
            return true()
        return col.in_(allowed)

    # apply loader-level criteria to all subclasses of BaseEntity
    orm_execute_state.statement = orm_execute_state.statement.options(
        with_loader_criteria(BaseEntity, criteria_fn, include_aliases=True)
    )

@event.listens_for(Session, "before_flush")
def _audit_before_flush(session, flush_context, instances):
    user = DataScopeHandle.get_user_info()
    effective_user = (user or "").strip() or SYSTEM_USER
    now = datetime.now()

    # new -> set created_* and updated_*
    for obj in list(session.new):
        if isinstance(obj, BaseEntity):
            if getattr(obj, "created_at", None) is None:
                obj.created_at = now
            if getattr(obj, "created_by", None) is None:
                obj.created_by = effective_user
            obj.updated_at = now
            obj.updated_by = effective_user
            # ensure SQLAlchemy sees changes
            try:
                flag_modified(obj, "created_by")
            except Exception:
                pass

    # dirty -> set updated_*
    for obj in list(session.dirty):
        if isinstance(obj, BaseEntity):
            obj.updated_at = now
            obj.updated_by = effective_user
            try:
                flag_modified(obj, "updated_by")
            except Exception:
                pass

async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """获取数据库会话"""
    async with AsyncSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close()