from dataclasses import dataclass
from enum import Enum
from typing import List, Dict, Any, Optional, Callable
from datetime import datetime
import traceback
from appPublic.log import debug, error, info


class OperationType(Enum):
    """操作类型枚举"""
    FILE_LOAD = "file_load"
    EMBEDDING = "embedding"
    VDB_INSERT = "vdb_insert"
    TRIPLES_EXTRACT = "triples_extract"
    GDB_INSERT = "gdb_insert"
    VDB_DELETE = "vdb_delete"
    GDB_DELETE = "gdb_delete"
    ENTITY_EXTRACT = "entity_extract"
    TRIPLET_MATCH = "triplet_match"
    VECTOR_SEARCH = "vector_search"
    RERANK = "rerank"

@dataclass
class RollbackOperation:
    """回滚操作记录"""
    operation_type: OperationType
    data: Dict[str, Any]
    timestamp: str
    rollback_func: Optional[Callable] = None


class TransactionManager:
    """事务管理器"""

    def __init__(self):
        self.rollback_ops: List[RollbackOperation] = []
        self.transaction_id: str = datetime.now().strftime("%Y%m%d_%H%M%S_%f")

    def add_operation(self, operation_type: OperationType, data: Dict[str, Any],
                      rollback_func: Optional[Callable] = None):
        """添加操作记录"""
        operation = RollbackOperation(
            operation_type=operation_type,
            data=data,
            timestamp=datetime.now().isoformat(),
            rollback_func=rollback_func
        )
        self.rollback_ops.append(operation)
        debug(f"事务 {self.transaction_id} 添加操作: {operation_type.value}")

    def get_operations_count(self) -> int:
        """获取操作数量"""
        return len(self.rollback_ops)

    async def rollback_all(self, rollback_context: Dict[str, Any] = None) -> List[str]:
        """执行所有回滚操作"""
        debug(f"事务 {self.transaction_id} 开始回滚，共 {len(self.rollback_ops)} 个操作")
        rollback_results = []
        rollback_context = rollback_context or {}

        # 按相反顺序执行回滚
        for op in reversed(self.rollback_ops):
            try:
                debug(f"回滚操作: {op.operation_type.value}")

                if op.rollback_func:
                    # 验证上下文和数据
                    required_context_keys = ['request', 'service_params', 'userid']
                    missing_context = [k for k in required_context_keys if
                                       k not in rollback_context or rollback_context[k] is None]
                    if missing_context:
                        raise ValueError(f"回滚上下文缺少字段: {', '.join(missing_context)}")

                    required_data_keys = ['orgid', 'realpath', 'fiid', 'id', 'db_type']
                    if op.operation_type == OperationType.VDB_INSERT:
                        missing_data = [k for k in required_data_keys if k not in op.data or op.data[k] is None]
                        if missing_data:
                            raise ValueError(f"VDB_INSERT 数据缺少字段: {', '.join(missing_data)}")

                    # 执行自定义回滚函数
                    result = await op.rollback_func(op.data, rollback_context)
                    rollback_results.append(f"成功回滚 {op.operation_type.value}: {result}")
                else:
                    # 默认回滚处理
                    result = await self._default_rollback(op, rollback_context)
                    rollback_results.append(f"成功回滚 {op.operation_type.value}: {result}")

            except Exception as e:
                error(f"回滚操作失败 {op.operation_type.value}: {str(e)}, 堆栈: {traceback.format_exc()}")
                rollback_results.append(f"回滚失败: {op.operation_type.value} - {str(e)}")

        return rollback_results

    async def _default_rollback(self, op: RollbackOperation, context: Dict[str, Any]) -> str:
        """默认回滚处理"""
        if op.operation_type in [OperationType.FILE_LOAD, OperationType.EMBEDDING,
                                 OperationType.TRIPLES_EXTRACT, OperationType.ENTITY_EXTRACT]:
            return "内存操作，已自动清理"
        elif op.operation_type in [OperationType.VDB_INSERT, OperationType.VDB_DELETE,
                                   OperationType.GDB_INSERT, OperationType.GDB_DELETE]:
            return "需要在具体实现中处理数据库回滚"
        else:
            return f"跳过回滚操作: {op.operation_type.value}"

    def clear(self):
        """清空事务记录"""
        self.rollback_ops.clear()
        debug(f"事务 {self.transaction_id} 已清空")


class TransactionContext:
    """事务上下文管理器"""

    def __init__(self, transaction_name: str = ""):
        self.transaction_name = transaction_name
        self.transaction_manager = TransactionManager()
        self.start_time = None
        self.end_time = None

    async def __aenter__(self):
        self.start_time = datetime.now()
        info(f"开始事务: {self.transaction_name} [{self.transaction_manager.transaction_id}]")
        return self.transaction_manager

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self.end_time = datetime.now()
        duration = (self.end_time - self.start_time).total_seconds()

        if exc_type is not None:
            # 发生异常，执行回滚
            error(f"事务失败: {self.transaction_name}, 异常: {exc_val}")
            # 使用 transaction_manager 的 rollback_context
            rollback_context = getattr(self.transaction_manager, 'transaction_context', {})
            rollback_results = await self.transaction_manager.rollback_all(rollback_context)
            error(f"回滚结果: {rollback_results}")
            info(f"事务回滚完成: {self.transaction_name}, 耗时: {duration:.2f}秒")
        else:
            # 正常完成
            info(f"事务成功: {self.transaction_name}, 耗时: {duration:.2f}秒")

        self.transaction_manager.clear()


def create_transaction_manager() -> TransactionManager:
    """创建事务管理器"""
    return TransactionManager()


def create_transaction_context(name: str = "") -> TransactionContext:
    """创建事务上下文"""
    return TransactionContext(name)