# -*- coding: utf-8 -*-
"""数据库批量操作"""
import psycopg2.extras
import re

class DatabaseOperations:
    """数据库批量操作封装"""
    
    def __init__(self, connection):
        self._connection = connection
        self.conn = connection.conn
    
    def batch_execute(self, sql: str, rows: list, page_size: int = 1000):
        """批量执行SQL"""
        if not rows:
            return
        with self.conn.cursor() as c:
            psycopg2.extras.execute_batch(c, sql, rows, page_size=page_size)
    
    def batch_upsert_with_returning(self, sql: str, rows: list, 
                                     page_size: int = 1000) -> tuple:
        """批量UPSERT并返回插入/更新计数"""
        if not rows:
            return (0, 0)
        
        use_returning = "RETURNING" in sql.upper()
        
        # 不带 RETURNING：直接批量执行即可
        if not use_returning:
            with self.conn.cursor() as c:
                psycopg2.extras.execute_batch(c, sql, rows, page_size=page_size)
            return (0, 0)

        # 尝试向量化执行（execute_values + fetch returning）
        vectorized_failed = False
        m = re.search(r"VALUES\s*\((.*?)\)", sql, flags=re.IGNORECASE | re.DOTALL)
        if m:
            tpl = "(" + m.group(1) + ")"
            base_sql = sql[:m.start()] + "VALUES %s" + sql[m.end():]
            try:
                with self.conn.cursor() as c:
                    ret = psycopg2.extras.execute_values(
                        c, base_sql, rows, template=tpl, page_size=page_size, fetch=True
                    )
                if not ret:
                    return (0, 0)
                inserted = sum(1 for rec in ret if self._is_inserted(rec))
                return (inserted, len(ret) - inserted)
            except Exception:
                # 向量化失败后，事务通常处于 aborted 状态，需要先 rollback 才能继续执行。
                vectorized_failed = True

        if vectorized_failed:
            try:
                self.conn.rollback()
            except Exception:
                pass

        # 回退：逐行执行
        inserted = 0
        updated = 0
        with self.conn.cursor() as c:
            for r in rows:
                c.execute(sql, r)
                try:
                    rec = c.fetchone()
                except Exception:
                    rec = None

                if self._is_inserted(rec):
                    inserted += 1
                else:
                    updated += 1

        return (inserted, updated)
    
    @staticmethod
    def _is_inserted(rec) -> bool:
        """判断是否为插入操作"""
        if rec is None:
            return False
        if isinstance(rec, tuple):
            return bool(rec[0])
        if isinstance(rec, dict):
            return bool(rec.get("inserted"))
        return False

    # --- pass-through helpers -------------------------------------------------
    def commit(self):
        """提交事务（委托给底层连接）"""
        self._connection.commit()

    def rollback(self):
        """回滚事务（委托给底层连接）"""
        self._connection.rollback()

    def query(self, sql: str, args=None):
        """执行查询并返回结果"""
        return self._connection.query(sql, args)

    def execute(self, sql: str, args=None):
        """执行任意 SQL"""
        self._connection.execute(sql, args)

    def cursor(self):
        """暴露原生 cursor，供特殊操作使用"""
        return self.conn.cursor()
