"""
Database-level locking for preventing concurrent execution.
"""

import sys
from datetime import date
from django.db import transaction
from django_cronjob_utils.models import CronExecution
from django_cronjob_utils.exceptions import ConcurrentExecutionError


class DatabaseLock:
    """Context manager for database-level locking."""
    
    def __init__(self, task_code: str, execution_date: date):
        self.task_code = task_code
        self.execution_date = execution_date
    
    def __enter__(self):
        # Use atomic to ensure we're in a transaction for SELECT FOR UPDATE
        # The transaction will remain open until __exit__ is called
        self._transaction = transaction.atomic()
        self._transaction.__enter__()
        
        try:
            # Lock any running executions for this task/date
            running = CronExecution.objects.select_for_update().filter(
                task_code=self.task_code,
                execution_date=self.execution_date,
                completed=False
            ).first()
            
            if running:
                # Exit transaction before raising exception
                self._transaction.__exit__(None, None, None)
                self._transaction = None
                raise ConcurrentExecutionError(
                    f"Task {self.task_code} already running for {self.execution_date} "
                    f"(started at {running.started})"
                )
        except ConcurrentExecutionError:
            # Re-raise ConcurrentExecutionError without exiting transaction again
            if self._transaction:
                self._transaction.__exit__(None, None, None)
                self._transaction = None
            raise
        except Exception:
            # For other exceptions, properly exit transaction
            if self._transaction:
                self._transaction.__exit__(*sys.exc_info())
                self._transaction = None
            raise
        
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # Commit or rollback the transaction, releasing the lock
        if hasattr(self, '_transaction') and self._transaction is not None:
            self._transaction.__exit__(exc_type, exc_val, exc_tb)
            self._transaction = None