from __future__ import annotations

import asyncio
import copy
import datetime
import os
import queue
import time
from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple

import structlog
from apscheduler.schedulers.background import BackgroundScheduler
from jsonpatch import JsonPatchException
from jsonpointer import JsonPointerException

from rasa.core.tracker_stores.tracker_store import TrackerStore
from rasa.privacy.constants import (
    TEXT_KEY,
    USER_CHAT_INACTIVITY_IN_MINUTES_ENV_VAR_NAME,
)
from rasa.privacy.event_broker_utils import create_event_brokers
from rasa.privacy.privacy_config import (
    PrivacyConfig,
    PrivacyPolicy,
    validate_sensitive_slots,
)
from rasa.privacy.privacy_filter import PrivacyFilter
from rasa.shared.core.constants import ACTION_SESSION_START_NAME
from rasa.shared.core.events import (
    ActionExecuted,
    Event,
    SlotSet,
    UserUttered,
    split_events,
)
from rasa.shared.core.trackers import DialogueStateTracker, EventVerbosity
from rasa.shared.exceptions import RasaException

if TYPE_CHECKING:
    from asyncio import AbstractEventLoop

    from rasa.core.brokers.broker import EventBroker
    from rasa.core.config.available_endpoints import AvailableEndpoints
    from rasa.core.lock_store import LockStore
    from rasa.shared.core.domain import Domain


structlogger = structlog.get_logger(__name__)


def wrap_async(func: Callable) -> Callable:
    """Wraps a function to be used as an async job in the background scheduler."""

    def wrapper(*args: Any, **kwargs: Any) -> Any:
        return asyncio.run(func(*args, **kwargs))

    return wrapper


class BackgroundPrivacyManager:
    """Manages privacy-related tasks in the background.

    This class handles the anonymization and deletion of sensitive information
    in dialogue state trackers, as well as the streaming of anonymized events
    to event brokers. It uses background schedulers to periodically run these
    tasks and processes trackers from a queue to ensure that sensitive information
    is handled in a timely manner.
    """

    TRACKER_QUEUE_PROCESSING_TIMEOUT_IN_SECONDS = 2.0

    def __init__(
        self,
        endpoints: Optional["AvailableEndpoints"],
        event_loop: Optional["AbstractEventLoop"] = None,
        in_memory_tracker_store: Optional[TrackerStore] = None,
        lock_store: Optional["LockStore"] = None,
    ):
        if lock_store is None:
            raise RasaException(
                "LockStore is required for BackgroundPrivacyManager. "
                "Pass lock_store when creating the privacy manager."
            )
        self.lock_store = lock_store
        self.config = (
            PrivacyConfig.from_dict(endpoints.privacy)
            if endpoints and endpoints.privacy
            else None
        )
        self.privacy_filter = (
            PrivacyFilter(self.config.anonymization_rules) if self.config else None
        )
        self.user_chat_inactivity_in_minutes = int(
            os.getenv(USER_CHAT_INACTIVITY_IN_MINUTES_ENV_VAR_NAME, 30)
        )

        if in_memory_tracker_store is not None:
            # if an in-memory tracker store is provided,
            # we need to keep the reference to it
            # so that the background jobs can access it.
            # We also set the event broker to None
            # to prevent it from publishing events
            # during the tracker store background jobs
            in_memory_tracker_store.event_broker = None
            tracker_store = in_memory_tracker_store
        else:
            # we recreate the tracker store here to ensure
            # that this instance has no event brokers
            # that could publish events during the tracker store
            # background jobs
            tracker_store = (
                TrackerStore.create(endpoints.tracker_store)
                if endpoints
                else TrackerStore.create(None)
            )

        self.tracker_store = tracker_store

        self.event_brokers: List["EventBroker"] = []
        self.event_loop = event_loop

        # Order of the initialisation is important
        # The tracker queue must be created before the scheduler
        # The can_consume_tracker_queue must be set to True before the scheduler starts
        self.tracker_queue: queue.Queue = queue.Queue()

        # This flag is used to stop the scheduler
        self.can_consume_from_tracker_queue = True
        self.background_scheduler = BackgroundScheduler()
        self.background_scheduler.add_job(
            self._consumer_queue, max_instances=1, id="event_broker_job"
        )

        self.previous_fire_time_deletion = datetime.datetime.now(
            tz=datetime.timezone.utc
        )
        self._configure_background_scheduler()
        self.background_scheduler.start()

    async def initialize(
        self, endpoints: Optional["AvailableEndpoints"]
    ) -> BackgroundPrivacyManager:
        """Initialize async attributes of the BackgroundPrivacyManager."""
        event_broker_endpoints = endpoints.event_broker if endpoints else None
        self.event_brokers = (
            await create_event_brokers(event_broker_endpoints, self.event_loop)
            if event_broker_endpoints
            else []
        )

        return self

    @classmethod
    async def create_instance(
        cls,
        endpoints: Optional["AvailableEndpoints"],
        event_loop: Optional["AbstractEventLoop"] = None,
        in_memory_tracker_store: Optional[TrackerStore] = None,
        lock_store: Optional["LockStore"] = None,
    ) -> BackgroundPrivacyManager:
        """Create an instance of BackgroundPrivacyManager."""
        instance = cls(endpoints, event_loop, in_memory_tracker_store, lock_store)
        return await instance.initialize(endpoints)

    def stop(self) -> None:
        structlogger.debug("rasa.privacy_manager.stop_schedulers")
        self.can_consume_from_tracker_queue = False
        self.background_scheduler.shutdown(wait=False)

    def run(self, tracker: "DialogueStateTracker") -> None:
        self.tracker_queue.put(tracker)

    def process(
        self, tracker: "DialogueStateTracker", process_all: bool = False
    ) -> None:
        """Process the tracker to identify and anonymize sensitive information.

        Args:
            tracker: The tracker to process.
            process_all: If True, process all events in the tracker.
        """
        events = self.process_events(tracker, process_all=process_all)
        events_to_stream = events if events else tracker.events
        self.stream_events(events_to_stream, tracker.sender_id)

    def process_events_from_segment(
        self,
        events: List[Event],
        prior_sensitive_slot_events: Optional[List[Event]] = None,
    ) -> List[Event]:
        """Anonymize a segment of events without replaying into a tracker.

        Use this when the segment may contain dialogue stack updates that
        assume a different stack state (e.g. segments from the middle of a
        run). Replaying such segments into a tracker can raise when applying
        stack patches. Passing the events directly avoids that.
        """
        if not events:
            return []

        return self.privacy_filter.anonymize(  # type: ignore[union-attr]
            events, prior_sensitive_slot_events or []
        )

    def process_events(
        self, tracker: DialogueStateTracker, process_all: bool = False
    ) -> List[Event]:
        """Anonymize tracker events."""
        if (latest_message := self._get_latest_user_message(tracker)) is None:
            return []

        processed_events = list(tracker.events)
        prior_sensitive_slot_events: List[Event] = []

        if not process_all:
            additional_splitting_conditions = {
                TEXT_KEY: latest_message.text,
                "timestamp": latest_message.timestamp,
            }

            resulting_events = split_events(
                processed_events,
                UserUttered,
                additional_splitting_conditions=additional_splitting_conditions,
                include_splitting_event=True,
            )

            processed_events = resulting_events[1]
            prior_events = resulting_events[0]
            # Derive prior_sensitive_slot_events from raw prior_events instead of
            # replaying into a tracker (DialogueStateTracker.from_events). Replaying
            # can raise JsonPointerException when prior_events contain dialogue stack
            # patches that assume a different stack state (e.g. segments spanning
            # multiple sessions). Scanning events directly avoids that.
            prior_sensitive_slot_events = [
                event
                for event in prior_events
                if isinstance(event, SlotSet)
                and event.key in self.config.anonymization_rules  # type: ignore[union-attr]
            ]

        return self.process_events_from_segment(
            processed_events, prior_sensitive_slot_events
        )

    def stream_events(
        self,
        anonymized_events: List[Event],
        sender_id: str,
    ) -> None:
        """Stream anonymized events to the event broker."""
        if not self.event_brokers:
            structlogger.debug(
                "rasa.privacy_manager.no_event_broker_configured",
            )
            return None

        for event in anonymized_events:
            body = {"sender_id": sender_id}
            body.update(event.as_dict())
            for broker in self.event_brokers:
                broker.publish(body)

        return None

    def validate_sensitive_slots_in_domain(self, domain: "Domain") -> None:
        """Validate the sensitive slots defined in the privacy config against the domain."""  # noqa: E501
        if not self.config:
            structlogger.debug(
                "rasa.privacy_manager.no_sensitive_slots_configured",
            )
            return None

        # we need to set the domain in the tracker store
        # to prevent errors being raised about slots not found in the domain
        # during the background jobs
        self.tracker_store.domain = domain
        sensitive_slots = list(self.config.anonymization_rules.keys())
        return validate_sensitive_slots(sensitive_slots, domain)

    def _consumer_queue(self) -> None:
        while self.can_consume_from_tracker_queue:
            try:
                # Wait for 2 seconds for an event to be added to the queue
                # If no event is added to the queue, continue
                # This is done to avoid the scheduler to be stuck in the while loop
                # when we want to stop the scheduler
                tracker = self.tracker_queue.get(
                    timeout=self.TRACKER_QUEUE_PROCESSING_TIMEOUT_IN_SECONDS
                )
                self.process(tracker)
                self.tracker_queue.task_done()
            except queue.Empty:
                continue

    def _get_latest_user_message(
        self, tracker: DialogueStateTracker
    ) -> Optional[UserUttered]:
        """Check if a tracker should be processed."""
        if self.privacy_filter is None:
            structlogger.debug(
                "rasa.privacy_manager.no_privacy_rules_configured",
            )
            return None

        latest_user_message_event = tracker.get_last_event_for(
            UserUttered, event_verbosity=EventVerbosity.ALL
        )
        if latest_user_message_event is None or not isinstance(
            latest_user_message_event, UserUttered
        ):
            structlogger.debug(
                "rasa.privacy_manager.no_user_message.skipping_processing",
            )
            return None

        latest_user_message: UserUttered = latest_user_message_event
        if not latest_user_message.text:
            structlogger.debug(
                "rasa.privacy_manager.no_user_message.skipping_processing",
            )
            return None

        return latest_user_message

    @staticmethod
    def _has_session_been_anonymized(events: List[Event]) -> bool:
        """Check if the session has already been anonymized."""
        if not events:
            return False
        for event in reversed(events):
            if (
                hasattr(event, "anonymized_at")
                and getattr(event, "anonymized_at") is not None
            ):
                return True

        return False

    async def _run_tracker_store_anonymization(self) -> None:
        """Anonymize eligible tracker sessions in the tracker store."""
        structlogger.info(
            "rasa.privacy_manager.starting_tracker_store_anonymization",
            triggered_by="anonymization_cron_job",
        )

        keys = await self.tracker_store.keys()
        keys_copy = copy.deepcopy(list(keys))

        for key in keys_copy:
            async with self.lock_store.lock(key):
                try:
                    await self._process_one_key_anonymization(key)
                except Exception as e:
                    structlogger.error(
                        "rasa.privacy_manager.error_anonymizing_tracker",
                        sender_id=key,
                        error=str(e),
                    )
                    continue

    async def _run_tracker_store_deletion(self) -> None:
        """Delete eligible tracker sessions from the tracker store."""
        structlogger.info(
            "rasa.privacy_manager.starting_tracker_store_deletion",
            triggered_by="deletion_cron_job",
        )
        keys = await self.tracker_store.keys()

        # Make a copy of the keys to avoid modifying the list while iterating
        keys_copy = copy.deepcopy(list(keys))

        for key in keys_copy:
            async with self.lock_store.lock(key):
                try:
                    await self._process_one_key_deletion(key)
                except Exception as e:
                    structlogger.error(
                        "rasa.privacy_manager.error_deleting_tracker",
                        sender_id=key,
                        error=str(e),
                    )

    async def _run_tracker_store_background_jobs_sequentially(self) -> None:
        """Run the tracker store background jobs.

        If both anonymization and deletion policies are configured,
        we need to ensure that the background job timings do not
        overlap to prevent race conditions when accessing the
        tracker store.

        The scheduler will run the anonymization job first,
        and then the deletion job if the current time is past
        the next scheduled time for deletion.
        """
        await self._run_tracker_store_anonymization()

        now = datetime.datetime.now(tz=datetime.timezone.utc)
        next_fire_time = (
            self.config.tracker_store_settings.deletion_policy.cron.get_next_fire_time(  # type: ignore[union-attr]
                self.previous_fire_time_deletion,
                now=now,
            )
        )

        if next_fire_time and now >= next_fire_time:
            await self._run_tracker_store_deletion()
            self.previous_fire_time_deletion = next_fire_time

        return None

    def _add_anonymization_job(self) -> None:
        wrapped_anonymization = wrap_async(self._run_tracker_store_anonymization)
        self.background_scheduler.add_job(
            wrapped_anonymization,
            trigger=self.config.tracker_store_settings.anonymization_policy.cron,  # type: ignore[union-attr]
            max_instances=1,
            id="anonymization_cron_job",
        )

    def _add_deletion_job(self) -> None:
        wrapped_deletion = wrap_async(self._run_tracker_store_deletion)
        self.background_scheduler.add_job(
            wrapped_deletion,
            trigger=self.config.tracker_store_settings.deletion_policy.cron,  # type: ignore[union-attr]
            max_instances=1,
            id="deletion_cron_job",
        )

    def _add_sequential_job(self) -> None:
        sequential_dispatcher = wrap_async(
            self._run_tracker_store_background_jobs_sequentially
        )
        self.background_scheduler.add_job(
            sequential_dispatcher,
            trigger=self.config.tracker_store_settings.anonymization_policy.cron,  # type: ignore[union-attr]
            max_instances=1,
            id="anonymization_and_deletion_cron_job",
        )

    def _configure_background_scheduler(self) -> None:
        """Configure the background scheduler."""
        tracker_store_settings_configured = (
            self.config is not None and self.config.tracker_store_settings is not None
        )
        anonymization_policy = (
            self.config.tracker_store_settings.anonymization_policy  # type: ignore[union-attr]
            if tracker_store_settings_configured
            else None
        )
        deletion_policy = (
            self.config.tracker_store_settings.deletion_policy  # type: ignore[union-attr]
            if tracker_store_settings_configured
            else None
        )

        if (
            tracker_store_settings_configured
            and anonymization_policy is not None
            and deletion_policy is not None
        ):
            next_fire_time_anonymization = get_next_fire_time(anonymization_policy)
            next_fire_time_deletion = get_next_fire_time(deletion_policy)

            # If both anonymization and deletion policies are configured
            # to start on the same date,
            # we need to run them sequentially to avoid race conditions
            if (
                next_fire_time_anonymization is not None
                and next_fire_time_deletion is not None
                and next_fire_time_anonymization.date()
                == next_fire_time_deletion.date()
            ):
                self._add_sequential_job()
            else:
                self._add_anonymization_job()
                self._add_deletion_job()

        elif tracker_store_settings_configured and anonymization_policy is not None:
            self._add_anonymization_job()

        elif tracker_store_settings_configured and deletion_policy is not None:
            self._add_deletion_job()

    def _get_processed_events_after_anonymization(
        self,
        full_tracker: DialogueStateTracker,
    ) -> Tuple[List[Event], List[Event], List[Event]]:
        """Get processed events after anonymization job."""
        split_sessions = split_events(
            full_tracker.events,
            ActionExecuted,
            {"action_name": ACTION_SESSION_START_NAME},
            include_splitting_event=True,
        )

        processed_events = []
        already_anonymized_events = []
        uneligible_events = []

        for session_events in split_sessions:
            has_session_been_anonymized = self._has_session_been_anonymized(
                session_events
            )

            if has_session_been_anonymized:
                structlogger.debug(
                    "rasa.privacy_manager.session_already_anonymized",
                    sender_id=full_tracker.sender_id,
                )
                already_anonymized_events.extend(session_events)
                continue

            current_time = time.time()

            last_event_timestamp = (
                str(datetime.datetime.fromtimestamp(session_events[-1].timestamp))
                if session_events
                else "N/A"
            )

            if session_events and current_time - session_events[-1].timestamp > (
                self.user_chat_inactivity_in_minutes * 60
                + self.config.tracker_store_settings.anonymization_policy.min_after_session_end  # type: ignore[union-attr] # noqa: E501
                * 60
            ):
                structlogger.info(
                    "rasa.privacy_manager.anonymizing_tracker_session",
                    sender_id=full_tracker.sender_id,
                    last_event_timestamp=last_event_timestamp,
                    triggered_by="anonymization_cron_job",
                )
                prior_sensitive_slot_events: List[Event] = []
                events = self.process_events_from_segment(
                    session_events, prior_sensitive_slot_events
                )
                processed_events.extend(events)
            else:
                # If the session is not valid for anonymization,
                # we still want to write them back to the tracker store
                uneligible_events.extend(session_events)
                structlogger.debug(
                    "rasa.privacy_manager.session_not_valid_for_anonymization",
                    sender_id=full_tracker.sender_id,
                    last_event_timestamp=last_event_timestamp,
                )
        return processed_events, already_anonymized_events, uneligible_events

    def _get_events_to_be_retained_after_deletion(
        self, full_tracker: DialogueStateTracker
    ) -> List[Event]:
        """Get the events to be retained after deletion."""
        split_sessions = split_events(
            full_tracker.events,
            ActionExecuted,
            {"action_name": ACTION_SESSION_START_NAME},
            include_splitting_event=True,
        )
        events_to_be_retained: List[Event] = []
        for session_events in split_sessions:
            current_time = time.time()
            if session_events and (
                current_time - session_events[-1].timestamp
                <= (
                    self.user_chat_inactivity_in_minutes * 60
                    + self.config.tracker_store_settings.deletion_policy.min_after_session_end  # type: ignore[union-attr] # noqa: E501
                    * 60
                )
            ):
                events_to_be_retained.extend(session_events)
            else:
                last_event_timestamp = (
                    str(datetime.datetime.fromtimestamp(session_events[-1].timestamp))
                    if session_events
                    else "N/A"
                )

                structlogger.info(
                    "rasa.privacy_manager.tracker_session_scheduled_for_deletion",
                    sender_id=full_tracker.sender_id,
                    last_event_timestamp=last_event_timestamp,
                    triggered_by="deletion_cron_job",
                )

        return events_to_be_retained

    async def _process_one_key_anonymization(self, key: str) -> None:
        """Process one tracker for anonymization (call under lock)."""
        full_tracker = await self.tracker_store.retrieve_full_tracker(key)

        if not full_tracker:
            structlogger.debug(
                "rasa.privacy_manager.no_tracker_found_for_sender_id",
                sender_id=key,
            )
            return

        processed_events, already_anonymized_events, uneligible_events = (
            self._get_processed_events_after_anonymization(full_tracker)
        )

        if not processed_events:
            structlogger.debug(
                "rasa.privacy_manager.no_events_to_anonymize_for_tracker",
                sender_id=key,
            )
            return

        all_events = already_anonymized_events + processed_events + uneligible_events
        # Ensure chronological order so tracker replay is consistent.
        # Use index as tie-breaker for identical timestamps (preserves relative order).
        all_events = [
            e
            for _, e in sorted(
                enumerate(all_events), key=lambda item: (item[1].timestamp, item[0])
            )
        ]
        updated_tracker = DialogueStateTracker.from_events(
            sender_id=key,
            evts=all_events,
            slots=full_tracker.slots.values(),
        )
        await self.tracker_store.update(updated_tracker, apply_deletion_only=False)

        structlogger.info(
            "rasa.privacy_manager.saved_tracker_after_anonymization",
            sender_id=key,
        )

    async def _process_one_key_deletion(self, key: str) -> None:
        """Process one tracker for deletion (call under lock)."""
        full_tracker = await self.tracker_store.retrieve_full_tracker(key)

        if not full_tracker:
            structlogger.debug(
                "rasa.privacy_manager.no_tracker_found_for_sender_id",
                sender_id=key,
            )
            return

        events_to_be_retained = self._get_events_to_be_retained_after_deletion(
            full_tracker
        )

        if not events_to_be_retained:
            await self.tracker_store.delete(sender_id=key)
            structlogger.info(
                "rasa.privacy_manager.tracker_session_deleted",
                sender_id=full_tracker.sender_id,
                triggered_by="deletion_cron_job",
            )
            return

        try:
            tracker = DialogueStateTracker.from_events(
                sender_id=key,
                evts=events_to_be_retained,
                slots=full_tracker.slots.values(),
            )
        except (JsonPatchException, JsonPointerException) as exception:
            structlogger.error(
                "rasa.privacy_manager.error_reconstructing_tracker_after_deletion",
                sender_id=key,
                error=str(exception),
                event_info="Could not reconstruct tracker with events "
                "to be retained after deletion. "
                "To avoid data loss, the tracker will not be deleted. "
                "To reset the tracker's dialogue stack and enable "
                "deletion of sessions scheduled for deletion "
                "you can append a Restarted Event to the tracker.",
            )
            return

        await self.tracker_store.update(tracker)

        structlogger.info(
            "rasa.privacy_manager.overwritten_tracker",
            sender_id=key,
            event_info="Deleted eligible events and saved "
            "tracker with events not scheduled "
            "for deletion yet.",
        )


def get_next_fire_time(
    privacy_policy: PrivacyPolicy,
) -> Optional[datetime.datetime]:
    """Get the next fire time for the privacy policy."""
    now = datetime.datetime.now(tz=datetime.timezone.utc)
    return privacy_policy.cron.get_next_fire_time(None, now=now)
