"""高级查询 API

提供复杂的查询操作：
- MongoDB 风格的查询操作符
- 聚合查询
- 排序和分页
- 流式迭代（内存优化）
"""

from collections.abc import AsyncIterator
import time
from typing import Any, Generic, TypeVar, Self

from sqlalchemy import and_, asc, desc, func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

# 导入 CRUD 辅助函数以避免重复定义
from src.kernel.db.api.crud import (
    _dict_to_model,
    _get_session_ctx,
    _model_to_dict,
    _normalize_cache_value,
)
from src.kernel.db.core.cache import (
    LIST_RESULT_CACHE_LIMIT,
    get_cache_namespace,
    get_db_read_cache,
)
from src.kernel.db.telemetry import record_db_operation_event
from src.kernel.logger import get_logger

logger = get_logger("database.query", display="DB Query")

T = TypeVar("T", bound=Any)


def _build_compiled_query_cache_key(
    session: AsyncSession,
    stmt: Any,
    *,
    operation: str,
    **parts: Any,
) -> tuple[Any, ...]:
    bind = session.get_bind()
    dialect = bind.dialect if bind is not None else None
    compiled = stmt.compile(
        dialect=dialect,
        compile_kwargs={"render_postcompile": True},
    )
    normalized_params = tuple(
        sorted(
            (str(key), _normalize_cache_value(value))
            for key, value in compiled.params.items()
        )
    )
    normalized_parts = tuple(
        sorted((key, _normalize_cache_value(value)) for key, value in parts.items())
    )
    return (
        operation,
        str(compiled),
        normalized_params,
        normalized_parts,
    )


async def _record_query_operation(
    *,
    operation: str,
    model_name: str,
    started_at: float,
    session_factory: async_sessionmaker | None,
    success: bool,
    result_count: int | None = None,
    error: Exception | None = None,
    attributes: dict[str, Any] | None = None,
) -> None:
    """记录 QueryBuilder/AggregateQuery 遥测。"""
    await record_db_operation_event(
        operation=operation,
        model_name=model_name,
        duration_ms=(time.perf_counter() - started_at) * 1000,
        custom_session_factory=session_factory is not None,
        success=success,
        result_count=result_count,
        error=error,
        attributes=attributes,
    )


class QueryBuilder(Generic[T]):
    """查询构建器

    支持链式调用，构建复杂查询
    """

    def __init__(self, model: type[T], *, session_factory: async_sessionmaker | None = None):
        """初始化查询构建器

        Args:
            model: SQLAlchemy 模型类
            session_factory: 可选的自定义异步会话工厂。为 None 时使用全局主数据库会话；
                传入自定义工厂时使用该工厂（适用于插件独立数据库等场景）。
        """
        self.model = model
        self.model_name = model.__tablename__
        self._stmt = select(model)
        self._session_factory = session_factory
        self._cache = get_db_read_cache()
        self._cache_namespace = get_cache_namespace(
            model,
            session_factory=session_factory,
        )

    def filter(self, **conditions: Any) -> Self:
        """添加过滤条件

        支持的操作符：
        - 直接相等: field=value
        - 大于: field__gt=value
        - 小于: field__lt=value
        - 大于等于: field__gte=value
        - 小于等于: field__lte=value
        - 不等于: field__ne=value
        - 包含: field__in=[values]
        - 不包含: field__nin=[values]
        - 模糊匹配: field__like='%pattern%'
        - 为空: field__isnull=True

        Args:
            **conditions: 过滤条件

        Returns:
            self，支持链式调用
        """
        for key, value in conditions.items():
            # 解析字段和操作符
            if "__" in key:
                field_name, operator = key.rsplit("__", 1)
            else:
                field_name, operator = key, "eq"

            if not hasattr(self.model, field_name):
                logger.warning(f"模型 {self.model_name} 没有字段 {field_name}")
                continue

            field = getattr(self.model, field_name)

            # 应用操作符
            if operator == "eq":
                self._stmt = self._stmt.where(field == value)
            elif operator == "gt":
                self._stmt = self._stmt.where(field > value)
            elif operator == "lt":
                self._stmt = self._stmt.where(field < value)
            elif operator == "gte":
                self._stmt = self._stmt.where(field >= value)
            elif operator == "lte":
                self._stmt = self._stmt.where(field <= value)
            elif operator == "ne":
                self._stmt = self._stmt.where(field != value)
            elif operator == "in":
                self._stmt = self._stmt.where(field.in_(value))
            elif operator == "nin":
                self._stmt = self._stmt.where(~field.in_(value))
            elif operator == "like":
                self._stmt = self._stmt.where(field.like(value))
            elif operator == "isnull":
                if value:
                    self._stmt = self._stmt.where(field.is_(None))
                else:
                    self._stmt = self._stmt.where(field.isnot(None))
            else:
                logger.warning(f"未知操作符: {operator}")

        return self

    def filter_or(self, **conditions: Any) -> Self:
        """添加 OR 过滤条件

        Args:
            **conditions: OR 条件

        Returns:
            self，支持链式调用
        """
        or_conditions = []
        for key, value in conditions.items():
            if hasattr(self.model, key):
                field = getattr(self.model, key)
                or_conditions.append(field == value)

        if or_conditions:
            self._stmt = self._stmt.where(or_(*or_conditions))

        return self

    def order_by(self, *fields: str) -> Self:
        """添加排序

        Args:
            *fields: 排序字段，'-' 前缀表示降序

        Returns:
            self，支持链式调用
        """
        for field_name in fields:
            if field_name.startswith("-"):
                field_name = field_name[1:]
                if hasattr(self.model, field_name):
                    self._stmt = self._stmt.order_by(desc(getattr(self.model, field_name)))
            else:
                if hasattr(self.model, field_name):
                    self._stmt = self._stmt.order_by(asc(getattr(self.model, field_name)))

        return self

    def limit(self, limit: int) -> Self:
        """限制结果数量

        Args:
            limit: 最大数量

        Returns:
            self，支持链式调用
        """
        self._stmt = self._stmt.limit(limit)
        return self

    def offset(self, offset: int) -> Self:
        """跳过指定数量

        Args:
            offset: 跳过数量

        Returns:
            self，支持链式调用
        """
        self._stmt = self._stmt.offset(offset)
        return self

    async def iter_batches(
        self,
        batch_size: int = 1000,
        *,
        as_dict: bool = True,
    ) -> AsyncIterator[list[T] | list[dict[str, Any]]]:
        """分批迭代获取结果（内存优化）

        使用 LIMIT/OFFSET 分页策略，避免一次性加载全部数据到内存。
        适用于大数据量的统计、导出等场景。

        Args:
            batch_size: 每批获取的记录数，默认 1000
            as_dict: 为 True 时返回字典格式

        Yields:
            每批的模型实例列表或字典列表

        Example:
            async for batch in query_builder.iter_batches(batch_size=500):
                for record in batch:
                    process(record)
        """
        offset = 0
        started_at = time.perf_counter()

        try:
            while True:
                paginated_stmt = self._stmt.offset(offset).limit(batch_size)

                async with _get_session_ctx(self._session_factory) as session:
                    result = await session.execute(paginated_stmt)
                    instances = result.scalars().all()

                    if not instances:
                        break

                    instances_dicts = [_model_to_dict(inst) for inst in instances]

                if as_dict:
                    yield instances_dicts
                else:
                    yield [_dict_to_model(self.model, row) for row in instances_dicts]

                if len(instances) < batch_size:
                    break

                offset += batch_size

            await _record_query_operation(
                operation="iter_batches",
                model_name=self.model_name,
                started_at=started_at,
                session_factory=self._session_factory,
                success=True,
                attributes={"batch_size": batch_size},
            )
        except Exception as exc:
            await _record_query_operation(
                operation="iter_batches",
                model_name=self.model_name,
                started_at=started_at,
                session_factory=self._session_factory,
                success=False,
                error=exc,
                attributes={"batch_size": batch_size},
            )
            raise

    async def iter_all(
        self,
        batch_size: int = 1000,
        *,
        as_dict: bool = True,
    ) -> AsyncIterator[T | dict[str, Any]]:
        """逐条迭代所有结果（内存优化）

        内部使用分批获取，但对外提供逐条迭代的接口。
        适用于需要逐条处理但数据量很大的场景。

        Args:
            batch_size: 内部分批大小，默认 1000
            as_dict: 为 True 时返回字典格式

        Yields:
            单个模型实例或字典

        Example:
            async for record in query_builder.iter_all():
                process(record)
        """
        async for batch in self.iter_batches(batch_size=batch_size, as_dict=as_dict):
            for item in batch:
                yield item

    async def all(self, *, as_dict: bool = False) -> list[T] | list[dict[str, Any]]:
        """获取所有结果

        Args:
            as_dict: 为 True 时返回字典格式

        Returns:
            模型实例列表或字典列表
        """
        started_at = time.perf_counter()
        try:
            async with _get_session_ctx(self._session_factory) as session:
                cache_key = _build_compiled_query_cache_key(
                    session,
                    self._stmt,
                    operation="all",
                    as_dict=as_dict,
                )
                hit, cached = self._cache.get(self._cache_namespace, cache_key)
                if hit:
                    if as_dict:
                        return cached or []
                    return [_dict_to_model(self.model, row) for row in (cached or [])]

                result = await session.execute(self._stmt)
                instances = list(result.scalars().all())
                instances_dicts = [_model_to_dict(inst) for inst in instances]

                await _record_query_operation(
                    operation="all",
                    model_name=self.model_name,
                    started_at=started_at,
                    session_factory=self._session_factory,
                    success=True,
                    result_count=len(instances_dicts),
                )
                if len(instances_dicts) <= LIST_RESULT_CACHE_LIMIT:
                    self._cache.set(self._cache_namespace, cache_key, instances_dicts)

                if as_dict:
                    return instances_dicts
                return [_dict_to_model(self.model, row) for row in instances_dicts]
        except Exception as exc:
            await _record_query_operation(
                operation="all",
                model_name=self.model_name,
                started_at=started_at,
                session_factory=self._session_factory,
                success=False,
                error=exc,
            )
            raise

    async def first(self, *, as_dict: bool = False) -> T | dict[str, Any] | None:
        """获取第一条结果

        Args:
            as_dict: 为 True 时返回字典格式

        Returns:
            模型实例、字典或 None
        """
        started_at = time.perf_counter()
        try:
            async with _get_session_ctx(self._session_factory) as session:
                cache_key = _build_compiled_query_cache_key(
                    session,
                    self._stmt,
                    operation="first",
                    as_dict=as_dict,
                )
                hit, cached = self._cache.get(self._cache_namespace, cache_key)
                if hit:
                    if cached is None:
                        return None
                    if as_dict:
                        return cached
                    return _dict_to_model(self.model, cached)

                result = await session.execute(self._stmt)
                instance = result.scalars().first()

                if instance is not None:
                    instance_dict = _model_to_dict(instance)
                    await _record_query_operation(
                        operation="first",
                        model_name=self.model_name,
                        started_at=started_at,
                        session_factory=self._session_factory,
                        success=True,
                        result_count=1,
                    )
                    self._cache.set(self._cache_namespace, cache_key, instance_dict)
                    if as_dict:
                        return instance_dict
                    return _dict_to_model(self.model, instance_dict)

                await _record_query_operation(
                    operation="first",
                    model_name=self.model_name,
                    started_at=started_at,
                    session_factory=self._session_factory,
                    success=True,
                    result_count=0,
                )
                return None
        except Exception as exc:
            await _record_query_operation(
                operation="first",
                model_name=self.model_name,
                started_at=started_at,
                session_factory=self._session_factory,
                success=False,
                error=exc,
            )
            raise

    async def count(self) -> int:
        """统计数量

        Returns:
            记录数量
        """
        count_stmt = select(func.count()).select_from(self._stmt.subquery())
        started_at = time.perf_counter()

        try:
            async with _get_session_ctx(self._session_factory) as session:
                cache_key = _build_compiled_query_cache_key(
                    session,
                    count_stmt,
                    operation="count",
                )
                hit, cached = self._cache.get(self._cache_namespace, cache_key)
                if hit:
                    return int(cached or 0)

                result = await session.execute(count_stmt)
                count_value = result.scalar() or 0
                await _record_query_operation(
                    operation="count",
                    model_name=self.model_name,
                    started_at=started_at,
                    session_factory=self._session_factory,
                    success=True,
                    result_count=int(count_value),
                )
                self._cache.set(self._cache_namespace, cache_key, int(count_value))
                return count_value
        except Exception as exc:
            await _record_query_operation(
                operation="count",
                model_name=self.model_name,
                started_at=started_at,
                session_factory=self._session_factory,
                success=False,
                error=exc,
            )
            raise

    async def exists(self) -> bool:
        """检查是否存在

        Returns:
            是否存在记录
        """
        count = await self.count()
        return count > 0

    async def paginate(
        self,
        page: int = 1,
        page_size: int = 20,
    ) -> tuple[list[T], int]:
        """分页查询

        Args:
            page: 页码（从 1 开始）
            page_size: 每页数量

        Returns:
            (结果列表, 总数量)
        """
        offset = (page - 1) * page_size
        started_at = time.perf_counter()

        try:
            total = await self.count()
            paginated_stmt = self._stmt.offset(offset).limit(page_size)

            async with _get_session_ctx(self._session_factory) as session:
                cache_key = _build_compiled_query_cache_key(
                    session,
                    paginated_stmt,
                    operation="paginate",
                    page=page,
                    page_size=page_size,
                )
                hit, cached = self._cache.get(self._cache_namespace, cache_key)
                if hit:
                    items = [_dict_to_model(self.model, row) for row in (cached or [])]
                    return items, total  # type: ignore[return-value]

                result = await session.execute(paginated_stmt)
                instances = list(result.scalars().all())
                instances_dicts = [_model_to_dict(inst) for inst in instances]
                items = [_dict_to_model(self.model, row) for row in instances_dicts]
                if len(instances_dicts) <= LIST_RESULT_CACHE_LIMIT:
                    self._cache.set(self._cache_namespace, cache_key, instances_dicts)

            await _record_query_operation(
                operation="paginate",
                model_name=self.model_name,
                started_at=started_at,
                session_factory=self._session_factory,
                success=True,
                result_count=len(items),
                attributes={"page": page, "page_size": page_size, "total": total},
            )
            return items, total  # type: ignore
        except Exception as exc:
            await _record_query_operation(
                operation="paginate",
                model_name=self.model_name,
                started_at=started_at,
                session_factory=self._session_factory,
                success=False,
                error=exc,
                attributes={"page": page, "page_size": page_size},
            )
            raise


class AggregateQuery:
    """聚合查询

    提供聚合操作如 sum、avg、max、min 等
    """

    def __init__(self, model: type[T], *, session_factory: async_sessionmaker | None = None):
        """初始化聚合查询

        Args:
            model: SQLAlchemy 模型类
            session_factory: 可选的自定义异步会话工厂。为 None 时使用全局主数据库会话；
                传入自定义工厂时使用该工厂（适用于插件独立数据库等场景）。
        """
        self.model = model
        self.model_name = model.__tablename__
        self._conditions = []
        self._session_factory = session_factory

    def filter(self, **conditions: Any) -> Self:
        """添加过滤条件

        Args:
            **conditions: 过滤条件

        Returns:
            self，支持链式调用
        """
        for key, value in conditions.items():
            if hasattr(self.model, key):
                field = getattr(self.model, key)
                self._conditions.append(field == value)
        return self

    async def sum(self, field: str) -> float:
        """求和

        Args:
            field: 字段名

        Returns:
            总和
        """
        if not hasattr(self.model, field):
            raise ValueError(f"字段 {field} 不存在")

        async with _get_session_ctx(self._session_factory) as session:
            stmt = select(func.sum(getattr(self.model, field)))

            if self._conditions:
                stmt = stmt.where(and_(*self._conditions))

            result = await session.execute(stmt)
            return result.scalar() or 0

    async def avg(self, field: str) -> float:
        """求平均值

        Args:
            field: 字段名

        Returns:
            平均值
        """
        if not hasattr(self.model, field):
            raise ValueError(f"字段 {field} 不存在")

        async with _get_session_ctx(self._session_factory) as session:
            stmt = select(func.avg(getattr(self.model, field)))

            if self._conditions:
                stmt = stmt.where(and_(*self._conditions))

            result = await session.execute(stmt)
            return result.scalar() or 0

    async def max(self, field: str) -> Any:
        """求最大值

        Args:
            field: 字段名

        Returns:
            最大值
        """
        if not hasattr(self.model, field):
            raise ValueError(f"字段 {field} 不存在")

        async with _get_session_ctx(self._session_factory) as session:
            stmt = select(func.max(getattr(self.model, field)))

            if self._conditions:
                stmt = stmt.where(and_(*self._conditions))

            result = await session.execute(stmt)
            return result.scalar()

    async def min(self, field: str) -> Any:
        """求最小值

        Args:
            field: 字段名

        Returns:
            最小值
        """
        if not hasattr(self.model, field):
            raise ValueError(f"字段 {field} 不存在")

        async with _get_session_ctx(self._session_factory) as session:
            stmt = select(func.min(getattr(self.model, field)))

            if self._conditions:
                stmt = stmt.where(and_(*self._conditions))

            result = await session.execute(stmt)
            return result.scalar()

    async def group_by_count(
        self,
        *fields: str,
    ) -> list[tuple[Any, ...]]:
        """分组统计

        Args:
            *fields: 分组字段

        Returns:
            [(分组值1, 分组值2, ..., 数量), ...] 
        """
        if not fields:
            raise ValueError("至少需要一个分组字段")

        group_columns = [
            getattr(self.model, field_name)
            for field_name in fields
            if hasattr(self.model, field_name)
        ]

        if not group_columns:
            return []

        async with _get_session_ctx(self._session_factory) as session:
            stmt = select(*group_columns, func.count(self.model.id))

            if self._conditions:
                stmt = stmt.where(and_(*self._conditions))

            stmt = stmt.group_by(*group_columns)

            result = await session.execute(stmt)
            return [tuple(row) for row in result.all()]
