"""On-disk SQLite message store used by the replay flow to persist messages
that failed to reach the Opik server. DBManager exposes helpers for the full
message lifecycle: register_message(s), fetch_failed_messages_batched, and
replay_failed_messages (which accepts a ReplayCallback, typically Streamer.put,
to re-inject messages). The manager tracks three states — undefined, initialized, closed,
and error — and if the underlying database becomes unavailable, it marks itself
as failed (error) and logs that resiliency features are disabled. The standalone helper
db_message_to_message raises ValueError for unsupported message types.
"""

import logging
import os
import shutil
import sqlite3
import tempfile
import threading
import time
from enum import unique, IntEnum
from typing import List, NamedTuple, Optional, Iterator

from opik.message_processing import messages
from opik.message_processing.replay import message_serialization

from .types import ReplayCallback


DEFAULT_DB_FILE = "opik_messages.db"

LOGGER = logging.getLogger(__name__)


@unique
class MessageStatus(IntEnum):
    """
    Represents the status of a message.

    This enumeration defines the potential statuses that a message can
    have in a messaging system. Used for categorizing messages based on
    their lifecycle stage or outcome.

    Attributes:
        registered: Represents a message that has been successfully
            registered but not yet processed or delivered.
        delivered : Represents a message that has been successfully
            delivered to its recipient.
        failed: Represents a message for which delivery has failed.
    """

    registered = 1
    delivered = 2
    failed = 3


@unique
class DBManagerStatus(IntEnum):
    """Represents various status values for a manager.

    Defines a set of distinct states that a manager can occupy during its lifecycle.
    This can be used to track and manage the status of a manager in an application.
    """

    undefined = 1
    initialized = 2
    closed = 3
    error = 4


class DBMessage(NamedTuple):
    """
    Represents a database message entity.

    This class is used to encapsulate information about messages stored in
    or retrieved from a database. It provides a structured format for handling
    message data, including its unique identifier, type, JSON content, and
    current status.

    Attributes:
        id: The unique identifier of the message.
        type: The type/category of the message.
        json: The JSON-encoded content of the message.
        status: The current status of the message.
    """

    id: int
    type: str
    json: str
    status: MessageStatus


def _preprocess_registered_message(message: messages.BaseMessage) -> str:
    if message.message_id is None:
        raise ValueError("Message ID expected")

    return message_serialization.serialize_message(message)


class DBManager:
    """
    Manages message storage, batch processing, and database operations within a replay
    system.

    This class provides functionalities to handle message registrations, batch updates,
    and database schema management. It ensures thread-safe operations, cleans up
    temporary files, and handles database connections as part of its lifecycle.

    """

    def __init__(
        self,
        batch_size: int,
        batch_replay_delay: float,
        db_file: Optional[str] = None,
        conn: Optional[sqlite3.Connection] = None,
        sync_lock: Optional[threading.RLock] = None,
    ) -> None:
        """
        Initializes the Manager class, setting up the database connection, temporary
        directory, and other necessary properties for managing batch operations.

        Args:
            batch_size: The size of batches for processing.
            batch_replay_delay: The delay (in seconds) between replaying batches of messages.
            db_file: Path to the database file. If not provided, a
                temporary file will be created in a temporary directory.
            conn: A pre-existing database connection.
                If not provided, a new connection is created.
        """
        self.batch_size = batch_size
        self.batch_replay_delay = batch_replay_delay
        self.status = DBManagerStatus.undefined
        self.tmp_dir = tempfile.mkdtemp()
        if db_file is None:
            db_file = os.path.join(self.tmp_dir, DEFAULT_DB_FILE)

        self.db_file = db_file
        # open DB connection if appropriate
        if conn is None:
            conn = sqlite3.connect(self.db_file, check_same_thread=False)
        self.conn = conn

        if sync_lock is None:
            self.__lock__ = threading.RLock()
        else:
            self.__lock__ = sync_lock

        # Non-reentrant mutex that serializes concurrent replay_failed_messages
        # calls. Producers (register_message / update_message) use self.__lock__,
        # so they are never affected by this mutex.
        self._replay_mutex = threading.Lock()

        self._create_db_schema()

    def _create_db_schema(self) -> None:
        try:
            with self.__lock__:
                with self.conn:
                    self.conn.execute(
                        """CREATE TABLE IF NOT EXISTS messages
                                            (message_id INTEGER NOT NULL PRIMARY KEY,
                                            status INTEGER NOT NULL,
                                            message_type TEXT NOT NULL,
                                            message_json TEXT NOT NULL)"""
                    )
                    self.status = DBManagerStatus.initialized
        except Exception as ex:
            msg = f"Database schema creation failed, reason: {ex}"
            self._mark_as_db_failed(msg)
            LOGGER.warning(msg, exc_info=True)

    def close(self) -> None:
        """
        Closes the current manager, releasing all associated resources.

        This method ensures the proper cleanup of resources such as database connections
        and temporary directories. It performs the operations within a thread-safe context
        by acquiring an internal lock. If the manager is already closed, the method will
        return immediately without performing any actions.
        """
        if self.closed:
            return

        with self.__lock__:
            self.status = DBManagerStatus.closed

            try:
                LOGGER.debug("Closing messages DB connection")
                self.conn.close()
            except Exception as e:
                LOGGER.debug(
                    "Failed to close messages DB connection: %s", e, exc_info=True
                )

            # delete temporary data
            if self.tmp_dir is not None:
                try:
                    LOGGER.debug("Cleaning temporary data dir: %r", self.tmp_dir)
                    shutil.rmtree(self.tmp_dir)
                except Exception as e:
                    LOGGER.debug(
                        "Failed to clean temporary data dir: %r, reason: %s",
                        self.tmp_dir,
                        e,
                        exc_info=True,
                    )

    def register_message(
        self,
        message: messages.BaseMessage,
        status: MessageStatus = MessageStatus.registered,
    ) -> None:
        """
        Registers a message into the database if the system is initialized and not closed.

        This method processes the given message, converts it into a JSON format, and
        inserts it into the database with the provided or default status. If the system
        is either not initialized or already closed, the message registration will be
        ignored.

        Args:
            message: The message object to be registered. This
                object must have attributes such as `message_id` and `message_type`.
            status: The status of the message to be registered.
                Defaults to `MessageStatus.registered`.
        """
        if not self.initialized:
            LOGGER.debug("Not initialized - register message ignored")
            return

        with self.__lock__:
            if self.closed:
                LOGGER.warning("Already closed - register message ignored")
                return

            message_json = _preprocess_registered_message(message)
            # insert into DB
            values = (
                message.message_id,
                status,
                message.message_type,
                message_json,
            )
            try:
                with self.conn:
                    self.conn.execute(
                        "INSERT INTO messages(message_id, status, message_type, message_json) VALUES (?,?,?,?)"
                        " ON CONFLICT(message_id) DO UPDATE SET status = excluded.status, message_type = excluded.message_type, message_json = excluded.message_json",
                        values,
                    )
            except Exception as ex:
                self._mark_as_db_failed(
                    f"register_message: failed to insert message into DB, reason: {ex}"
                )
                raise

    def register_messages(
        self,
        messages_batch: List[messages.BaseMessage],
        status: MessageStatus = MessageStatus.registered,
    ) -> None:
        """
        Registers a batch of messages in the database. If the system is not initialized or
        already closed, the operation is ignored.

        Args:
            messages_batch: A list of message objects to be
                registered in the database.
            status: The status to be associated with the registered messages.
                Defaults to MessageStatus.registered.
        """
        if not self.initialized:
            LOGGER.debug("Not initialized - register messages list ignored")
            return

        with self.__lock__:
            if self.closed:
                LOGGER.warning("Already closed - register messages list ignored")
                return

            values = []
            for message in messages_batch:
                message_json = _preprocess_registered_message(message)
                values.append(
                    (
                        message.message_id,
                        status,
                        message.message_type,
                        message_json,
                    )
                )

            try:
                with self.conn:
                    self.conn.executemany(
                        "INSERT INTO messages(message_id, status, message_type, message_json) VALUES (?,?,?,?)"
                        " ON CONFLICT(message_id) DO UPDATE SET status = excluded.status, message_type = excluded.message_type, message_json = excluded.message_json",
                        values,
                    )
            except Exception as ex:
                self._mark_as_db_failed(
                    f"register_messages: failed to insert messages into DB, reason: {ex}"
                )
                raise

    def update_messages_batch(
        self, message_ids: List[int], status: MessageStatus
    ) -> None:
        """
        Updates the status of a batch of messages in the messages database table or deletes
        the messages if their status is 'delivered'. The function ensures thread safety
        and performs operations only if the instance is initialized and not closed.

        Args:
            message_ids: A list of message IDs to be updated or deleted.
            status: The new status to be assigned to the messages or to
                determine if messages should be deleted when the status is 'delivered'.

        """
        if not self.initialized:
            LOGGER.debug(
                "Not initialized - messages batch update ignored, size: %d, status: %r",
                len(message_ids),
                status,
            )
            return

        with self.__lock__:
            if self.closed:
                LOGGER.warning(
                    "Already closed - messages batch update ignored, size: %d, status: %r",
                    len(message_ids),
                    status,
                )
                return

            try:
                with self.conn:
                    if status == MessageStatus.delivered:
                        db_ids = [(message_id,) for message_id in message_ids]
                        c = self.conn.executemany(
                            "DELETE FROM messages WHERE message_id = ?",
                            db_ids,
                        )
                        LOGGER.debug(
                            "Deleted %d DB message records for %d delivered messages",
                            c.rowcount,
                            len(message_ids),
                        )
                    else:
                        db_status_ids = [
                            (status, message_id) for message_id in message_ids
                        ]
                        c = self.conn.executemany(
                            "UPDATE messages SET status = ? WHERE message_id = ?",
                            db_status_ids,
                        )
                        LOGGER.debug(
                            "Updated %d DB message records for %d messages to have status: %r",
                            c.rowcount,
                            len(message_ids),
                            status,
                        )
            except Exception as ex:
                self._mark_as_db_failed(
                    f"update_messages_batch: failed to update messages batch in the DB, reason: {ex}"
                )
                raise

    def update_message(self, message_id: int, status: MessageStatus) -> None:
        """
        Updates the status of a message in the database or removes it if delivered.

        This method is responsible for updating the status of a message in the database
        or removing delivered messages along with their associated data. The operation
        is ignored if the instance is not initialized or already closed.

        Args:
            message_id: The unique identifier of the message to be updated.
            status: The new status to set for the message. If the status
                is `MessageStatus.delivered`, the message will be removed.
        """
        if not self.initialized:
            LOGGER.debug(
                "Not initialized - message update ignored, id: %d, status: %r",
                message_id,
                status,
            )
            return

        with self.__lock__:
            if self.closed:
                LOGGER.warning(
                    "Already closed - message update ignored, id: %d, status: %r",
                    message_id,
                    status,
                )
                return
            try:
                with self.conn:
                    if status == MessageStatus.delivered:
                        # remove delivered messages
                        self.conn.execute(
                            "DELETE FROM messages WHERE message_id = ?", (message_id,)
                        )
                    else:
                        self.conn.execute(
                            "UPDATE messages SET status = ? WHERE message_id = ?",
                            (status, message_id),
                        )
            except Exception as ex:
                self._mark_as_db_failed(
                    f"update_message: failed to update message in the DB, reason: {ex}"
                )
                raise

    def replay_failed_messages(self, replay_callback: ReplayCallback) -> int:
        """
        Replays previously failed messages by fetching them in batches, marking them as
        "in progress," and invoking the provided callback for processing.

        This method processes messages marked as failed in the database by updating
        their status and invoking the replay callback. It supports batch processing to
        limit memory usage. If the system is not initialized or already closed, the
        replay process is skipped. Errors encountered during the replay or while
        updating the database are logged, and the replay process halts.

        The lock is held only for the minimal critical sections (closed check and the
        DB status update that marks each batch as in-progress). It is released before
        calling the replay callback and before sleeping between batches so that
        producers (register_message / update_message) are never blocked during replay.

        Concurrent calls are serialized by a non-blocking mutex: if a replay is
        already in progress, the second caller returns 0 immediately, preventing
        duplicate re-fetching of the same failed messages.

        Args:
            replay_callback: A callback function that processes the
                replayed messages.

        Returns:
            int: The total count of successfully replayed messages.
        """
        if self.closed:
            LOGGER.debug("DBManager already closed")
            return 0

        if not self.initialized:
            LOGGER.debug("Not initialized - messages replay ignored")
            return 0

        # Prevent concurrent replay calls from re-fetching the same failed
        # messages.  The SELECT in fetch_failed_messages_batched runs outside
        # self.__lock__, so without this guard two simultaneous callers could
        # both see the same rows before either has had a chance to mark them
        # as in-progress.
        if not self._replay_mutex.acquire(blocking=False):
            LOGGER.debug("Replay already in progress - skipping concurrent call")
            return 0

        try:
            return self._replay_failed_messages_impl(replay_callback)
        finally:
            self._replay_mutex.release()

    def _replay_failed_messages_impl(self, replay_callback: ReplayCallback) -> int:
        # The initial closed check — short critical section, no loop held inside.
        with self.__lock__:
            if self.closed:
                LOGGER.warning("Already closed - messages replay ignored")
                return 0

        total_replayed = 0
        for db_messages in self.fetch_failed_messages_batched(self.batch_size):
            params = [
                (int(MessageStatus.registered), message.id) for message in db_messages
            ]
            # Critical section: check closed and mark the batch as "in progress"
            # so the same messages are not re-fetched on a concurrent replay call.
            with self.__lock__:
                if self.closed:
                    LOGGER.debug(
                        "Closed during replay - stopping after %d replayed",
                        total_replayed,
                    )
                    break
                try:
                    with self.conn:
                        c = self.conn.executemany(
                            "UPDATE messages SET status = ? WHERE message_id = ?",
                            params,
                        )
                        LOGGER.debug(
                            "Updated %d DB message records for %d failed messages",
                            c.rowcount,
                            len(db_messages),
                        )
                except Exception as ex:
                    self._mark_as_db_failed(
                        f"replay_failed_messages: failed to update messages in the DB, reason: {ex}"
                    )
                    raise

            # Lock released — replay callback and inter-batch sleep run without
            # blocking producers (register_message / update_message).
            replayed = self._replay_messages(
                db_messages=db_messages, replay_callback=replay_callback
            )
            total_replayed += replayed

            # Sleep without the lock so producers are not blocked during the pause.
            # Applied only after a full batch to avoid unnecessary delays on the
            # last (partial) batch.
            if replayed >= self.batch_size:
                time.sleep(self.batch_replay_delay)

        return total_replayed

    def _replay_messages(
        self, db_messages: List[DBMessage], replay_callback: ReplayCallback
    ) -> int:
        LOGGER.debug("Replaying %d failed messages to streamer", len(db_messages))
        replayed = 0
        for message in db_messages:
            try:
                base_message = db_message_to_message(message)
                replay_callback(base_message)
                replayed += 1
            except Exception as e:
                LOGGER.error(
                    "Failed to replay message with id=%r, type=%r, status=%r, reason: %s",
                    message.id,
                    message.type,
                    message.status,
                    e,
                    exc_info=True,
                )
                # mark the message as failed in the DB
                self.update_message(message.id, MessageStatus.failed)

        return replayed

    def get_message(self, message_id: int) -> Optional[messages.BaseMessage]:
        """
        Fetches a message by its unique identifier.

        This method retrieves a message from the database using the provided message ID. If a corresponding
        database message is found, it converts the database message to a message object and returns it.
        If no message is found for the given ID, the method returns None.

        Args:
            message_id: The unique identifier of the message to retrieve.

        Returns:
            The converted message object if found, otherwise None.
        """
        db_message = self.get_db_message(message_id)
        if db_message is not None:
            return db_message_to_message(db_message)
        else:
            return None

    def get_db_message(self, message_id: int) -> Optional[DBMessage]:
        """
        Retrieves a database message based on the provided message ID.

        This method queries the database for a message record that corresponds to the
        given message ID. If a matching message record is found, it constructs and
        returns a `DBMessage` object containing the message's details. If no record is
        found, it returns `None`.

        Args:
            message_id: The ID of the message to retrieve.

        Returns:
            A `DBMessage` object containing the details of the
            retrieved message if found, otherwise `None`.
        """
        with self.conn:
            c = self.conn.execute(
                "SELECT message_id, message_type, message_json, status FROM messages WHERE message_id = ?",
                (message_id,),
            )
            row = c.fetchone()
            if row is not None:
                return DBMessage(
                    id=row[0], type=row[1], json=row[2], status=MessageStatus(row[3])
                )
            else:
                return None

    @property
    def closed(self) -> bool:
        """
        Determines if the manager is currently closed.

        This property evaluates whether the manager's `status` is equivalent to
        `DBManagerStatus.closed`. If so, it indicates that the manager is not
        active or operational.

        Returns:
            bool: True if the manager's status is `DBManagerStatus.closed`, False otherwise.
        """
        return self.status == DBManagerStatus.closed

    @property
    def initialized(self) -> bool:
        """
        Indicates whether the manager's status is currently set to 'initialized'.

        This property provides a lookup to determine if the manager is currently in the
        initialized state, based on its status attribute.

        Returns:
            bool: True if the manager's status is 'initialized', False otherwise.
        """
        return self.status == DBManagerStatus.initialized

    @property
    def failed(self) -> bool:
        """
        Checks if the current status indicates a failure.

        The `failed` property evaluates whether the `status` attribute of the
        manager is equal to `ManagerStatus.error`. This serves as an indicator
        of whether an error state has been encountered.

        Returns:
            bool: True if the status is `ManagerStatus.error`, otherwise False.
        """
        return self.status == DBManagerStatus.error

    def fetch_failed_messages_batched(
        self, batch_size: int
    ) -> Iterator[List[DBMessage]]:
        """Fetch failed messages from DB in bounded batches to avoid OOM.

        Uses cursor-based pagination with message_id for efficient iteration.
        Yields batches of DBMessage objects until no more failed messages remain.
        """
        last_seen_id = -1
        while True:
            batch: List[DBMessage] = []
            with self.__lock__:
                try:
                    rows = self.conn.execute(
                        "SELECT message_id, message_type, message_json FROM messages "
                        "WHERE status = ? AND message_id > ? ORDER BY message_id LIMIT ?",
                        (MessageStatus.failed, last_seen_id, batch_size),
                    )
                except Exception as ex:
                    self._mark_as_db_failed(
                        f"fetch_failed_messages_batched: failed to fetch failed messages from the DB, reason: {ex}"
                    )
                    raise

                for row in rows:
                    batch.append(
                        DBMessage(
                            id=row[0],
                            type=row[1],
                            json=row[2],
                            status=MessageStatus.failed,
                        )
                    )

            if not batch:
                break

            last_seen_id = batch[-1].id
            yield batch

    def failed_messages_count(self) -> int:
        """Returns the number of failed messages in the DB or -1 if not initialized, closed, or failed."""
        if not self.initialized:
            LOGGER.debug("Not initialized - failed messages count ignored")
            return -1

        with self.__lock__:
            if self.closed:
                LOGGER.warning("Already closed - failed messages count ignored")
                return -1

            try:
                return self.conn.execute(
                    "SELECT COUNT(*) FROM messages WHERE status = ?",
                    (MessageStatus.failed,),
                ).fetchone()[0]
            except Exception as ex:
                msg = f"failed_messages_count: failed to get failed messages count, reason: {ex}"
                LOGGER.error(msg)
                self._mark_as_db_failed(msg)
                return -1

    def _mark_as_db_failed(self, message: str) -> None:
        self.status = DBManagerStatus.error
        LOGGER.error(
            "Due to an internal error, some network resiliency features were disabled "
            "which could lead to data loss. Contact us at support@comet.com. Error details: %r",
            message,
        )


SUPPORTED_MESSAGE_TYPES = {
    messages.CreateTraceMessage.message_type: messages.CreateTraceMessage,
    messages.UpdateTraceMessage.message_type: messages.UpdateTraceMessage,
    messages.CreateSpanMessage.message_type: messages.CreateSpanMessage,
    messages.UpdateSpanMessage.message_type: messages.UpdateSpanMessage,
    messages.AddTraceFeedbackScoresBatchMessage.message_type: messages.AddTraceFeedbackScoresBatchMessage,
    messages.AddSpanFeedbackScoresBatchMessage.message_type: messages.AddSpanFeedbackScoresBatchMessage,
    messages.AddThreadsFeedbackScoresBatchMessage.message_type: messages.AddThreadsFeedbackScoresBatchMessage,
    messages.CreateSpansBatchMessage.message_type: messages.CreateSpansBatchMessage,
    messages.CreateTraceBatchMessage.message_type: messages.CreateTraceBatchMessage,
    messages.GuardrailBatchMessage.message_type: messages.GuardrailBatchMessage,
    messages.CreateExperimentItemsBatchMessage.message_type: messages.CreateExperimentItemsBatchMessage,
    messages.CreateAttachmentMessage.message_type: messages.CreateAttachmentMessage,
}


def db_message_to_message(db_message: DBMessage) -> messages.BaseMessage:
    """
    Converts a database message object to a corresponding application message object.

    This function maps a database message type to its associated application message
    class using a predefined dictionary of supported message types. If the type is not
    supported, an exception is raised. The message is then deserialized from its JSON
    representation into the appropriate application message class.

    Args:
        db_message: The database message object containing the type and serialized JSON
            data to be deserialized.

    Returns:
        The deserialized application message object corresponding to the provided
        database message type.

    Raises:
        ValueError: If the database message contains an unsupported or unrecognized
        message type.
    """
    message_class = SUPPORTED_MESSAGE_TYPES.get(db_message.type)
    if message_class is None:
        raise ValueError(f"Unsupported message type: {db_message.type}")

    return message_serialization.deserialize_message(
        message_class, json_str=db_message.json
    )
