"""
Transaction Decorators and Context Managers
Provides convenient transaction handling for service methods
"""

from functools import wraps
from typing import Callable, Any, Optional
from contextlib import asynccontextmanager
import structlog

from .base import DatabaseManager
from .unit_of_work import UnitOfWork
from .exceptions import TransactionError

logger = structlog.get_logger()


def transactional(database_manager: DatabaseManager, auto_commit: bool = True):
    """
    Decorator that wraps a method in a database transaction
    
    Args:
        database_manager: DatabaseManager instance
        auto_commit: Whether to auto-commit on success
    
    Usage:
        @transactional(database_manager)
        async def create_user_with_profile(self, user_data, profile_data):
            # Your business logic here
            # Transaction is automatically managed
            pass
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            async with database_manager.get_background_session() as session:
                try:
                    # Inject session into kwargs if not present
                    if 'session' not in kwargs:
                        kwargs['session'] = session
                    
                    result = await func(*args, **kwargs)
                    
                    # Session is auto-committed by get_background_session
                    logger.debug(f"Transaction completed successfully for {func.__name__}")
                    return result
                    
                except Exception as e:
                    # Session is auto-rolled back by get_background_session
                    logger.error(f"Transaction failed for {func.__name__}", error=str(e))
                    raise TransactionError(f"Transaction failed: {str(e)}")
        
        return wrapper
    return decorator


def unit_of_work_transactional(database_manager: DatabaseManager):
    """
    Decorator that provides Unit of Work pattern for complex operations
    
    Usage:
        @unit_of_work_transactional(database_manager)
        async def complex_business_operation(self, data, uow: UnitOfWork):
            user_repo = uow.register_repository("users", UserRepository, User)
            sales_repo = uow.register_repository("sales", SalesRepository, SalesData)
            
            user = await user_repo.create(data.user)
            sale = await sales_repo.create(data.sale)
            
            # UnitOfWork automatically commits
            return {"user": user, "sale": sale}
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            async with database_manager.get_background_session() as session:
                async with UnitOfWork(session, auto_commit=True) as uow:
                    try:
                        # Inject UnitOfWork into kwargs
                        kwargs['uow'] = uow
                        
                        result = await func(*args, **kwargs)
                        
                        logger.debug(f"Unit of Work transaction completed for {func.__name__}")
                        return result
                        
                    except Exception as e:
                        logger.error(f"Unit of Work transaction failed for {func.__name__}", 
                                   error=str(e))
                        raise TransactionError(f"Transaction failed: {str(e)}")
        
        return wrapper
    return decorator


@asynccontextmanager
async def managed_transaction(database_manager: DatabaseManager):
    """
    Context manager for explicit transaction control
    
    Usage:
        async with managed_transaction(database_manager) as session:
            # Your database operations here
            user = User(name="John")
            session.add(user)
            # Auto-commits on exit, rolls back on exception
    """
    async with database_manager.get_background_session() as session:
        try:
            logger.debug("Starting managed transaction")
            yield session
            logger.debug("Managed transaction completed successfully")
        except Exception as e:
            logger.error("Managed transaction failed", error=str(e))
            raise


@asynccontextmanager
async def managed_unit_of_work(database_manager: DatabaseManager, event_publisher=None):
    """
    Context manager for explicit Unit of Work control
    
    Usage:
        async with managed_unit_of_work(database_manager) as uow:
            user_repo = uow.register_repository("users", UserRepository, User)
            user = await user_repo.create(user_data)
            await uow.commit()
    """
    async with database_manager.get_background_session() as session:
        uow = UnitOfWork(session)
        try:
            logger.debug("Starting managed Unit of Work")
            yield uow
            
            if not uow._committed:
                await uow.commit()
            
            logger.debug("Managed Unit of Work completed successfully")
            
        except Exception as e:
            if not uow._rolled_back:
                await uow.rollback()
            logger.error("Managed Unit of Work failed", error=str(e))
            raise


class TransactionManager:
    """
    Advanced transaction manager for complex scenarios
    
    Usage:
        tx_manager = TransactionManager(database_manager)
        
        async with tx_manager.create_transaction() as tx:
            await tx.execute_in_transaction(my_business_logic, data)
    """
    
    def __init__(self, database_manager: DatabaseManager):
        self.database_manager = database_manager
    
    @asynccontextmanager
    async def create_transaction(self, isolation_level: Optional[str] = None):
        """Create a transaction with optional isolation level"""
        async with self.database_manager.get_background_session() as session:
            transaction_context = TransactionContext(session, isolation_level)
            try:
                yield transaction_context
            except Exception as e:
                logger.error("Transaction manager failed", error=str(e))
                raise
    
    async def execute_with_retry(
        self, 
        func: Callable, 
        max_retries: int = 3, 
        *args, 
        **kwargs
    ):
        """Execute function with transaction retry on failure"""
        last_error = None
        
        for attempt in range(max_retries):
            try:
                async with managed_transaction(self.database_manager) as session:
                    kwargs['session'] = session
                    result = await func(*args, **kwargs)
                    logger.debug(f"Transaction succeeded on attempt {attempt + 1}")
                    return result
                    
            except Exception as e:
                last_error = e
                logger.warning(f"Transaction attempt {attempt + 1} failed", 
                             error=str(e), remaining_attempts=max_retries - attempt - 1)
                
                if attempt == max_retries - 1:
                    break
        
        logger.error(f"All transaction attempts failed after {max_retries} tries")
        raise TransactionError(f"Transaction failed after {max_retries} retries: {str(last_error)}")


class TransactionContext:
    """Context for managing individual transactions"""
    
    def __init__(self, session, isolation_level: Optional[str] = None):
        self.session = session
        self.isolation_level = isolation_level
    
    async def execute_in_transaction(self, func: Callable, *args, **kwargs):
        """Execute function within the transaction context"""
        try:
            kwargs['session'] = self.session
            result = await func(*args, **kwargs)
            return result
        except Exception as e:
            logger.error("Function execution failed in transaction context", error=str(e))
            raise


# ===== UTILITY FUNCTIONS =====

async def run_in_transaction(database_manager: DatabaseManager, func: Callable, *args, **kwargs):
    """
    Utility function to run any async function in a transaction
    
    Usage:
        result = await run_in_transaction(
            database_manager, 
            my_async_function, 
            arg1, arg2, 
            kwarg1="value"
        )
    """
    async with managed_transaction(database_manager) as session:
        kwargs['session'] = session
        return await func(*args, **kwargs)


async def run_with_unit_of_work(
    database_manager: DatabaseManager, 
    func: Callable, 
    *args, 
    **kwargs
):
    """
    Utility function to run any async function with Unit of Work
    
    Usage:
        result = await run_with_unit_of_work(
            database_manager, 
            my_complex_function, 
            arg1, arg2
        )
    """
    async with managed_unit_of_work(database_manager) as uow:
        kwargs['uow'] = uow
        return await func(*args, **kwargs)


# ===== BATCH OPERATIONS =====

@asynccontextmanager
async def batch_operation(database_manager: DatabaseManager, batch_size: int = 1000):
    """
    Context manager for batch operations with automatic commit batching
    
    Usage:
        async with batch_operation(database_manager, batch_size=500) as batch:
            for item in large_dataset:
                await batch.add_operation(create_record, item)
    """
    async with database_manager.get_background_session() as session:
        batch_context = BatchOperationContext(session, batch_size)
        try:
            yield batch_context
            await batch_context.flush_remaining()
        except Exception as e:
            logger.error("Batch operation failed", error=str(e))
            raise


class BatchOperationContext:
    """Context for managing batch database operations"""
    
    def __init__(self, session, batch_size: int):
        self.session = session
        self.batch_size = batch_size
        self.operation_count = 0
    
    async def add_operation(self, func: Callable, *args, **kwargs):
        """Add operation to batch"""
        kwargs['session'] = self.session
        await func(*args, **kwargs)
        
        self.operation_count += 1
        
        if self.operation_count >= self.batch_size:
            await self.session.commit()
            self.operation_count = 0
            logger.debug(f"Batch committed at {self.batch_size} operations")
    
    async def flush_remaining(self):
        """Commit any remaining operations"""
        if self.operation_count > 0:
            await self.session.commit()
            logger.debug(f"Final batch committed with {self.operation_count} operations")