"""
数据库批量调度器
实现多个数据库请求的智能合并和批量处理，减少数据库连接竞争
"""

import asyncio
import time
from collections import defaultdict, deque
from collections.abc import Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any, TypeVar

from sqlalchemy import delete, insert, select, update

from src.common.database.sqlalchemy_database_api import get_db_session
from src.common.logger import get_logger

logger = get_logger("db_batch_scheduler")

T = TypeVar("T")


@dataclass
class BatchOperation:
    """批量操作基础类"""

    operation_type: str  # 'select', 'insert', 'update', 'delete'
    model_class: Any
    conditions: dict[str, Any]
    data: dict[str, Any] | None = None
    callback: Callable | None = None
    future: asyncio.Future | None = None
    timestamp: float = 0.0

    def __post_init__(self):
        if self.timestamp == 0.0:
            self.timestamp = time.time()


@dataclass
class BatchResult:
    """批量操作结果"""

    success: bool
    data: Any = None
    error: str | None = None


class DatabaseBatchScheduler:
    """数据库批量调度器"""

    def __init__(
        self,
        batch_size: int = 50,
        max_wait_time: float = 0.1,  # 100ms
        max_queue_size: int = 1000,
    ):
        self.batch_size = batch_size
        self.max_wait_time = max_wait_time
        self.max_queue_size = max_queue_size

        # 操作队列，按操作类型和模型分类
        self.operation_queues: dict[str, deque] = defaultdict(deque)

        # 调度控制
        self._scheduler_task: asyncio.Task | None = None
        self._is_running = False
        self._lock = asyncio.Lock()

        # 统计信息
        self.stats = {"total_operations": 0, "batched_operations": 0, "cache_hits": 0, "execution_time": 0.0}

        # 简单的结果缓存（用于频繁的查询）
        self._result_cache: dict[str, tuple[Any, float]] = {}
        self._cache_ttl = 5.0  # 5秒缓存

    async def start(self):
        """启动调度器"""
        if self._is_running:
            return

        self._is_running = True
        self._scheduler_task = asyncio.create_task(self._scheduler_loop())
        logger.info("数据库批量调度器已启动")

    async def stop(self):
        """停止调度器"""
        if not self._is_running:
            return

        self._is_running = False
        if self._scheduler_task:
            self._scheduler_task.cancel()
            try:
                await self._scheduler_task
            except asyncio.CancelledError:
                pass

        # 处理剩余的操作
        await self._flush_all_queues()
        logger.info("数据库批量调度器已停止")

    def _generate_cache_key(self, operation_type: str, model_class: Any, conditions: dict[str, Any]) -> str:
        """生成缓存键"""
        # 简单的缓存键生成，实际可以根据需要优化
        key_parts = [operation_type, model_class.__name__, str(sorted(conditions.items()))]
        return "|".join(key_parts)

    def _get_from_cache(self, cache_key: str) -> Any | None:
        """从缓存获取结果"""
        if cache_key in self._result_cache:
            result, timestamp = self._result_cache[cache_key]
            if time.time() - timestamp < self._cache_ttl:
                self.stats["cache_hits"] += 1
                return result
            else:
                # 清理过期缓存
                del self._result_cache[cache_key]
        return None

    def _set_cache(self, cache_key: str, result: Any):
        """设置缓存"""
        self._result_cache[cache_key] = (result, time.time())

    async def add_operation(self, operation: BatchOperation) -> asyncio.Future:
        """添加操作到队列"""
        # 检查是否可以立即返回缓存结果
        if operation.operation_type == "select":
            cache_key = self._generate_cache_key(operation.operation_type, operation.model_class, operation.conditions)
            cached_result = self._get_from_cache(cache_key)
            if cached_result is not None:
                if operation.callback:
                    operation.callback(cached_result)
                future = asyncio.get_event_loop().create_future()
                future.set_result(cached_result)
                return future

        # 创建future用于返回结果
        future = asyncio.get_event_loop().create_future()
        operation.future = future

        # 添加到队列
        queue_key = f"{operation.operation_type}_{operation.model_class.__name__}"

        async with self._lock:
            if len(self.operation_queues[queue_key]) >= self.max_queue_size:
                # 队列满了，直接执行
                await self._execute_operations([operation])
            else:
                self.operation_queues[queue_key].append(operation)
                self.stats["total_operations"] += 1

        return future

    async def _scheduler_loop(self):
        """调度器主循环"""
        while self._is_running:
            try:
                await asyncio.sleep(self.max_wait_time)
                await self._flush_all_queues()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"调度器循环异常: {e}", exc_info=True)

    async def _flush_all_queues(self):
        """刷新所有队列"""
        async with self._lock:
            if not any(self.operation_queues.values()):
                return

            # 复制队列内容，避免长时间占用锁
            queues_copy = {key: deque(operations) for key, operations in self.operation_queues.items()}
            # 清空原队列
            for queue in self.operation_queues.values():
                queue.clear()

        # 批量执行各队列的操作
        for operations in queues_copy.values():
            if operations:
                await self._execute_operations(list(operations))

    async def _execute_operations(self, operations: list[BatchOperation]):
        """执行批量操作"""
        if not operations:
            return

        start_time = time.time()

        try:
            # 按操作类型分组
            op_groups = defaultdict(list)
            for op in operations:
                op_groups[op.operation_type].append(op)

            # 为每种操作类型创建批量执行任务
            tasks = []
            for op_type, ops in op_groups.items():
                if op_type == "select":
                    tasks.append(self._execute_select_batch(ops))
                elif op_type == "insert":
                    tasks.append(self._execute_insert_batch(ops))
                elif op_type == "update":
                    tasks.append(self._execute_update_batch(ops))
                elif op_type == "delete":
                    tasks.append(self._execute_delete_batch(ops))

            # 并发执行所有操作
            results = await asyncio.gather(*tasks, return_exceptions=True)

            # 处理结果
            for i, result in enumerate(results):
                operation = operations[i]
                if isinstance(result, Exception):
                    if operation.future and not operation.future.done():
                        operation.future.set_exception(result)
                else:
                    if operation.callback:
                        try:
                            operation.callback(result)
                        except Exception as e:
                            logger.warning(f"操作回调执行失败: {e}")

                    if operation.future and not operation.future.done():
                        operation.future.set_result(result)

                    # 缓存查询结果
                    if operation.operation_type == "select":
                        cache_key = self._generate_cache_key(
                            operation.operation_type, operation.model_class, operation.conditions
                        )
                        self._set_cache(cache_key, result)

            self.stats["batched_operations"] += len(operations)

        except Exception as e:
            logger.error(f"批量操作执行失败: {e}", exc_info="")
            # 设置所有future的异常状态
            for operation in operations:
                if operation.future and not operation.future.done():
                    operation.future.set_exception(e)
        finally:
            self.stats["execution_time"] += time.time() - start_time

    async def _execute_select_batch(self, operations: list[BatchOperation]):
        """批量执行查询操作"""
        # 合并相似的查询条件
        merged_conditions = self._merge_select_conditions(operations)

        async with get_db_session() as session:
            results = []
            for conditions, ops in merged_conditions.items():
                try:
                    # 构建查询
                    query = select(ops[0].model_class)
                    for field_name, value in conditions.items():
                        model_attr = getattr(ops[0].model_class, field_name)
                        if isinstance(value, list | tuple | set):
                            query = query.where(model_attr.in_(value))
                        else:
                            query = query.where(model_attr == value)

                    # 执行查询
                    result = await session.execute(query)
                    data = result.scalars().all()

                    # 分发结果到各个操作
                    for op in ops:
                        if len(conditions) == 1 and len(ops) == 1:
                            # 单个查询，直接返回所有结果
                            op_result = data
                        else:
                            # 需要根据条件过滤结果
                            op_result = [
                                item
                                for item in data
                                if all(getattr(item, k) == v for k, v in op.conditions.items() if hasattr(item, k))
                            ]
                        results.append(op_result)

                except Exception as e:
                    logger.error(f"批量查询失败: {e}", exc_info=True)
                    results.append([])

            return results if len(results) > 1 else results[0] if results else []

    async def _execute_insert_batch(self, operations: list[BatchOperation]):
        """批量执行插入操作"""
        async with get_db_session() as session:
            try:
                # 收集所有要插入的数据
                all_data = [op.data for op in operations if op.data]
                if not all_data:
                    return []

                # 批量插入
                stmt = insert(operations[0].model_class).values(all_data)
                result = await session.execute(stmt)
                await session.commit()

                return [result.rowcount] * len(operations)

            except Exception as e:
                await session.rollback()
                logger.error(f"批量插入失败: {e}", exc_info=True)
                return [0] * len(operations)

    async def _execute_update_batch(self, operations: list[BatchOperation]):
        """批量执行更新操作"""
        async with get_db_session() as session:
            try:
                results = []
                for op in operations:
                    if not op.data or not op.conditions:
                        results.append(0)
                        continue

                    stmt = update(op.model_class)
                    for field_name, value in op.conditions.items():
                        model_attr = getattr(op.model_class, field_name)
                        if isinstance(value, list | tuple | set):
                            stmt = stmt.where(model_attr.in_(value))
                        else:
                            stmt = stmt.where(model_attr == value)

                    stmt = stmt.values(**op.data)
                    result = await session.execute(stmt)
                    results.append(result.rowcount)

                await session.commit()
                return results

            except Exception as e:
                await session.rollback()
                logger.error(f"批量更新失败: {e}", exc_info=True)
                return [0] * len(operations)

    async def _execute_delete_batch(self, operations: list[BatchOperation]):
        """批量执行删除操作"""
        async with get_db_session() as session:
            try:
                results = []
                for op in operations:
                    if not op.conditions:
                        results.append(0)
                        continue

                    stmt = delete(op.model_class)
                    for field_name, value in op.conditions.items():
                        model_attr = getattr(op.model_class, field_name)
                        if isinstance(value, list | tuple | set):
                            stmt = stmt.where(model_attr.in_(value))
                        else:
                            stmt = stmt.where(model_attr == value)

                    result = await session.execute(stmt)
                    results.append(result.rowcount)

                await session.commit()
                return results

            except Exception as e:
                await session.rollback()
                logger.error(f"批量删除失败: {e}", exc_info=True)
                return [0] * len(operations)

    def _merge_select_conditions(self, operations: list[BatchOperation]) -> dict[tuple, list[BatchOperation]]:
        """合并相似的查询条件"""
        merged = {}

        for op in operations:
            # 生成条件键
            condition_key = tuple(sorted(op.conditions.keys()))

            if condition_key not in merged:
                merged[condition_key] = {}

            # 尝试合并相同字段的值
            for field_name, value in op.conditions.items():
                if field_name not in merged[condition_key]:
                    merged[condition_key][field_name] = []

                if isinstance(value, list | tuple | set):
                    merged[condition_key][field_name].extend(value)
                else:
                    merged[condition_key][field_name].append(value)

            # 记录操作
            if condition_key not in merged:
                merged[condition_key] = {"_operations": []}
            if "_operations" not in merged[condition_key]:
                merged[condition_key]["_operations"] = []
            merged[condition_key]["_operations"].append(op)

        # 去重并构建最终条件
        final_merged = {}
        for condition_key, conditions in merged.items():
            operations = conditions.pop("_operations")

            # 去重
            for field_name, values in conditions.items():
                conditions[field_name] = list(set(values))

            final_merged[condition_key] = operations

        return final_merged

    def get_stats(self) -> dict[str, Any]:
        """获取统计信息"""
        return {
            **self.stats,
            "cache_size": len(self._result_cache),
            "queue_sizes": {k: len(v) for k, v in self.operation_queues.items()},
            "is_running": self._is_running,
        }


# 全局数据库批量调度器实例
db_batch_scheduler = DatabaseBatchScheduler()


@asynccontextmanager
async def get_batch_session():
    """获取批量会话上下文管理器"""
    if not db_batch_scheduler._is_running:
        await db_batch_scheduler.start()

    try:
        yield db_batch_scheduler
    finally:
        pass


# 便捷函数
async def batch_select(model_class: Any, conditions: dict[str, Any]) -> Any:
    """批量查询"""
    operation = BatchOperation(operation_type="select", model_class=model_class, conditions=conditions)
    return await db_batch_scheduler.add_operation(operation)


async def batch_insert(model_class: Any, data: dict[str, Any]) -> int:
    """批量插入"""
    operation = BatchOperation(operation_type="insert", model_class=model_class, conditions={}, data=data)
    return await db_batch_scheduler.add_operation(operation)


async def batch_update(model_class: Any, conditions: dict[str, Any], data: dict[str, Any]) -> int:
    """批量更新"""
    operation = BatchOperation(operation_type="update", model_class=model_class, conditions=conditions, data=data)
    return await db_batch_scheduler.add_operation(operation)


async def batch_delete(model_class: Any, conditions: dict[str, Any]) -> int:
    """批量删除"""
    operation = BatchOperation(operation_type="delete", model_class=model_class, conditions=conditions)
    return await db_batch_scheduler.add_operation(operation)


def get_db_batch_scheduler() -> DatabaseBatchScheduler:
    """获取数据库批量调度器实例"""
    return db_batch_scheduler