#!/usr/bin/env python3
"""db_pool.py — 轻量 MySQL 连接池（基于 pymysql + queue，无需额外依赖）"""

import pymysql
from queue import Empty, Queue
from threading import Lock

from common import MYSQL_CFG

_pool = None
_lock = Lock()

# 连接超时配置（避免无密码/错误配置时挂起）
_CONNECT_TIMEOUT = 5  # 秒
_READ_TIMEOUT = 10    # 秒

# 从池中取连接的等待上限。
# 池容量 8；若取连接无 timeout，一旦有调用方未归还导致池被耗尽，
# 所有请求都会永久阻塞在 _pool.get() 上（进程只剩线程堆积、不报错），
# 设上限才能把"静默挂死"降级为"可观测的报错"。
_POOL_GET_TIMEOUT = 10  # 秒


def _connect_cfg():
    """返回带超时参数的连接配置"""
    cfg = dict(MYSQL_CFG)
    cfg.setdefault("connect_timeout", _CONNECT_TIMEOUT)
    cfg.setdefault("read_timeout", _READ_TIMEOUT)
    return cfg


def _init_pool(max_conn=8):
    """初始化连接池（双检锁：裸判 _pool is None 在并发首调时会建出多个池）"""
    global _pool
    if _pool is None:
        with _lock:
            if _pool is None:
                q = Queue(maxsize=max_conn)
                for _ in range(max_conn):
                    q.put(None)  # None 表示需要时才创建
                _pool = q


def get_conn(max_retries=3, retry_delay=2):
    """从池中获取连接（惰性创建 + 自动重连 + 重试）"""
    global _pool
    if _pool is None:
        _init_pool()
    try:
        conn = _pool.get(timeout=_POOL_GET_TIMEOUT)
    except Empty:
        raise TimeoutError(
            f"连接池已耗尽（容量 {_pool.maxsize}），等待 {_POOL_GET_TIMEOUT}s 未取到连接。"
            "通常意味着有调用方持连接未归还（异常路径缺 put_conn）。"
        )
    cfg = _connect_cfg()

    last_err = None
    for attempt in range(max_retries):
        try:
            if conn is None:
                conn = pymysql.connect(**cfg)
            else:
                conn.ping(reconnect=True)
            return conn
        except Exception as e:
            last_err = e
            # 坏连接必须关掉再置 None，别带着它继续重试
            if conn is not None:
                try:
                    conn.close()
                except Exception:
                    pass
            conn = None
            if attempt < max_retries - 1:
                import time
                time.sleep(retry_delay)

    # 全部重试失败：必须把已占用的池位归还，否则每失败一次就永久少一个连接，
    # 池会被失败请求逐步耗尽。
    put_conn(None)
    raise last_err or Exception("数据库连接失败")


def put_conn(conn):
    """归还连接到池中

    归还前必须 rollback：写路径（INSERT/UPDATE）若半途抛错，连接会带着
    未提交事务与行锁回到池里，下一个请求拿到它就会读到脏数据甚至死锁。
    """
    global _pool
    if _pool is None:
        _init_pool()
    try:
        if conn and conn.open:
            try:
                conn.rollback()
            except Exception:
                # 回滚都失败说明连接已不可用，丢弃之，用 None 占位
                try:
                    conn.close()
                except Exception:
                    pass
                _pool.put_nowait(None)
                return
            _pool.put_nowait(conn)
        else:
            _pool.put_nowait(None)
    except Exception:
        # 池满了，直接关闭
        try:
            if conn:
                conn.close()
        except Exception:
            pass


class DBContext:
    """上下文管理器，自动获取和归还连接"""
    def __enter__(self):
        self.conn = get_conn()
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 异常出栈时显式回滚（put_conn 内也有兜底 rollback），
        # 避免把未提交事务留在连接上带回池。
        if exc_type is not None and self.conn is not None:
            try:
                self.conn.rollback()
            except Exception:
                pass
        put_conn(self.conn)
        self.conn = None
        return False