from datetime import datetime
from typing import Dict, List, Optional

from sqlalchemy import select

from mirix.orm.errors import NoResultFound
from mirix.orm.message import Message as MessageModel
from mirix.schemas.client import Client as PydanticClient
from mirix.schemas.enums import MessageRole
from mirix.schemas.message import Message as PydanticMessage
from mirix.schemas.message import MessageUpdate
from mirix.services.utils import update_timezone
from mirix.utils import enforce_types


class MessageManager:
    """Manager class to handle business logic related to Messages."""

    def __init__(self):
        from mirix.server.server import db_context

        self.session_maker = db_context

    @update_timezone
    @enforce_types
    async def get_message_by_id(
        self, message_id: str, actor: PydanticClient, use_cache: bool = True
    ) -> Optional[PydanticMessage]:
        """Fetch a message by ID (with cache - Redis or IPS Cache)."""
        from mirix.log import get_logger

        logger = get_logger(__name__)
        cache_provider = None
        try:
            from mirix.database.cache_provider import get_cache_provider

            cache_provider = get_cache_provider()

            if use_cache and cache_provider:
                cache_key = f"{cache_provider.MESSAGE_PREFIX}{message_id}"
                cached_data = await cache_provider.get_hash(cache_key)
                if cached_data:
                    return PydanticMessage(**cached_data)
        except Exception as e:
            logger.warning("Cache read failed for message %s: %s", message_id, e)

        # Cache MISS or no cache - fetch from PostgreSQL
        async with self.session_maker() as session:
            try:
                message = await MessageModel.read(db_session=session, identifier=message_id, actor=actor)
                pydantic_message = message.to_pydantic()

                try:
                    if cache_provider:
                        from mirix.settings import settings

                        cache_key = f"{cache_provider.MESSAGE_PREFIX}{message_id}"
                        data = pydantic_message.model_dump(mode="json")
                        await cache_provider.set_hash(cache_key, data, ttl=settings.redis_ttl_messages)
                except Exception as e:
                    logger.warning("Failed to populate cache for message %s: %s", message_id, e)

                return pydantic_message
            except NoResultFound:
                return None

    @update_timezone
    @enforce_types
    async def get_messages_by_ids(
        self, message_ids: List[str], actor: PydanticClient, use_cache: bool = True
    ) -> List[PydanticMessage]:
        """Fetch messages by ID and return them in the requested order."""
        # TODO: Add Redis pipeline support for batch retrieval when use_cache=True
        async with self.session_maker() as session:
            results = await MessageModel.list(
                db_session=session,
                id=message_ids,
                organization_id=actor.organization_id,
                limit=len(message_ids),
            )

            # Return messages in requested order, skipping any missing IDs
            # (messages may be missing due to concurrent summarization/cleanup by other workers)
            result_dict = {msg.id: msg.to_pydantic() for msg in results}
            return [result_dict[msg_id] for msg_id in message_ids if msg_id in result_dict]

    @enforce_types
    async def create_message(
        self,
        pydantic_msg: PydanticMessage,
        actor: PydanticClient,
        use_cache: bool = True,
        client_id: Optional[str] = None,
        user_id: Optional[str] = None,
    ) -> PydanticMessage:
        """Create a new message (with Redis Hash caching).

        Args:
            pydantic_msg: The message data to create
            actor: Client performing the operation (for audit trail)
            use_cache: If True, cache in Redis. If False, skip caching.
            client_id: Optional client_id for data scoping (defaults to actor.id)
            user_id: Optional user_id for data scoping (defaults to None)
        """
        async with self.session_maker() as session:
            # Set the organization id of the Pydantic message
            pydantic_msg.organization_id = actor.organization_id

            # Set client_id: use provided value, or from message, or from actor
            if client_id:
                pydantic_msg.client_id = client_id
            elif not pydantic_msg.client_id:
                pydantic_msg.client_id = actor.id

            # Set user_id: use provided value, or from message, or fallback to default user
            if user_id is not None:
                pydantic_msg.user_id = user_id
            elif not pydantic_msg.user_id:
                # Fallback: use default system user for client-level operations
                from mirix.services.user_manager import UserManager

                pydantic_msg.user_id = UserManager.ADMIN_USER_ID

            msg_data = pydantic_msg.model_dump()
            msg = MessageModel(**msg_data)
            await msg.create_with_redis(session, actor=actor, use_cache=use_cache)
            return msg.to_pydantic()

    @enforce_types
    async def create_many_messages(
        self,
        pydantic_msgs: List[PydanticMessage],
        actor: PydanticClient,
        client_id: Optional[str] = None,
        user_id: Optional[str] = None,
    ) -> List[PydanticMessage]:
        """Create multiple messages."""
        return [await self.create_message(m, actor=actor, client_id=client_id, user_id=user_id) for m in pydantic_msgs]

    @enforce_types
    async def update_message_by_id(
        self, message_id: str, message_update: MessageUpdate, actor: PydanticClient
    ) -> PydanticMessage:
        """
        Updates an existing record in the database with values from the provided record object.
        """
        async with self.session_maker() as session:
            # Fetch existing message from database
            message = await MessageModel.read(
                db_session=session,
                identifier=message_id,
                actor=actor,
            )

            # Some safety checks specific to messages
            if message_update.tool_calls and message.role != MessageRole.assistant:
                raise ValueError(
                    f"Tool calls {message_update.tool_calls} can only be added to assistant messages. Message {message_id} has role {message.role}."
                )
            if message_update.tool_call_id and message.role != MessageRole.tool:
                raise ValueError(
                    f"Tool call IDs {message_update.tool_call_id} can only be added to tool messages. Message {message_id} has role {message.role}."
                )

            # get update dictionary
            update_data = message_update.model_dump(exclude_unset=True, exclude_none=True)
            # Remove redundant update fields
            update_data = {key: value for key, value in update_data.items() if getattr(message, key) != value}

            for key, value in update_data.items():
                setattr(message, key, value)
            await message.update_with_redis(db_session=session, actor=actor)  # Update PostgreSQL + Redis

            return message.to_pydantic()

    @enforce_types
    async def delete_message_by_id(self, message_id: str, actor: PydanticClient) -> bool:
        """Delete a message (removes from cache)."""
        async with self.session_maker() as session:
            try:
                msg = await MessageModel.read(
                    db_session=session,
                    identifier=message_id,
                    actor=actor,
                )
                # Remove from cache before hard delete
                from mirix.database.cache_provider import get_cache_provider

                cache_provider = get_cache_provider()
                if cache_provider:
                    cache_key = f"{cache_provider.MESSAGE_PREFIX}{message_id}"
                    await cache_provider.delete(cache_key)
                await msg.hard_delete(session, actor=actor)
            except NoResultFound:
                raise ValueError(f"Message with id {message_id} not found.")

    @enforce_types
    async def delete_by_client_id(self, actor: PydanticClient) -> int:
        """
        Bulk delete all NON-SYSTEM messages for a client (removes from Redis cache).
        System messages are preserved as they are essential for agent functionality.
        Optimized with single DB query and batch Redis deletion.

        Args:
            actor: Client whose messages to delete (uses actor.id as client_id)

        Returns:
            Number of records deleted
        """
        from mirix.database.redis_client import get_redis_client
        from mirix.schemas.message import MessageRole

        async with self.session_maker() as session:
            # Get IDs for non-system messages only (preserve system messages)
            stmt = select(MessageModel.id).where(
                MessageModel.client_id == actor.id,
                MessageModel.role != MessageRole.system,
            )
            result = await session.execute(stmt)
            message_ids = [row[0] for row in result.all()]

            count = len(message_ids)
            if count == 0:
                return 0

            # Bulk delete in single query (exclude system messages)
            from sqlalchemy import delete

            await session.execute(
                delete(MessageModel).where(
                    MessageModel.client_id == actor.id,
                    MessageModel.role != MessageRole.system,
                )
            )
            await session.commit()

        # Batch delete from Redis cache (outside of session context)
        redis_client = get_redis_client()
        if redis_client and message_ids:
            redis_keys = [f"{redis_client.MESSAGE_PREFIX}{msg_id}" for msg_id in message_ids]

            # Delete in batches to avoid command size limits
            BATCH_SIZE = 1000
            for i in range(0, len(redis_keys), BATCH_SIZE):
                batch = redis_keys[i : i + BATCH_SIZE]
                await redis_client.client.delete(*batch)

        return count

    async def soft_delete_by_client_id(self, actor: PydanticClient) -> int:
        """
        Bulk soft delete all messages for a client (updates Redis cache).

        Args:
            actor: Client whose messages to soft delete (uses actor.id as client_id)

        Returns:
            Number of records soft deleted
        """
        from mirix.database.redis_client import get_redis_client

        async with self.session_maker() as session:
            # Query all non-deleted records for this client (use actor.id)
            stmt = select(MessageModel).where(
                MessageModel.client_id == actor.id,
                MessageModel.is_deleted == False,
            )
            result = await session.execute(stmt)
            messages = result.scalars().all()

            count = len(messages)
            if count == 0:
                return 0

            # Extract IDs BEFORE committing (to avoid detached instance errors)
            message_ids = [msg.id for msg in messages]

            # Soft delete from database (set is_deleted = True directly, don't call msg.delete())
            for msg in messages:
                msg.is_deleted = True
                msg.set_updated_at()

            await session.commit()

        # Update Redis cache with is_deleted=true (outside session)
        redis_client = get_redis_client()
        if redis_client:
            for msg_id in message_ids:
                redis_key = f"{redis_client.MESSAGE_PREFIX}{msg_id}"
                try:
                    await redis_client.client.hset(redis_key, "is_deleted", "true")
                except Exception:
                    # If update fails, remove from cache
                    await redis_client.delete(redis_key)

        return count

    async def soft_delete_by_user_id(self, user_id: str) -> int:
        """
        Bulk soft delete all messages for a user (updates Redis cache).

        Args:
            user_id: ID of the user whose messages to soft delete

        Returns:
            Number of records soft deleted
        """
        from mirix.database.redis_client import get_redis_client

        async with self.session_maker() as session:
            # Query all non-deleted records for this user
            stmt = select(MessageModel).where(
                MessageModel.user_id == user_id,
                MessageModel.is_deleted == False,
            )
            result = await session.execute(stmt)
            messages = result.scalars().all()

            count = len(messages)
            if count == 0:
                return 0

            # Extract IDs BEFORE committing (to avoid detached instance errors)
            message_ids = [msg.id for msg in messages]

            # Soft delete from database (set is_deleted = True directly, don't call msg.delete())
            for msg in messages:
                msg.is_deleted = True
                msg.set_updated_at()

            await session.commit()

        # Update Redis cache with is_deleted=true (outside session)
        redis_client = get_redis_client()
        if redis_client:
            for msg_id in message_ids:
                redis_key = f"{redis_client.MESSAGE_PREFIX}{msg_id}"
                try:
                    await redis_client.client.hset(redis_key, "is_deleted", "true")
                except Exception:
                    # If update fails, remove from cache
                    await redis_client.delete(redis_key)

        return count

    async def delete_by_user_id(self, user_id: str) -> int:
        """
        Bulk hard delete all NON-SYSTEM messages for a user (removes from Redis cache).
        System messages are preserved as they are essential for agent functionality.
        Optimized with single DB query and batch Redis deletion.

        Args:
            user_id: ID of the user whose messages to delete

        Returns:
            Number of records deleted
        """
        from sqlalchemy import delete

        from mirix.database.redis_client import get_redis_client
        from mirix.schemas.message import MessageRole

        async with self.session_maker() as session:
            # Get IDs for non-system messages only (preserve system messages)
            stmt = select(MessageModel.id).where(
                MessageModel.user_id == user_id,
                MessageModel.role != MessageRole.system,
            )
            result = await session.execute(stmt)
            message_ids = [row[0] for row in result.all()]

            count = len(message_ids)
            if count == 0:
                return 0

            # Bulk delete in single query (exclude system messages)
            await session.execute(
                delete(MessageModel).where(
                    MessageModel.user_id == user_id,
                    MessageModel.role != MessageRole.system,
                )
            )
            await session.commit()

        # Batch delete from Redis cache (outside of session context)
        redis_client = get_redis_client()
        if redis_client and message_ids:
            redis_keys = [f"{redis_client.MESSAGE_PREFIX}{msg_id}" for msg_id in message_ids]

            # Delete in batches to avoid command size limits
            BATCH_SIZE = 1000
            for i in range(0, len(redis_keys), BATCH_SIZE):
                batch = redis_keys[i : i + BATCH_SIZE]
                await redis_client.client.delete(*batch)

        return count

    @enforce_types
    async def size(
        self,
        actor: PydanticClient,
        role: Optional[MessageRole] = None,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> int:
        """Get the total count of messages with optional filters.

        Args:
            actor: The client requesting the count
            role: The role of the message
        """
        async with self.session_maker() as session:
            return await MessageModel.size(
                db_session=session,
                actor=actor,
                role=role,
                user_id=user_id,
                agent_id=agent_id,
            )

    @update_timezone
    @enforce_types
    async def list_user_messages_for_agent(
        self,
        agent_id: str,
        actor: Optional[PydanticClient] = None,
        cursor: Optional[str] = None,
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None,
        limit: Optional[int] = 50,
        filters: Optional[Dict] = None,
        query_text: Optional[str] = None,
        ascending: bool = True,
    ) -> List[PydanticMessage]:
        """List user messages with flexible filtering and pagination options.

        Args:
            cursor: Cursor-based pagination - return records after this ID (exclusive)
            start_date: Filter records created after this date
            end_date: Filter records created before this date
            limit: Maximum number of records to return
            filters: Additional filters to apply
            query_text: Optional text to search for in message content

        Returns:
            List[PydanticMessage] - List of messages matching the criteria
        """
        message_filters = {"role": "user"}
        if filters:
            message_filters.update(filters)

        return await self.list_messages_for_agent(
            agent_id=agent_id,
            actor=actor,
            cursor=cursor,
            start_date=start_date,
            end_date=end_date,
            limit=limit,
            filters=message_filters,
            query_text=query_text,
            ascending=ascending,
        )

    @update_timezone
    @enforce_types
    async def list_messages_for_agent(
        self,
        agent_id: str,
        actor: Optional[PydanticClient] = None,
        cursor: Optional[str] = None,
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None,
        limit: Optional[int] = 50,
        filters: Optional[Dict] = None,
        query_text: Optional[str] = None,
        ascending: bool = True,
        use_cache: bool = True,
    ) -> List[PydanticMessage]:
        """List messages with flexible filtering and pagination options.

        Args:
            cursor: Cursor-based pagination - return records after this ID (exclusive)
            start_date: Filter records created after this date
            end_date: Filter records created before this date
            limit: Maximum number of records to return
            filters: Additional filters to apply
            query_text: Optional text to search for in message content
            use_cache: Control Redis cache behavior (default: True)
                      Note: Currently this method only queries PostgreSQL.
                      Redis list retrieval for messages is not yet implemented.

        Returns:
            List[PydanticMessage] - List of messages matching the criteria
        """
        async with self.session_maker() as session:
            # Start with base filters
            message_filters = {"agent_id": agent_id}
            if actor:
                message_filters.update({"organization_id": actor.organization_id})
            if filters:
                message_filters.update(filters)

            results = await MessageModel.list(
                db_session=session,
                cursor=cursor,
                start_date=start_date,
                end_date=end_date,
                limit=limit,
                query_text=query_text,
                ascending=ascending,
                **message_filters,
            )

            return [msg.to_pydantic() for msg in results]

    async def get_messages_for_agent_user(
        self,
        agent_id: str,
        user_id: str,
        actor: PydanticClient,
        limit: int = 10,
    ) -> List[PydanticMessage]:
        """
        Fetch the most recent N messages for a given (agent, user) pair, returned in
        chronological order (oldest first).

        Uses the composite index ix_messages_agent_user_created_at for efficient retrieval.

        Args:
            agent_id: The agent whose messages to retrieve
            user_id: The user whose messages to retrieve
            actor: Client performing the operation (for org scoping)
            limit: Maximum number of messages to return (newest N, then reversed)

        Returns:
            List of messages in chronological order
        """
        from sqlalchemy import desc

        async with self.session_maker() as session:
            stmt = (
                select(MessageModel)
                .where(
                    MessageModel.agent_id == agent_id,
                    MessageModel.user_id == user_id,
                    MessageModel.organization_id == actor.organization_id,
                    MessageModel.is_deleted == False,
                )
                .order_by(desc(MessageModel.created_at), desc(MessageModel.id))
                .limit(limit)
            )
            result = await session.execute(stmt)
            messages = result.scalars().all()

        # Reverse to chronological order
        return [msg.to_pydantic() for msg in reversed(messages)]

    async def hard_delete_user_messages_for_agent(
        self,
        agent_id: str,
        user_id: str,
        actor: PydanticClient,
        keep_newest_n: int = 0,
    ) -> int:
        """
        Hard-delete messages for a (agent, user) pair, optionally keeping the newest N.

        Deletes from both the database and Redis cache.

        Args:
            agent_id: The agent whose messages to prune
            user_id: The user whose messages to prune
            actor: Client performing the operation (for org scoping)
            keep_newest_n: How many of the most-recent messages to retain. 0 = delete all.

        Returns:
            Number of messages deleted
        """
        from sqlalchemy import delete, desc

        from mirix.database.redis_client import get_redis_client

        async with self.session_maker() as session:
            # Identify IDs to keep (the newest N)
            keep_ids: set = set()
            if keep_newest_n > 0:
                keep_stmt = (
                    select(MessageModel.id)
                    .where(
                        MessageModel.agent_id == agent_id,
                        MessageModel.user_id == user_id,
                        MessageModel.organization_id == actor.organization_id,
                        MessageModel.is_deleted == False,
                    )
                    .order_by(desc(MessageModel.created_at), desc(MessageModel.id))
                    .limit(keep_newest_n)
                )
                keep_result = await session.execute(keep_stmt)
                keep_ids = {row[0] for row in keep_result.all()}

            # Collect IDs that will be deleted (for cache invalidation)
            select_stmt = select(MessageModel.id).where(
                MessageModel.agent_id == agent_id,
                MessageModel.user_id == user_id,
                MessageModel.organization_id == actor.organization_id,
                MessageModel.is_deleted == False,
            )
            if keep_ids:
                select_stmt = select_stmt.where(MessageModel.id.not_in(keep_ids))

            id_result = await session.execute(select_stmt)
            delete_ids = [row[0] for row in id_result.all()]

            count = len(delete_ids)
            if count == 0:
                return 0

            # Bulk delete
            del_stmt = delete(MessageModel).where(
                MessageModel.agent_id == agent_id,
                MessageModel.user_id == user_id,
                MessageModel.organization_id == actor.organization_id,
                MessageModel.is_deleted == False,
            )
            if keep_ids:
                del_stmt = del_stmt.where(MessageModel.id.not_in(keep_ids))

            await session.execute(del_stmt)
            await session.commit()

        # Evict from Redis cache
        redis_client = get_redis_client()
        if redis_client and delete_ids:
            BATCH_SIZE = 1000
            for i in range(0, len(delete_ids), BATCH_SIZE):
                batch = [f"{redis_client.MESSAGE_PREFIX}{mid}" for mid in delete_ids[i : i + BATCH_SIZE]]
                await redis_client.client.delete(*batch)

        return count
