"""
Core state management functionality for Label Studio.

Provides high-performance state management with caching and batch operations
that can be extended by Label Studio Enterprise with additional features.
"""

import logging
from datetime import datetime
from typing import Any, Dict, List, Optional, Type

from core.current_request import CurrentContext
from core.feature_flags import flag_set
from django.conf import settings
from django.core.cache import cache, caches
from django.db.models import Model, QuerySet
from fsm.registry import get_state_model_for_entity
from fsm.state_models import BaseState
from fsm.transition_executor import execute_transition_with_state_manager

logger = logging.getLogger(__name__)


_fsm_cache = None


def get_fsm_cache():
    """
    Get the cache backend for FSM operations.

    Uses REDIS_CACHE_ALIAS when available (LSE) to ensure FSM caching works
    even when DUMMY_CACHE is enabled. Falls back to default cache for LSO.

    The result is memoized so the cache lookup only happens once per process.

    Returns:
        Cache backend instance
    """
    global _fsm_cache
    if _fsm_cache is not None:
        return _fsm_cache

    redis_cache_alias = getattr(settings, 'REDIS_CACHE_ALIAS', None)
    if redis_cache_alias and redis_cache_alias in settings.CACHES:
        _fsm_cache = caches[redis_cache_alias]
    else:
        _fsm_cache = cache
    return _fsm_cache


class StateManagerError(Exception):
    """Base exception for StateManager operations"""

    pass


class InvalidTransitionError(StateManagerError):
    """Raised when an invalid state transition is attempted"""

    pass


class StateManager:
    """
    Core state management system for Label Studio.

    Provides the foundation for state management that can be extended
    by Label Studio Enterprise with additional features like:
    - Advanced caching strategies
    - Bulk operations optimization
    - Complex transition validation
    - Enterprise-specific state models

    Features:
    - INSERT-only architecture with UUID7 for maximum performance
    - Basic caching for current state lookups
    - Simple state transitions with audit trails
    - Extensible design for enterprise features
    """

    CACHE_TTL = getattr(settings, 'FSM_CACHE_TTL', 300)  # 5 minutes default
    CACHE_PREFIX = 'fsm:current'

    @classmethod
    def clear_fsm_cache(cls):
        """
        Clear all FSM-related cache keys.

        Uses delete_pattern if available (django-redis), otherwise logs a warning.
        This is primarily used for test isolation.
        """
        fsm_cache = get_fsm_cache()
        pattern = f'{cls.CACHE_PREFIX}:*'
        if hasattr(fsm_cache, 'delete_pattern'):
            fsm_cache.delete_pattern(pattern)
        else:
            logger.warning(
                'FSM cache clear requested but cache backend does not support delete_pattern. '
                'FSM cache keys may persist.'
            )

    @classmethod
    def _is_fsm_enabled(cls, user='auto') -> bool:
        """Check if FSM feature is enabled via feature flag."""
        if user is None or user == 'auto':
            return CurrentContext.is_fsm_enabled()
        return flag_set('fflag_feat_fit_568_finite_state_management', user=user)

    @classmethod
    def get_cache_key(cls, entity: Model) -> str:
        """Generate cache key for entity's current state"""
        return f'{cls.CACHE_PREFIX}:{entity._meta.label_lower}:{entity.pk}'

    @classmethod
    def get_current_state_value(cls, entity: Model) -> Optional[str]:
        """
        Get current state with basic caching.

        Args:
            entity: The entity to get current state for

        Returns:
            Current state string
        Raises:
            StateManagerError: If no state model found

        Example:
            task = Task.objects.get(id=123)
            current_state = StateManager.get_current_state_value(task)
            if current_state == 'COMPLETED':
                # Task is finished
                pass
        """
        if not cls._is_fsm_enabled():
            return None  # Feature disabled, return no state

        cache_key = cls.get_cache_key(entity)
        fsm_cache = get_fsm_cache()

        # Try cache first
        cached_state = fsm_cache.get(cache_key)
        if cached_state is not None:
            logger.info(
                'FSM: Cache hit',
                extra={
                    'event': 'fsm.cache_hit',
                    'entity_type': entity._meta.label_lower,
                    'entity_id': entity.pk,
                    'organization_id': CurrentContext.get_organization_id(),
                    'state': cached_state,
                },
            )
            return cached_state

        # Query database using state model registry
        state_model = get_state_model_for_entity(entity)
        if not state_model:
            raise StateManagerError(f'No state model found for {entity._meta.model_name} when getting current state')

        try:
            current_state = state_model.get_current_state_value(entity)

            # Cache result
            if current_state is not None:
                fsm_cache.set(cache_key, current_state, cls.CACHE_TTL)
                logger.info(
                    'FSM: Cache miss',
                    extra={
                        'event': 'fsm.cache_miss',
                        'entity_type': entity._meta.label_lower,
                        'entity_id': entity.pk,
                        'organization_id': CurrentContext.get_organization_id(),
                    },
                )

            return current_state

        except Exception as e:
            logger.error(
                'FSM: Error getting current state',
                extra={
                    'event': 'fsm.get_state_error',
                    'entity_type': entity._meta.label_lower,
                    'entity_id': entity.pk,
                    'organization_id': CurrentContext.get_organization_id(),
                    'error': str(e),
                },
                exc_info=True,
            )
            raise StateManagerError(f'Error getting current state: {e}') from e

    @classmethod
    def get_current_state_object(cls, entity: Model) -> BaseState:
        """
        Get current state object with full audit information.

        Args:
            entity: The entity to get current state object for

        Returns:
            Latest BaseState instance

        Raises:
            StateManagerError: If no state model found
        """
        state_model = get_state_model_for_entity(entity)
        if not state_model:
            raise StateManagerError(
                f'No state model found for {entity._meta.model_name} when getting current state object'
            )

        return state_model.get_current_state(entity)

    @classmethod
    def transition_state(
        cls,
        entity: Model,
        new_state: str,
        transition_name: str = None,
        user=None,
        organization_id=None,
        context: Dict[str, Any] = None,
        reason: str = '',
        force_state_record: bool = False,
    ) -> bool:
        """
        Perform state transition with audit trail.

        Uses INSERT-only approach for maximum performance:
        - No UPDATE operations or row locks
        - Complete audit trail by design
        - Basic cache update for consistency

        Args:
            entity: The entity to transition
            new_state: Target state
            transition_name: Name of transition method (for audit)
            user: User triggering the transition
            organization_id: Organization ID
            context: Additional context data
            reason: Human-readable reason for transition
            force_state_record: If True, creates state record even if state doesn't change (for audit trails)

        Returns:
            True if transition succeeded, False otherwise

        Raises:
            InvalidTransitionError: If transition is not valid
            StateManagerError: If transition fails

        Example:
            success = StateManager.transition_state(
                entity=task,
                new_state='IN_PROGRESS',
                transition_name='start_annotation',
                user=request.user,
                organization_id=request.user.active_organization_id,
                context={'assignment_id': assignment.id},
                reason='User started annotation work'
            )
        """
        if not cls._is_fsm_enabled(user=user):
            return True  # Feature disabled, silently succeed

        # Skip if FSM is temporarily disabled (e.g., during cleanup or bulk operations)
        if CurrentContext.is_fsm_disabled():
            return True  # FSM disabled, silently succeed

        state_model = get_state_model_for_entity(entity)
        if not state_model:
            raise StateManagerError(f'No state model found for {entity._meta.model_name} when transitioning state')

        current_state = cls.get_current_state_value(entity)

        # Prevent same-state transitions - only create state records for actual state changes
        # This avoids creating redundant data when the effective state doesn't change
        # However, allow forced state records for audit trails (e.g., annotation updates)
        # IMPORTANT: Also check if a state record exists in DB - if not, we must create one
        # even if inferred state matches target state (to persist the inferred state)
        if current_state == new_state and not force_state_record:
            # Verify a state record actually exists in DB (not just inferred)
            state_record_exists = state_model.objects.filter(**{entity._meta.model_name: entity}).exists()
            if state_record_exists:
                return True  # Skip transition - record exists and state unchanged
            # else: No record exists (state was inferred), continue to create record

        # Optimistic concurrency control using cache-based locking
        cache_key = cls.get_cache_key(entity)
        lock_key = f'{cache_key}:lock'
        fsm_cache = get_fsm_cache()

        if organization_id is None:
            organization_id = CurrentContext.get_organization_id()

        try:
            # Try to acquire an optimistic lock using cache add (atomic operation)
            # add() only succeeds if the key doesn't exist
            lock_acquired = fsm_cache.add(lock_key, 'locked', timeout=5)  # 5 second timeout

            if not lock_acquired:
                # Another process is currently transitioning this entity
                logger.info(
                    'FSM: Concurrent transition detected, skipping',
                    extra={
                        'event': 'fsm.concurrent_transition_skipped',
                        'entity_type': entity._meta.label_lower,
                        'entity_id': entity.pk,
                        'target_state': new_state,
                        'organization_id': organization_id,
                    },
                )
                return True

            try:
                # INSERT-only approach - no UPDATE operations needed
                # Get denormalized fields from the state model class
                denormalized_fields = state_model.get_denormalized_fields(entity)

                # Get organization from entity or denormalized fields, or user's active organization
                if organization_id is None:
                    organization_id = getattr(
                        entity, 'organization_id', getattr(denormalized_fields, 'organization_id', None)
                    )
                    if organization_id is not None:
                        CurrentContext.set_organization_id(organization_id)

                if not organization_id and user and hasattr(user, 'active_organization') and user.active_organization:
                    organization_id = user.active_organization.id
                    if organization_id is not None:
                        CurrentContext.set_organization_id(organization_id)

                logger.info(
                    'FSM: State transition starting',
                    extra={
                        'event': 'fsm.transition_state_start',
                        'entity_type': entity._meta.label_lower,
                        'entity_id': entity.pk,
                        'from_state': current_state,
                        'to_state': new_state,
                        'transition_name': transition_name,
                        **{
                            'user_id': user.id if user else None,
                            'organization_id': organization_id if organization_id else None,
                        },
                    },
                )

                # CRITICAL FIX: Use state model's correct field name instead of entity._meta.model_name
                # This fixes the architectural entity field mapping issue where entity._meta.model_name
                # doesn't always match the actual field name defined in FSM state models
                entity_field_name = state_model._get_entity_field_name()

                new_state_record = state_model.objects.create(
                    **{entity_field_name: entity},
                    state=new_state,
                    previous_state=current_state,
                    transition_name=transition_name,
                    triggered_by=user if getattr(user, 'is_authenticated', False) else None,
                    context_data=context or {},
                    reason=reason,
                    organization_id=organization_id,
                    **denormalized_fields,
                )

                # Write-through cache: Update immediately
                # This ensures the cache is updated atomically with the database
                fsm_cache.set(cache_key, new_state, cls.CACHE_TTL)

                logger.info(
                    'FSM: Cache updated for transition state',
                    extra={
                        'event': 'fsm.transition_state_cache_updated',
                        'entity_type': entity._meta.label_lower,
                        'entity_id': entity.pk,
                        'state': new_state,
                        **{
                            'user_id': user.id if user else None,
                            'organization_id': organization_id if organization_id else None,
                        },
                    },
                )

                logger.info(
                    'FSM: State transition successful',
                    extra={
                        'event': 'fsm.transition_state_success',
                        'entity_type': entity._meta.label_lower,
                        'entity_id': entity.pk,
                        'state': new_state,
                        'state_record_id': str(new_state_record.id),
                        **{
                            'user_id': user.id if user else None,
                            'organization_id': organization_id if organization_id else None,
                        },
                    },
                )
                return True

            finally:
                # Always release the lock, regardless of success or failure
                fsm_cache.delete(lock_key)

        except Exception as e:
            # On failure, clean up lock and invalidate potentially stale cache
            fsm_cache.delete(lock_key)
            fsm_cache.delete(cache_key)

            # Get organization_id for error logging if it wasn't set earlier
            organization_id = CurrentContext.get_organization_id()

            logger.error(
                'FSM: State transition failed',
                extra={
                    'event': 'fsm.transition_state_failed',
                    'entity_type': entity._meta.label_lower,
                    'entity_id': entity.pk,
                    'from_state': current_state,
                    'to_state': new_state,
                    'error': str(e),
                    **{
                        'user_id': user.id if user else None,
                        'organization_id': organization_id if organization_id else None,
                    },
                },
                exc_info=True,
            )
            raise StateManagerError(f'Failed to transition state: {e}') from e

    @classmethod
    def get_state_history(cls, entity: Model) -> QuerySet[BaseState]:
        """
        Get complete state history for an entity.

        Args:
            entity: Entity to get history for

        Returns:
            QuerySet of state records ordered by most recent first
        """
        state_model = get_state_model_for_entity(entity)
        if not state_model:
            raise StateManagerError(
                f'No state model registered for {entity._meta.model_name} when getting state history'
            )

        return state_model.get_state_history(entity)

    @classmethod
    def get_states_in_time_range(
        cls, entity: Model, start_time: datetime, end_time: Optional[datetime] = None
    ) -> List[BaseState]:
        """
        Get states within a time range using UUID7 time-based queries.

        Args:
            entity: Entity to get states for
            start_time: Start of time range
            end_time: End of time range (defaults to now)

        Returns:
            List of states within the time range
        """
        state_model = get_state_model_for_entity(entity)
        if not state_model:
            raise StateManagerError(
                f'No state model registered for {entity._meta.model_name} when getting states in time range'
            )

        return state_model.get_states_in_range(entity, start_time, end_time or datetime.now())

    @classmethod
    def invalidate_cache(cls, entity: Model):
        """Invalidate cached state for an entity"""
        cache_key = cls.get_cache_key(entity)
        fsm_cache = get_fsm_cache()
        fsm_cache.delete(cache_key)
        organization_id = CurrentContext.get_organization_id()
        logger.info(
            'FSM: Cache invalidated',
            extra={
                'event': 'fsm.cache_invalidated',
                'entity_type': entity._meta.label_lower,
                'entity_id': entity.pk,
                **{'organization_id': organization_id if organization_id else None},
            },
        )

    @classmethod
    def warm_cache(cls, entities: List[Model]):
        """
        Warm cache with current states for a list of entities.

        Basic implementation that can be optimized by Enterprise with
        bulk queries and advanced caching strategies.
        """
        cache_updates = {}
        organization_id = CurrentContext.get_organization_id()
        for entity in entities:
            current_state = cls.get_current_state_value(entity)
            if current_state:
                cache_key = cls.get_cache_key(entity)
                cache_updates[cache_key] = current_state

        if cache_updates:
            fsm_cache = get_fsm_cache()
            fsm_cache.set_many(cache_updates, cls.CACHE_TTL)
            logger.info(
                'FSM: Cache warmed',
                extra={
                    'event': 'fsm.cache_warmed',
                    'entity_count': len(cache_updates),
                    **{'organization_id': organization_id if organization_id else None},
                },
            )

    @classmethod
    def execute_transition(
        cls, entity: Model, transition_name: str, transition_data: Dict[str, Any] = None, user=None, **context_kwargs
    ) -> BaseState:
        """
        Execute a registered transition by name.

        This is the main entry point for all state transitions using the declarative system.
        Enterprise implementations can override this method to add additional behavior.

        Args:
            entity: The entity to transition
            transition_name: Name of the registered transition
            transition_data: Data for the transition (validated by Pydantic)
            user: User executing the transition
            **context_kwargs: Additional context data

        Returns:
            The newly created state record

        Raises:
            ValueError: If transition is not found
            TransitionValidationError: If transition validation fails
        """
        # Delegate to transition executor, passing StateManager methods as parameters
        return execute_transition_with_state_manager(
            entity=entity,
            transition_name=transition_name,
            transition_data=transition_data,
            user=user,
            state_manager_class=cls,
            **context_kwargs,
        )


# Allow runtime configuration of which StateManager to use
# Enterprise can set this to their extended implementation
DEFAULT_STATE_MANAGER = StateManager
RESOLVED_STATE_MANAGER = None


def get_state_manager() -> Type[StateManager]:
    """
    Get the configured state manager class.

    Returns the StateManager class to use. Enterprise can override
    this by setting a different class in their configuration.
    """
    # Resolve once
    if RESOLVED_STATE_MANAGER is not None:
        return RESOLVED_STATE_MANAGER

    # Check if enterprise has configured a custom state manager
    if hasattr(settings, 'FSM_STATE_MANAGER_CLASS'):
        manager_path = settings.FSM_STATE_MANAGER_CLASS
        module_name, class_name = manager_path.rsplit('.', 1)
        module = __import__(module_name, fromlist=[class_name])
        return getattr(module, class_name)

    return DEFAULT_STATE_MANAGER