"""
Manages conversation history, token counting, and SQLite-backed persistence for any LLM provider.

The Conversation class is keyed by chat_id (Telegram chat ID) and maintains an in-memory message
history, delegating persistence to the database module. Supports async token counting via provider-
specific APIs and automatic pruning when token limits are approached.
"""
from .database import (
    insert_message,
    upsert_user,
    upsert_chat,
    load_full_user_context,
    get_max_cross_chat_message_id,
    load_new_cross_chat_context,
    get_shared_group_chat_ids,
    load_shared_group_context,
    delete_messages_for_user,
    delete_messages_for_chat,
)
from .models import TokenLimits
from .utils import generate_error_path, log_error

# Minimum remaining token headroom required before loading shared group context
# in a private chat. Prevents loading group history when the private context
# already fills nearly the entire budget.
MIN_GROUP_CONTEXT_TOKENS = 500

class Conversation:
    """
    Manages a chat conversation including message history, token counting, and SQLite persistence.

    Keyed by chat_id (Telegram chat ID): negative for groups, positive (== user_id) for private chats.
    Works with any LLM provider (OpenAI, Anthropic). When the token limit is approached, oldest
    messages are pruned from memory; the database retains the full history.
    """

    def __init__(self, chat_id: int, chat_type: str, system_content: str, system_model: str = "gpt-4o-mini", chat_title: str | None = None):
        """
        Initialize a conversation keyed by chat_id.

        Args:
            chat_id: Telegram chat ID (negative for groups, positive/==user_id for private chats).
            chat_type: Chat type ('private', 'group', 'supergroup', 'channel').
            system_content: System prompt (bot persona).
            system_model: LLM model to use for token counting (default: "gpt-4o-mini").
            chat_title: Display name of the group/channel, or None for private chats.
        """
        self.error_log = generate_error_path()

        self.chat_id = chat_id
        self.chat_type = chat_type
        self.chat_title = chat_title
        self.chat_print = f"{self.chat_type} Chat {self.chat_id}"

        self.system_content = system_content
        self.system_model = TokenLimits(system_model)
        self.messages = [{"role": "system", "content": system_content}]

        # Maps user_id -> max message id at the time their cross-chat context was loaded.
        # Used for both per-user idempotency (is the user loaded?) and mid-session refresh
        # detection (have new cross-chat messages appeared since the last load?).
        self._context_cursor: dict[int, int] = {}

        # Maps user_id -> len(self.messages) after their history was loaded.
        # Marks the boundary between historical context and live session messages so that
        # refresh_user_context() inserts delta rows at the right position (after history,
        # before live messages) rather than at index 1 (before all existing history).
        self._history_end: dict[int, int] = {}

        # Set of id()s for message dicts added while private mode was ON.
        # Used by purge_private_messages() to remove them from self.messages without
        # touching the database. Cleared on purge or clear_interaction().
        self._private_message_ids: set[int] = set()

    def set_system_content(self, new_content: str):
        """
        Replace the system content (bot's personality prompt) in the conversation.

        Updates both self.system_content and the first message in self.messages.

        Args:
            new_content: The new system prompt text.
        """
        self.system_content = new_content
        self.messages[0] = {"role": "system", "content": new_content}

    async def add_user_message(
        self,
        content: str,
        user_id: int,
        username: str | None,
        first_name: str | None = None,
        last_name: str | None = None,
        is_private: bool = False,
    ) -> int:
        """
        Add a user message to conversation memory and persist it to the database.

        Also upserts the sender into the users table and this chat into the chats table
        so that search_messages can resolve names and chat titles.

        Args:
            content: The message text.
            user_id: Telegram user ID of the sender.
            username: Telegram @handle (without @), may be None.
            first_name: Telegram first name, may be None.
            last_name: Telegram last name, may be None.
            is_private: If True, flags this message as private-mode excluded from
                        all context loads (user has /private ON in a private chat).

        Returns:
            The database id of the newly inserted row, for use as reply_to_id.
        """
        await upsert_user(user_id, username, first_name, last_name)
        await upsert_chat(self.chat_id, self.chat_type, self.chat_title)
        msg_dict = {"role": "user", "content": content}
        self.messages.append(msg_dict)
        if is_private:
            self._private_message_ids.add(id(msg_dict))
        return await insert_message(self.chat_id, user_id, "user", content, is_private)

    async def add_assistant_message(
        self,
        content: str,
        bot_user_id: int,
        bot_username: str | None,
        bot_first_name: str | None = None,
        bot_last_name: str | None = None,
        is_private: bool = False,
        reply_to_id: int | None = None,
    ):
        """
        Add an assistant message to conversation memory and persist it to the database.

        Also upserts the bot into the users table and this chat into the chats table.

        Args:
            content: The message text (LLM response).
            bot_user_id: Telegram user ID of the bot account.
            bot_username: Telegram username of the bot (may be None).
            bot_first_name: Bot's first name from Telegram (may be None).
            bot_last_name: Bot's last name from Telegram (may be None).
            is_private: If True, excludes this reply from group context loading. Should match
                        the is_private flag of the user message it responds to, so that both
                        sides of a private-mode exchange are excluded from shared contexts.
            reply_to_id: Database id of the user message this reply responds to.
        """
        await upsert_user(bot_user_id, bot_username, bot_first_name, bot_last_name)
        await upsert_chat(self.chat_id, self.chat_type, self.chat_title)
        msg_dict = {"role": "assistant", "content": content}
        self.messages.append(msg_dict)
        if is_private:
            self._private_message_ids.add(id(msg_dict))
        await insert_message(self.chat_id, bot_user_id, "assistant", content, is_private, reply_to_id)

    async def get_message_token_count(self) -> int:
        """
        Get the current token count of all messages in this conversation.

        Delegates to the configured model's provider to count tokens accurately.
        This is an async operation because it may contact the provider API (e.g., for Anthropic).

        Returns:
            The total token count for all messages (system, user, assistant).

        Raises:
            ProviderAuthError: If API credentials are missing or invalid.
            ProviderConnectionError: If the provider API cannot be reached.
        """
        return await self.system_model.num_tokens_from_messages(self.messages)

    async def prune_conversation(self, token_limit: int):
        """
        Remove oldest user-assistant messages until token count is below threshold.

        Recursively removes messages starting from index 1 (preserving the system message at index 0)
        until the conversation fits within the token limit. The database is not modified - pruning
        only affects the in-memory context window for the current session.

        Args:
            token_limit: The maximum token count threshold. Messages are removed until
                         the conversation is strictly below this limit.
        """
        try:
            self.messages.pop(1)
            if await self.get_message_token_count() > token_limit:
                await self.prune_conversation(token_limit)
        except Exception as e:
            log_error(e, f"{type(e).__name__} pruning under {token_limit} tokens for {self.chat_print}", self.error_log)

    def purge_private_messages(self):
        """
        Remove all in-session private-mode messages from the in-memory conversation.

        Strips any messages added while private mode was ON from self.messages without
        modifying the database. Called when a user disables private mode (/private off)
        so the bot can no longer reference messages sent during the private session.

        Side Effects:
            - Modifies self.messages in-place, removing tracked private message dicts.
            - Clears self._private_message_ids.
        """
        if not self._private_message_ids:
            return
        self.messages = [m for m in self.messages if id(m) not in self._private_message_ids]
        self._private_message_ids.clear()

    async def get_past_interaction(self, token_limit: int, user_id: int, bot_id: int | None = None) -> bool:
        """
        Load past conversation messages from the database up to a token limit.

        Implements bidirectional cross-pollination: loads all messages from the current chat
        (all participants) plus all messages from the user's private chat (both user and
        assistant turns), merged chronologically. This eliminates amnesia when a user switches
        between private and group contexts. Messages marked as private (is_private=1) are
        always excluded, regardless of chat type, ensuring private-mode exchanges never leak
        into any session context.

        For private chats, if bot_id is provided and remaining token budget allows, also loads
        full message history from any groups shared between the user and the bot. This lets the
        bot answer questions about group conversations from a private chat context.

        Per-user idempotent: once a user's cursor is set, subsequent calls return True immediately.
        This guard handles concurrent message arrival: if two messages from the same user arrive
        before the first load completes and both check the cursor, the second call hits the early
        return and skips a duplicate load, ensuring single-pass context loading even under race
        conditions. Calling for a new user_id merges that user's cross-chat context into the
        existing conversation. Context is considered "found" only if messages were actually
        inserted (rows_inserted > 0 or group_rows_inserted > 0).

        Args:
            token_limit: The maximum token count to load into the in-memory context.
            user_id: Telegram user ID of the requesting user; used to fetch bidirectional context.
            bot_id: Telegram user ID of the bot account. When provided in a private chat,
                    enables shared group context loading into any remaining token budget.

        Returns:
            True if past interactions were found (either on this call or previously),
            False if no past interactions exist for this user.

        Side Effects:
            - Modifies self.messages in-place by inserting historical messages after the system message
              (attempting to fit messages starting from most recent, stopping when limit is reached).
            - Sets self._context_cursor[user_id] to the max message id loaded once context is found.
            - Sets self._history_end[user_id] to len(self.messages) after loading, marking the
              boundary between historical context and live session messages for future refreshes.
            - Prints context loading summary (Full/Partial context with token count) if history found,
              or "No prior context" message if no history exists (removed: early-return and refresh
              diagnostics to reduce noisy logging).
        """
        if user_id in self._context_cursor:
            return True

        exclude_private = True
        rows = await load_full_user_context(user_id, self.chat_id, exclude_private=exclude_private)

        token_limit_reached = False
        token_count = 0
        rows_inserted = 0

        for row in reversed(rows):
            self.messages.insert(1, row)
            token_count = await self.get_message_token_count()
            if token_count > token_limit:
                self.messages.pop(1)
                token_count = await self.get_message_token_count()
                token_limit_reached = True
                break
            rows_inserted += 1

        # Load shared group context for private chats when token budget allows.
        # Private context fills the primary budget; group context fills the remainder.
        group_rows_inserted = 0
        group_limit_reached = False
        if self.chat_type == 'private' and bot_id is not None:
            if token_count == 0:
                token_count = await self.get_message_token_count()
            remaining = token_limit - token_count
            if remaining > MIN_GROUP_CONTEXT_TOKENS:
                shared_groups = await get_shared_group_chat_ids(user_id, bot_id)
                if shared_groups:
                    group_rows = await load_shared_group_context(shared_groups, exclude_private=True)
                    insert_at = 1 + rows_inserted
                    for row in reversed(group_rows):
                        self.messages.insert(insert_at, row)
                        token_count = await self.get_message_token_count()
                        if token_count > token_limit:
                            self.messages.pop(insert_at)
                            token_count = await self.get_message_token_count()
                            group_limit_reached = True
                            break
                        group_rows_inserted += 1

        has_context = rows_inserted > 0 or group_rows_inserted > 0

        if has_context:
            self._context_cursor[user_id] = await get_max_cross_chat_message_id(
                user_id, self.chat_id, exclude_private=exclude_private
            )
            self._history_end[user_id] = len(self.messages)
            label = "Full" if not token_limit_reached and not group_limit_reached else "Partial"
            group_note = f", plus {group_rows_inserted} group message(s)" if group_rows_inserted > 0 else ""
            print(
                f"{label} context loaded for User {user_id} in {self.chat_print} - "
                f"storing {token_count} token(s){group_note}"
            )
        else:
            print(f"No prior context found for User {user_id} in {self.chat_print}")

        return has_context

    async def refresh_user_context(self, user_id: int, token_limit: int) -> bool:
        """
        Merge any new cross-chat messages that arrived since this user's context was last loaded.

        Compares the current max message id in the cross-chat context against the stored
        cursor for this user. If new rows exist, loads only the delta and merges it into
        the in-memory conversation, respecting the token budget.

        This handles the mid-session staleness case: a user whose context was loaded
        at session start (with no private history) will pick up new private messages they
        send before returning to the group, without requiring a full reload.

        Args:
            user_id: Telegram user ID of the user to refresh.
            token_limit: Maximum token budget; new messages are discarded if they would
                         push the conversation over this limit.

        Returns:
            True if at least one new message was merged, False otherwise.

        Side Effects:
            - Modifies self.messages in-place, inserting new rows at self._history_end[user_id]
              (after historical context, before any live session messages).
            - Advances self._context_cursor[user_id] and self._history_end[user_id] on success.
        """
        if user_id not in self._context_cursor:
            return False

        exclude_private = True
        current_max = await get_max_cross_chat_message_id(user_id, self.chat_id, exclude_private)

        if current_max <= self._context_cursor[user_id]:
            return False

        new_rows = await load_new_cross_chat_context(
            user_id, self.chat_id, self._context_cursor[user_id], exclude_private
        )

        if not new_rows:
            return False

        # Insert at the tracked history boundary so delta rows land after existing historical
        # context and before any live session messages, preserving chronological ordering.
        insert_pos = self._history_end.get(user_id, 1)
        rows_inserted = 0
        for row in reversed(new_rows):
            self.messages.insert(insert_pos, row)
            token_count = await self.get_message_token_count()
            if token_count > token_limit:
                self.messages.pop(insert_pos)
                break
            rows_inserted += 1

        if rows_inserted > 0:
            self._context_cursor[user_id] = current_max
            self._history_end[user_id] = insert_pos + rows_inserted

        return rows_inserted > 0

    async def clear_interaction(self, user_id: int):
        """
        Clear this user's messages from memory and from the database.

        In private chats, all messages in the chat are deleted (including bot replies),
        fully clearing the conversation. In group chats, only this user's rows are removed;
        other members' messages remain. Resets both _context_cursor and _history_end.

        Args:
            user_id: Telegram user ID of the user requesting the forget.

        Side Effects:
            - Clears self.messages to contain only the system message (index 0).
            - Resets self._context_cursor to empty dict.
            - Resets self._history_end to empty dict.
            - Resets self._private_message_ids to empty set.
            - Deletes all messages for this chat (private) or this user (group) from database.
        """
        self.messages = [self.messages[0]]
        self._context_cursor = {}
        self._history_end = {}
        self._private_message_ids = set()
        if self.chat_type == 'private':
            await delete_messages_for_chat(self.chat_id)
        else:
            await delete_messages_for_user(user_id)
