from typing import List, Dict, Set

from sqlalchemy import insert, select, update, and_


# from sqlalchemy.dialects.mysql import insert as mysql_insert
# from sqlalchemy.dialects.postgresql import insert as pg_insert
# from sqlalchemy.dialects.sqlite import insert as sqlite_insert


class BatchUpsert:
    """
    同步版本的批量插入存在时更新工具类
    """

    def __init__(self, db):
        self.db = db
        self.engine = db.get_bind()

    def batch_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]) -> int:
        """同步批量upsert"""
        if not data_list:
            return 0

        dialect_name = self.engine.dialect.name

        if dialect_name == 'postgresql':
            return self._postgresql_upsert(data_list, table_class, unique_keys)
        elif dialect_name == 'mysql':
            return self._mysql_upsert(data_list, table_class, unique_keys)
        elif dialect_name == 'sqlite':
            return self._sqlite_upsert(data_list, table_class, unique_keys)
        else:
            return self._fallback_upsert(data_list, table_class, unique_keys)

    def _get_update_data(self, insert_column_collection, table, unique_keys: List[str], data_list: List[Dict]):
        if not data_list:
            return {}

        # 获取 data_list 中所有出现的字段（排除唯一键）
        data_keys: Set[str] = set()
        for data in data_list:
            for key in data.keys():
                if key not in unique_keys:
                    data_keys.add(key)
        update_data = {
            column.name: getattr(insert_column_collection, column.name)
            for column in table.columns
            if column.name not in unique_keys and column.name in data_keys
        }
        return update_data

    def _postgresql_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]) -> int:
        from sqlalchemy.dialects.postgresql import insert as pg_insert

        stmt = pg_insert(table_class).values(data_list)
        update_data = self._get_update_data(stmt.excluded, table_class.__table__, unique_keys, data_list)

        stmt = stmt.on_conflict_do_update(index_elements=unique_keys, set_=update_data)
        result = self.db.execute(stmt)
        return result.rowcount

    def _mysql_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]) -> int:
        from sqlalchemy.dialects.mysql import insert as mysql_insert
        stmt = mysql_insert(table_class).values(data_list)
        update_data = self._get_update_data(stmt.inserted, table_class.__table__, unique_keys, data_list)
        stmt = stmt.on_duplicate_key_update(**update_data)
        result = self.db.execute(stmt)
        return result.rowcount

    def _sqlite_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]) -> int:
        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
        stmt = sqlite_insert(table_class).values(data_list)
        update_data = self._get_update_data(stmt.excluded, table_class.__table__, unique_keys, data_list)

        stmt = stmt.on_conflict_do_update(index_elements=unique_keys, set_=update_data)
        result = self.db.execute(stmt)
        return result.rowcount

    def _fallback_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]) -> int:
        processed_count = 0
        for data in data_list:
            conditions = [
                getattr(table_class, key) == data[key]
                for key in unique_keys
                if key in data and data[key] is not None
            ]
            if not conditions:
                continue

            existing = self.db.execute(select(table_class).where(and_(*conditions))).scalar_one_or_none()
            if existing:
                update_data = {k: v for k, v in data.items() if k not in unique_keys}
                if update_data:
                    self.db.execute(
                        update(table_class)
                        .where(and_(*conditions))
                        .values(**update_data)
                    )
            else:
                self.db.execute(insert(table_class).values(**data))
            processed_count += 1

        self.db.commit()
        return processed_count


class AsyncBatchUpsert:
    def __init__(self, db):
        self.db = db
        self.engine = db.get_bind()

    async def batch_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]):
        """通用异步批量upsert方法，支持SQLite"""
        if not data_list:
            return 0

        dialect_name = self.engine.dialect.name

        if dialect_name == 'postgresql':
            return await self._postgresql_upsert(data_list, table_class, unique_keys)
        elif dialect_name == 'mysql':
            return await self._mysql_upsert(data_list, table_class, unique_keys)
        elif dialect_name == 'sqlite':
            return await self._sqlite_upsert(data_list, table_class, unique_keys)
        else:
            return await self._fallback_upsert(data_list, table_class, unique_keys)

    async def _sqlite_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]):
        """SQLite upsert - 使用 ON CONFLICT"""
        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
        try:
            stmt = sqlite_insert(table_class).values(data_list)

            # print("=== SQL INSERT 语句 ===")
            # compiled_stmt = stmt.compile(compile_kwargs={"literal_binds": True})
            # print(str(compiled_stmt))

            # 构建更新字段
            update_data = {}

            for column in table_class.__table__.columns:
                if column.name not in unique_keys:
                    # SQLite 使用 excluded 来引用冲突时的值
                    update_data[column.name] = getattr(stmt.excluded, column.name)

            stmt = stmt.on_conflict_do_update(
                index_elements=unique_keys,
                set_=update_data
            )

            result = await self.db.execute(stmt)
            return result.rowcount

        except Exception as e:
            print(f"SQLite upsert 错误: {e}")
            # 如果 ON CONFLICT 失败，回退到逐一处理
            return await self._fallback_upsert(data_list, table_class, unique_keys)

    async def _postgresql_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]):
        """PostgreSQL upsert"""
        from sqlalchemy.dialects.postgresql import insert as pg_insert

        stmt = pg_insert(table_class).values(data_list)
        update_data = {
            column.name: getattr(stmt.excluded, column.name)
            for column in table_class.__table__.columns
            if column.name not in unique_keys
        }

        stmt = stmt.on_conflict_do_update(
            index_elements=unique_keys,
            set_=update_data
        )

        result = await self.db.execute(stmt)
        return result.rowcount

    async def _mysql_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]):
        """MySQL upsert"""
        from sqlalchemy.dialects.mysql import insert as mysql_insert

        stmt = mysql_insert(table_class).values(data_list)
        update_data = {
            column.name: getattr(stmt.inserted, column.name)
            for column in table_class.__table__.columns
            if column.name not in unique_keys
        }

        stmt = stmt.on_duplicate_key_update(**update_data)
        result = await self.db.execute(stmt)
        return result.rowcount

    async def _fallback_upsert(self, data_list: List[Dict], table_class, unique_keys: List[str]):
        """回退方案：逐一处理"""
        processed_count = 0

        for data in data_list:
            # 构建查询条件
            conditions = []
            for key in unique_keys:
                if key in data:
                    conditions.append(getattr(table_class, key) == data[key])

            if not conditions:
                continue

            # 查询是否存在
            existing = await self.db.execute(
                select(table_class).where(and_(*conditions))
            )
            existing = existing.scalar_one_or_none()

            if existing:
                # 更新
                update_data = {k: v for k, v in data.items() if k not in unique_keys}
                if update_data:  # 确保有需要更新的字段
                    update_stmt = (
                        update(table_class)
                        .where(and_(*conditions))
                        .values(**update_data)
                    )
                    await self.db.execute(update_stmt)
            else:
                # 插入
                insert_stmt = insert(table_class).values(**data)
                await self.db.execute(insert_stmt)

            processed_count += 1

        return processed_count

# 使用示例
# from sqlalchemy import create_engine
#
# # 创建引擎
# engine = create_engine('postgresql://user:pass@localhost/db')
# upsert = BatchUpsert(engine)
#
# # 数据示例
# employees = [
#     {'employee_id': 'EMP001', 'name': '张三', 'department': '技术部', 'salary': 15000},
#     {'employee_id': 'EMP002', 'name': '李四', 'department': '产品部', 'salary': 12000},
#     {'employee_id': 'EMP003', 'name': '王五', 'department': '设计部', 'salary': 13000},
# ]
#
# # 执行批量upsert
# result = upsert.batch_upsert_generic(employees, 'sys_basic_info', ['employee_id'])
# print(f"处理了 {result} 条记录")
