# -*- coding: utf-8 -*-
"""
数据库操作（批量、RETURNING支持）
"""
import re
from typing import List, Dict, Tuple
import psycopg2.extras
from .connection import DatabaseConnection


class DatabaseOperations(DatabaseConnection):
    """扩展数据库操作（包含批量upsert和returning支持）"""
    
    def batch_execute(self, sql: str, rows: List[Dict], page_size: int = 1000):
        """批量执行SQL（不带RETURNING）"""
        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[Dict], page_size: int = 1000) -> Tuple[int, int]:
        """
        批量 UPSERT 并统计插入/更新数
        
        Args:
            sql: 包含RETURNING子句的SQL
            rows: 数据行列表
            page_size: 批次大小
            
        Returns:
            (inserted_count, updated_count) 元组
        """
        if not rows:
            return (0, 0)
        
        use_returning = "RETURNING" in sql.upper()
        
        with self.conn.cursor() as c:
            if not use_returning:
                psycopg2.extras.execute_batch(c, sql, rows, page_size=page_size)
                return (0, 0)
            
            # 优先尝试向量化执行
            try:
                inserted, updated = self._execute_with_returning_vectorized(c, sql, rows, page_size)
                return (inserted, updated)
            except Exception:
                # 回退到逐行执行
                return self._execute_with_returning_row_by_row(c, sql, rows)
    
    def _execute_with_returning_vectorized(self, cursor, sql: str, rows: List[Dict], page_size: int) -> Tuple[int, int]:
        """向量化执行（使用execute_values）"""
        # 解析VALUES子句
        m = re.search(r"VALUES\s*\((.*?)\)", sql, flags=re.IGNORECASE | re.DOTALL)
        if not m:
            raise ValueError("Cannot parse VALUES clause")
        
        tpl = "(" + m.group(1) + ")"
        base_sql = sql[:m.start()] + "VALUES %s" + sql[m.end():]
        
        ret = psycopg2.extras.execute_values(
            cursor, base_sql, rows, template=tpl, page_size=page_size, fetch=True
        )
        
        if not ret:
            return (0, 0)
        
        inserted = 0
        for rec in ret:
            flag = self._extract_inserted_flag(rec)
            if flag:
                inserted += 1
        
        return (inserted, len(ret) - inserted)
    
    def _execute_with_returning_row_by_row(self, cursor, sql: str, rows: List[Dict]) -> Tuple[int, int]:
        """逐行执行（回退方案）"""
        inserted = 0
        updated = 0
        
        for r in rows:
            cursor.execute(sql, r)
            try:
                rec = cursor.fetchone()
            except Exception:
                rec = None
            
            flag = self._extract_inserted_flag(rec) if rec else None
            
            if flag:
                inserted += 1
            else:
                updated += 1
        
        return (inserted, updated)
    
    @staticmethod
    def _extract_inserted_flag(rec) -> bool:
        """从返回记录中提取inserted标志"""
        if isinstance(rec, tuple):
            return bool(rec[0])
        elif isinstance(rec, dict):
            return bool(rec.get("inserted"))
        else:
            try:
                return bool(rec["inserted"])
            except Exception:
                return False


# 为了向后兼容，提供Pg别名
Pg = DatabaseOperations