"""高性能 SQLite 数据库管理器 - 线程安全的数据库操作封装.

本模块提供:
- DatabaseConnectionPool: 线程安全的连接池
- DatabaseManager: 高级数据库操作管理器
- 支持上下文管理器和装饰器进行事务控制
"""

from __future__ import annotations

import sqlite3
import threading
import types
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from queue import Empty, Queue
from typing import Callable, Generator, TypeVar

from ..logger import logger

__all__ = [
    "DatabaseConnectionPool",
    "DatabaseError",
    "DatabaseManager",
    "IsolationLevel",
    "RowFactory",
    "TransactionError",
]


# 类型别名
RowFactory = Callable[[sqlite3.Cursor], Callable[[sqlite3.Row], dict]]
IsolationLevel = str
T = TypeVar("T")


class DatabaseError(Exception):
    """数据库操作异常基类."""

    pass


class TransactionError(DatabaseError):
    """事务处理异常."""

    pass


@dataclass
class PoolConfig:
    """连接池配置.

    Attributes
    ----------
        max_size: 最大连接数
        min_size: 最小连接数
        timeout: 获取连接超时时间 (秒)
        recycle_timeout: 连接回收超时时间 (秒)
        check_connection_on_borrow: 借用连接时是否检查有效性
    """

    max_size: int = 10
    min_size: int = 2
    timeout: float = 30.0
    recycle_timeout: float = 3600.0  # 1 小时
    check_connection_on_borrow: bool = True


@dataclass
class ConnectionWrapper:
    """连接包装器，用于跟踪连接状态.

    Attributes
    ----------
        conn: SQLite 连接对象
        created_at: 创建时间戳
        last_used: 最后使用时间戳
        borrow_count: 借用次数
    """

    conn: sqlite3.Connection
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    last_used: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    borrow_count: int = 0


class DatabaseConnectionPool:
    """线程安全的 SQLite 连接池.

    特性:
    - 使用 RLock 保证线程安全
    - 支持最小/最大连接数限制
    - 自动连接回收机制
    - 连接借用计数和统计
    - 支持上下文管理器

    示例:
    --------
    >>> pool = DatabaseConnectionPool(":memory:")
    >>> with pool.get_connection() as conn:
    ...     cursor = conn.execute("SELECT 1")
    """

    def __init__(self, db_path: str | Path, config: PoolConfig | None = None) -> None:
        """初始化连接池.

        Args:
            db_path: 数据库文件路径
            config: 连接池配置,None 则使用默认配置
        """
        self.db_path = Path(db_path)
        self.config = config or PoolConfig()
        self._pool: Queue[ConnectionWrapper] = Queue()
        self._lock = threading.RLock()
        self._total_connections = 0
        self._shutdown = False

        # 预创建最小连接数
        self._initialize_min_connections()

    def _try_initialize_single_connection(self) -> bool:
        """尝试初始化单个连接.

        Returns
        -------
            初始化成功返回 True,失败返回 False
        """
        try:
            conn = self._create_connection()
            self._pool.put(ConnectionWrapper(conn=conn))
            self._total_connections += 1
        except sqlite3.Error:
            logger.exception("初始化连接失败")
            return False
        else:
            return True

    def _initialize_min_connections(self) -> None:
        """预创建最小连接数的连接."""
        with self._lock:
            while self._total_connections < self.config.min_size:
                if not self._try_initialize_single_connection():
                    break

    def _create_connection(self) -> sqlite3.Connection:
        """创建新的数据库连接.

        Returns
        -------
            新创建的 SQLite 连接
        """
        conn = sqlite3.connect(
            str(self.db_path),
            check_same_thread=False,  # 允许跨线程使用(由连接池管理)
            isolation_level=None,  # 自动提交模式,由事务管理器控制
        )
        # 启用 WAL 模式以获得更好的并发性能
        conn.execute("PRAGMA journal_mode=WAL")
        # 设置 busy timeout
        conn.execute("PRAGMA busy_timeout=5000")
        # 启用外键约束
        conn.execute("PRAGMA foreign_keys=ON")
        return conn

    def _is_connection_valid(self, wrapper: ConnectionWrapper) -> bool:
        """检查连接是否有效.

        Args:
            wrapper: 连接包装器

        Returns
        -------
            连接有效返回 True,否则返回 False
        """
        try:
            # 检查连接是否可以执行简单查询
            wrapper.conn.execute("SELECT 1")
            # 检查连接是否超时
            if self.config.recycle_timeout > 0:
                age = (datetime.now(timezone.utc) - wrapper.created_at).total_seconds()
                if age > self.config.recycle_timeout:
                    return False
        except sqlite3.Error:
            return False
        else:
            return True

    def _return_connection(self, wrapper: ConnectionWrapper) -> None:
        """归还连接到连接池.

        Args:
            wrapper: 连接包装器
        """
        with self._lock:
            if self._shutdown:
                self._close_connection(wrapper)
                return

            # 检查连接是否需要回收
            if not self._is_connection_valid(wrapper):
                self._close_connection(wrapper)
                # 创建新连接替换
                try:
                    new_conn = self._create_connection()
                    self._pool.put(ConnectionWrapper(conn=new_conn))
                except sqlite3.Error:
                    logger.exception("创建替换连接失败")
                    with self._lock:
                        self._total_connections -= 1
                return

            # 更新最后使用时间并归还到池中
            wrapper.last_used = datetime.now(timezone.utc)
            self._pool.put(wrapper)

    def _close_connection(self, wrapper: ConnectionWrapper) -> None:
        """关闭连接.

        Args:
            wrapper: 连接包装器
        """
        try:
            wrapper.conn.close()
        except sqlite3.Error:
            logger.exception("关闭连接失败")

    def _try_get_connection_from_pool(self) -> ConnectionWrapper | None:
        """尝试从连接池获取连接,如果失败则返回 None.

        Returns
        -------
            获取到的连接包装器,如果无法获取则返回 None
        """
        try:
            wrapper = self._pool.get_nowait()
            if self._is_connection_valid(wrapper):
                wrapper.borrow_count += 1
                return wrapper
            # 连接无效,关闭并创建新的
            self._close_connection(wrapper)
            with self._lock:
                self._total_connections -= 1
        except (sqlite3.Error, Empty):
            # 池为空或获取失败
            return None
        else:
            return None

    def _try_create_new_connection(self) -> ConnectionWrapper | None:
        """尝试创建新连接.

        Returns
        -------
            创建成功的连接包装器,如果失败则返回 None
        """
        try:
            conn = self._create_connection()
            return ConnectionWrapper(conn=conn, borrow_count=1)
        except sqlite3.Error:
            logger.exception("创建新连接失败")
            return None

    def _acquire_connection_with_wait(self) -> ConnectionWrapper:
        """在等待后获取连接.

        Returns
        -------
            获取到的连接包装器

        Raises
        ------
            DatabaseError: 当无法获取连接时
        """
        try:
            wrapper = self._pool.get(timeout=self.config.timeout)
        except Empty as e:
            msg = f"获取连接超时 ({self.config.timeout}秒)"
            raise DatabaseError(msg) from e

        if not self._is_connection_valid(wrapper):
            self._close_connection(wrapper)
            with self._lock:
                self._total_connections -= 1
            msg = "连接无效,重新获取"
            raise DatabaseError(msg)
        wrapper.borrow_count += 1
        return wrapper

    def _process_acquired_connection(self, wrapper: ConnectionWrapper | None) -> tuple[ConnectionWrapper | None, bool]:
        """处理已获取的连接.

        Args:
            wrapper: 连接包装器

        Returns
        -------
            tuple[Optional[ConnectionWrapper], bool]: (连接包装器,是否继续循环)
        """
        if wrapper is not None:
            return wrapper, False

        # 池为空,尝试创建新连接
        if self._total_connections < self.config.max_size:
            wrapper = self._try_create_new_connection()
            if wrapper is not None:
                with self._lock:
                    self._total_connections += 1
                return wrapper, False

        # 达到最大连接数,需要等待
        return None, True

    @contextmanager
    def get_connection(self) -> Generator[sqlite3.Connection, None, None]:
        """获取数据库连接的上下文管理器.

        Yields
        ------
            SQLite 连接对象

        Raises
        ------
            DatabaseError: 当无法获取连接时
        """
        if self._shutdown:
            msg = "连接池已关闭"
            raise DatabaseError(msg)

        wrapper: ConnectionWrapper | None = None

        try:
            # 尝试从池中获取连接
            with self._lock:
                while True:
                    # 尝试从池中获取有效连接
                    wrapper = self._try_get_connection_from_pool()
                    wrapper, should_continue = self._process_acquired_connection(wrapper)
                    if not should_continue:
                        break

                    # 需要等待其他线程归还连接
                    self._lock.release()
                    try:
                        wrapper = self._acquire_connection_with_wait()
                        break
                    except DatabaseError as e:
                        # 连接无效或超时,继续重试
                        if "超时" in str(e):
                            # 超时错误,重新抛出
                            raise
                        # 连接无效,继续重试
                        continue
                    finally:
                        self._lock.acquire()

            if wrapper is None:
                msg = "无法获取有效连接"
                raise DatabaseError(msg)

            yield wrapper.conn

        finally:
            # 归还连接
            if wrapper is not None:
                self._return_connection(wrapper)

    def get_stats(self) -> dict:
        """获取连接池统计信息.

        Returns
        -------
            包含连接池状态的字典
        """
        with self._lock:
            pool_size = self._pool.qsize()
            return {
                "total_connections": self._total_connections,
                "available_connections": pool_size,
                "in_use_connections": self._total_connections - pool_size,
                "max_size": self.config.max_size,
                "min_size": self.config.min_size,
                "db_path": str(self.db_path),
            }

    def close(self) -> None:
        """关闭连接池并释放所有连接."""
        with self._lock:
            self._shutdown = True
            while not self._pool.empty():
                wrapper = self._pool.get_nowait()
                self._close_connection(wrapper)
            self._total_connections = 0
            logger.info(f"连接池已关闭:{self.db_path}")

    def __enter__(self) -> DatabaseConnectionPool:
        """进入上下文管理器."""
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: types.TracebackType | None,
    ) -> None:
        """退出上下文管理器."""
        self.close()


class DatabaseManager:
    """高级数据库管理器,提供便捷的 CRUD 操作.

    特性:
    - 基于连接池的线程安全操作
    - 装饰器和上下文管理器风格的事务控制
    - 字典形式的结果集返回
    - 批量操作支持
    - 完整的类型注解

    示例:
    --------
    >>> db = DatabaseManager(":memory:")
    >>> db.execute('''
    ...     CREATE TABLE users (
    ...         id INTEGER PRIMARY KEY,
    ...         name TEXT NOT NULL,
    ...         email TEXT UNIQUE
    ...     )
    ... ''')
    >>> db.insert('users', {'name': 'Alice', 'email': 'alice@example.com'})
    >>> users = db.fetch_all('SELECT * FROM users')
    >>> db.close()
    """

    def __init__(self, db_path: str | Path, pool_config: PoolConfig | None = None) -> None:
        """初始化数据库管理器.

        Args:
            db_path: 数据库文件路径
            pool_config: 连接池配置,None 则使用默认配置
        """
        self.pool = DatabaseConnectionPool(db_path, pool_config)
        # 默认使用字典形式的行工厂
        self._row_factory: Callable[[sqlite3.Cursor], Callable[[sqlite3.Row], dict]] | None = self.dict_row_factory

    @property
    def row_factory(self) -> RowFactory | None:
        """获取行工厂函数."""
        return self._row_factory

    @row_factory.setter
    def row_factory(self, factory: RowFactory | None) -> None:
        """设置行工厂函数.

        Args:
            factory: 行工厂函数,接收 cursor 返回 dict
        """
        self._row_factory = factory

    @staticmethod
    def dict_row_factory(cursor: sqlite3.Cursor) -> Callable[[sqlite3.Row], dict]:
        """字典形式的行工厂.

        Args:
            cursor: SQLite 游标

        Returns
        -------
            将行转换为字典的函数
        """

        def row_factory(row: sqlite3.Row) -> dict:
            return dict(zip([description[0] for description in cursor.description], row))

        return row_factory

    def execute(
        self,
        query: str,
        params: tuple | dict | None = None,
        *,
        fetch: bool = False,
        commit: bool = False,
    ) -> list[dict] | int:
        """执行 SQL 查询.

        Args:
            query: SQL 查询语句
            params: 查询参数
            fetch: 是否返回结果
            commit: 是否提交事务

        Returns
        -------
            如果 fetch=True,返回结果列表;否则返回受影响的行数
        """
        with self.pool.get_connection() as conn:
            if self._row_factory:
                conn.row_factory = sqlite3.Row

            cursor = conn.cursor()

            try:
                if params:
                    cursor.execute(query, params)
                else:
                    cursor.execute(query)

                if commit:
                    conn.commit()

                if fetch:
                    rows = cursor.fetchall()
                    if self._row_factory and cursor.description:
                        # 使用 row factory 转换为字典
                        row_factory_func = self._row_factory(cursor)
                        return [row_factory_func(row) for row in rows]
                    return rows
            except sqlite3.Error as e:
                conn.rollback()
                msg = f"执行查询失败:{e}"
                raise DatabaseError(msg) from e
            else:
                return cursor.rowcount

    def fetch_one(self, query: str, params: tuple | dict | None = None) -> dict | None:
        """获取单条记录.

        Args:
            query: SQL 查询语句
            params: 查询参数

        Returns
        -------
            单条记录的字典形式,如果没有记录则返回 None
        """
        results = self.execute(query, params, fetch=True)
        if isinstance(results, list):
            return results[0] if results else None
        msg = "fetch_one 期望列表类型的返回值"
        raise DatabaseError(msg)

    def fetch_all(self, query: str, params: tuple | dict | None = None) -> list[dict]:
        """获取所有记录.

        Args:
            query: SQL 查询语句
            params: 查询参数

        Returns
        -------
            所有记录的字典列表
        """
        results = self.execute(query, params, fetch=True)
        if isinstance(results, list):
            return results
        msg = "fetch_all 期望列表类型的返回值"
        raise DatabaseError(msg)

    def insert(self, table: str, data: dict) -> int:
        """插入单条记录.

        Args:
            table: 表名
            data: 列名到值的映射

        Returns
        -------
            插入行的 rowid
        """
        columns = ", ".join(data.keys())
        placeholders = ", ".join(["?" for _ in data])
        query = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"

        with self.pool.get_connection() as conn:
            cursor = conn.cursor()
            try:
                cursor.execute(query, list(data.values()))
                conn.commit()
            except sqlite3.Error as e:
                conn.rollback()
                msg = f"插入记录失败:{e}"
                raise DatabaseError(msg) from e
            else:
                lastrowid = cursor.lastrowid
                return lastrowid if lastrowid is not None else 0

    def insert_many(self, table: str, data_list: list[dict]) -> int:
        """批量插入记录.

        Args:
            table: 表名
            data_list: 数据字典列表

        Returns
        -------
            插入的行数
        """
        if not data_list:
            return 0

        # 确保所有字典有相同的键
        columns = list(data_list[0].keys())
        placeholders = ", ".join(["?" for _ in columns])
        query = f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({placeholders})"

        with self.pool.get_connection() as conn:
            cursor = conn.cursor()
            try:
                values = [[data.get(col) for col in columns] for data in data_list]
                cursor.executemany(query, values)
                conn.commit()
            except sqlite3.Error as e:
                conn.rollback()
                msg = f"批量插入失败:{e}"
                raise DatabaseError(msg) from e
            else:
                return cursor.rowcount

    def count(self, table: str, where: str, where_params: tuple | None = None) -> int:
        """统计记录数.

        Args:
            table: 表名
            where: WHERE 子句(不包含 WHERE 关键字),使用占位符如 column=?
            where_params: WHERE 子句的参数值元组

        Returns
        -------
            记录数

        Raises
        ------
            DatabaseError: 当参数验证失败时
        """
        if where:
            query = f"SELECT COUNT(*) FROM {table} WHERE {where}"
            params = where_params or ()
            result = self.fetch_one(query, params)
        else:
            query = f"SELECT COUNT(*) FROM {table}"
            result = self.fetch_one(query)

        return result["COUNT(*)"] if result else 0

    def query(self, query: str, params: tuple | None = None) -> list[dict]:
        """执行 SELECT 查询并返回结果.

        为了向后兼容而提供的便捷方法,封装了 fetch_all.

        Args:
            query: SQL SELECT 查询
            params: 查询参数

        Returns
        -------
            结果字典列表
        """
        return self.fetch_all(query, params)

    def query_one(self, query: str, params: tuple | None = None) -> dict | None:
        """执行 SELECT 查询并返回单个结果.

        为了向后兼容而提供的便捷方法,封装了 fetch_one.

        Args:
            query: SQL SELECT 查询
            params: 查询参数

        Returns
        -------
            单个结果字典,如果没有结果则返回 None
        """
        return self.fetch_one(query, params)

    def update(self, table: str, data: dict, where: str, where_params: tuple | None = None) -> int:
        """更新记录.

        Args:
            table: 表名
            data: 要更新的列和值
            where: WHERE 子句(不包含 WHERE 关键字),使用占位符如 column=?
            where_params: WHERE 子句的参数值元组

        Returns
        -------
            受影响的行数

        Raises
        ------
            DatabaseError: 当参数验证失败时
        """
        # 验证表名(只允许字母,数字,下划线)
        if not all(c.isalnum() or c == "_" for c in table):
            msg = f"无效的表名:{table}"
            raise DatabaseError(msg)

        set_clause = ", ".join([f"{col} = ?" for col in data])
        query = f"UPDATE {table} SET {set_clause} WHERE {where}"

        params = list(data.values())
        if where_params:
            params.extend(where_params)

        with self.pool.get_connection() as conn:
            cursor = conn.cursor()
            try:
                cursor.execute(query, tuple(params))
                conn.commit()
            except sqlite3.Error as e:
                conn.rollback()
                msg = f"更新记录失败:{e}"
                raise DatabaseError(msg) from e
            else:
                return cursor.rowcount

    def delete(self, table: str, where: str, where_params: tuple | None = None) -> int:
        """删除记录.

        Args:
            table: 表名
            where: WHERE 子句(不包含 WHERE 关键字),使用占位符如 column=?
            where_params: WHERE 子句的参数值元组

        Returns
        -------
            删除的行数

        Raises
        ------
            DatabaseError: 当参数验证失败时
        """
        # 验证表名(只允许字母,数字,下划线)
        if not all(c.isalnum() or c == "_" for c in table):
            msg = f"无效的表名:{table}"
            raise DatabaseError(msg)

        query = f"DELETE FROM {table} WHERE {where}"

        with self.pool.get_connection() as conn:
            cursor = conn.cursor()
            try:
                if where_params:
                    cursor.execute(query, where_params)
                else:
                    cursor.execute(query)
                conn.commit()
            except sqlite3.Error as e:
                conn.rollback()
                msg = f"删除记录失败:{e}"
                raise DatabaseError(msg) from e
            else:
                return cursor.rowcount

    @contextmanager
    def transaction(self) -> Generator[sqlite3.Connection, None, None]:
        """事务上下文管理器.

        Yields
        ------
            数据库连接对象,在事务中

        Example:
        --------
        >>> with db.transaction() as conn:
        ...     conn.execute("INSERT INTO users (name) VALUES (?)", ("Alice",))
        ...     conn.execute("INSERT INTO users (name) VALUES (?)", ("Bob",))
        ...     # 自动提交,如果抛出异常则回滚
        """
        with self.pool.get_connection() as conn:
            try:
                # 在自动提交模式下,需要显式开始事务
                conn.execute("BEGIN")
                yield conn
                conn.commit()
            except sqlite3.Error as e:
                conn.rollback()
                msg = f"事务执行失败:{e}"
                raise TransactionError(msg) from e

    def transactional(self, func: Callable[..., T]) -> Callable[..., T]:
        """事务装饰器.

        Args:
            func: 要装饰的函数,可以接受 conn 作为第一个参数或不接受

        Returns
        -------
            包装后的函数,自动在事务中执行

        Example:
        --------
        >>> @db.transactional
        ... def my_func(conn):  # 带连接参数
        ...     conn.execute("INSERT INTO test VALUES (1, ?)", ("Alice",))

        >>> @db.transactional
        ... def my_func():  # 不带连接参数
        ...     db.insert("test", {"name": "Alice"})
        """

        def wrapper(*args: object, **kwargs: object) -> T:
            with self.transaction() as conn:
                # 检查函数是否接受 conn 参数
                import inspect

                sig = inspect.signature(func)
                params = list(sig.parameters.values())

                if params and params[0].name == "conn":
                    # 函数接受 conn 参数
                    return func(conn, *args[1:] if args else (), **kwargs)
                # 函数不接受 conn 参数
                return func(*args, **kwargs)

        return wrapper

    def table_exists(self, table_name: str) -> bool:
        """检查表是否存在.

        Args:
            table_name: 表名

        Returns
        -------
            表存在返回 True,否则返回 False
        """
        query = "SELECT name FROM sqlite_master WHERE type='table' AND name=?"
        result = self.fetch_one(query, (table_name,))
        return result is not None

    def get_table_info(self, table_name: str) -> list[dict]:
        """获取表结构信息.

        Args:
            table_name: 表名

        Returns
        -------
            表结构信息的字典列表

        Raises
        ------
            DatabaseError: 当表名无效时
        """
        # 验证表名(只允许字母,数字,下划线)
        if not all(c.isalnum() or c == "_" for c in table_name):
            msg = f"无效的表名:{table_name}"
            raise DatabaseError(msg)

        query = f"PRAGMA table_info({table_name})"
        return self.fetch_all(query)

    def get_stats(self) -> dict:
        """获取数据库统计信息.

        Returns
        -------
            包含数据库和连接池状态的字典
        """
        stats = self.pool.get_stats()

        # 添加数据库基本信息
        with self.pool.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT COUNT(*) FROM sqlite_master WHERE type='table'")
            table_count = cursor.fetchone()[0]
            stats["table_count"] = table_count

        return stats

    def close(self) -> None:
        """关闭数据库管理器并释放连接池."""
        self.pool.close()

    def __enter__(self) -> DatabaseManager:
        """进入上下文管理器."""
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: types.TracebackType | None,
    ) -> None:
        """退出上下文管理器."""
        self.close()
