# -*- coding: utf-8 -*-
"""数据库连接管理器（限制最大连接超时时间）。"""

import time

import psycopg2
import psycopg2.extras

# 连接重试参数：应对 PostgreSQL 瞬时不可用
_CONNECT_MAX_RETRIES = 3
_CONNECT_RETRY_DELAY = 1.0


class DatabaseConnection:
    """封装 psycopg2 连接，支持会话参数和超时保护。"""

    def __init__(self, dsn: str, session: dict = None, connect_timeout: int = None):
        self._dsn = dsn
        self._session = session or {}
        self._connect_timeout = connect_timeout
        self.conn = self._open_connection()

    def _open_connection(self):
        """创建并初始化连接（包含会话参数）。"""
        timeout_val = self._connect_timeout if self._connect_timeout is not None else 5
        # 生产环境要求：数据库连接超时不得超过 20 秒。
        timeout_val = max(1, min(int(timeout_val), 20))

        # CHANGE 2026-03-06 | intent: 修复 Windows GBK 环境下 psycopg2 连接握手的 UnicodeDecodeError
        # assumptions: libpq 默认使用系统 locale 的 client_encoding，Windows 中文系统为 GBK/CP936
        # 边界: 显式指定 client_encoding=utf8 确保连接层始终使用 UTF-8，与数据库 server_encoding 一致
        # 验证: web-admin 手动触发 ETL 全量 flow，不再出现 0xd6 解码错误
        last_exc = None
        for attempt in range(_CONNECT_MAX_RETRIES):
            try:
                conn = psycopg2.connect(
                    self._dsn,
                    connect_timeout=timeout_val,
                    options="-c client_encoding=utf8",
                    keepalives=1,
                    keepalives_idle=60,
                    keepalives_interval=10,
                    keepalives_count=3,
                )
                break
            except psycopg2.OperationalError as e:
                last_exc = e
                if attempt < _CONNECT_MAX_RETRIES - 1:
                    time.sleep(_CONNECT_RETRY_DELAY * (attempt + 1))
        else:
            raise last_exc
        conn.autocommit = False

        # 会话参数（时区、语句超时等）
        if self._session:
            with conn.cursor() as c:
                if self._session.get("timezone"):
                    c.execute("SET TIME ZONE %s", (self._session["timezone"],))
                if self._session.get("statement_timeout_ms") is not None:
                    c.execute(
                        "SET statement_timeout = %s",
                        (int(self._session["statement_timeout_ms"]),),
                    )
                if self._session.get("lock_timeout_ms") is not None:
                    c.execute(
                        "SET lock_timeout = %s", (int(self._session["lock_timeout_ms"]),)
                    )
                if self._session.get("idle_in_tx_timeout_ms") is not None:
                    c.execute(
                        "SET idle_in_transaction_session_timeout = %s",
                        (int(self._session["idle_in_tx_timeout_ms"]),),
                    )
        return conn

    def query(self, sql: str, args=None):
        """Execute a query and fetch all rows."""
        with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as c:
            c.execute(sql, args)
            return c.fetchall()

    def execute(self, sql: str, args=None) -> int:
        """Execute a SQL statement without returning rows. Returns rowcount."""
        with self.conn.cursor() as c:
            c.execute(sql, args)
            return c.rowcount

    def commit(self):
        """Commit current transaction."""
        self.conn.commit()

    def rollback(self):
        """Rollback current transaction."""
        self.conn.rollback()

    def close(self):
        """Safely close the connection."""
        try:
            self.conn.close()
        except Exception:
            pass

    def ensure_open(self) -> bool:
        """确保连接可用，若已关闭则尝试重连。"""
        try:
            if getattr(self.conn, "closed", 0):
                self.conn = self._open_connection()
            return True
        except Exception:
            return False
