"""
SQLite-backed conversation storage for TeLLMgramBot.

Provides async helpers for persisting and loading messages, managing per-user private mode,
and retrieving conversation history. The schema includes a messages table (indexed by chat_id
and user_id) and a private_mode table for tracking user preferences.
"""
import os
from typing import Optional

import aiosqlite
from .utils import execution_dir

# Allows tests to override the DB path without touching env vars
_DB_PATH: Optional[str] = None

# Maximum bound parameters per SQLite statement (SQLite default limit is 999).
# Defined at module level so tests can patch it to force batching with small inputs.
_SQLITE_MAX_PARAMS = 999

_DDL = """
PRAGMA journal_mode=WAL;

CREATE TABLE IF NOT EXISTS messages (
    id         INTEGER PRIMARY KEY AUTOINCREMENT,
    chat_id    INTEGER NOT NULL,
    user_id    INTEGER NOT NULL,
    username   TEXT,
    role       TEXT NOT NULL CHECK(role IN ('user', 'assistant', 'system')),
    content    TEXT NOT NULL,
    is_private INTEGER NOT NULL DEFAULT 0,
    created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
);

CREATE INDEX IF NOT EXISTS idx_messages_chat_id ON messages(chat_id, created_at);
CREATE INDEX IF NOT EXISTS idx_messages_user_id ON messages(user_id, created_at);

CREATE TABLE IF NOT EXISTS private_mode (
    user_id    INTEGER PRIMARY KEY,
    enabled    INTEGER NOT NULL DEFAULT 0,
    updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
);
"""

def get_db_path() -> str:
    """
    Return the SQLite database file path.

    Reads TELLMGRAMBOT_DATA_PATH env var (directory path) and appends conversations.db.
    Falls back to data/conversations.db relative to the execution directory if unset.
    Module-level _DB_PATH overrides everything - used by tests.
    """
    if _DB_PATH:
        return _DB_PATH
    data_dir = os.environ.get('TELLMGRAMBOT_DATA_PATH', os.path.join(execution_dir(), 'data'))
    return os.path.join(data_dir, 'conversations.db')

async def init_db() -> None:
    """
    Create the messages and private_mode tables if they do not exist.
    Enables WAL journal mode for safe concurrent reads and writes.
    Safe to call repeatedly - all statements use IF NOT EXISTS.
    """
    async with aiosqlite.connect(get_db_path()) as db:
        await db.executescript(_DDL)
        await db.commit()

async def insert_message(
    chat_id: int,
    user_id: int,
    username: Optional[str],
    role: str,
    content: str,
    is_private: bool = False,
) -> None:
    """
    Persist a single message row.

    Args:
        chat_id: Telegram chat ID (negative for groups, positive/==user_id for private).
        user_id: Telegram user ID of the sender.
        username: Telegram username (display only, may be None).
        role: 'user', 'assistant', or 'system'.
        content: Message text.
        is_private: If True, this message is excluded from group context loading.
    """
    async with aiosqlite.connect(get_db_path()) as db:
        await db.execute(
            "INSERT INTO messages (chat_id, user_id, username, role, content, is_private) "
            "VALUES (?, ?, ?, ?, ?, ?)",
            (chat_id, user_id, username, role, content, int(is_private)),
        )
        await db.commit()

async def load_messages_for_chat(
    chat_id: int,
    exclude_private: bool = False,
) -> list[dict]:
    """
    Load all messages for a given chat in chronological order.

    Args:
        chat_id: Telegram chat ID to query.
        exclude_private: If True, rows with is_private=1 are excluded.
            Used when loading group conversation context.

    Returns:
        List of {"role": ..., "content": ...} dicts.
    """
    query = "SELECT role, content FROM messages WHERE chat_id = ?"
    params: tuple = (chat_id,)
    if exclude_private:
        query += " AND is_private = 0"
    query += " ORDER BY created_at ASC, id ASC"
    async with aiosqlite.connect(get_db_path()) as db:
        async with db.execute(query, params) as cursor:
            rows = await cursor.fetchall()
    return [{"role": row[0], "content": row[1]} for row in rows]

async def load_messages_for_user(
    user_id: int,
    exclude_private: bool = False,
) -> list[dict]:
    """
    Load all messages from a user's private chat history.

    In Telegram, a user's private chat_id equals their user_id. This function loads messages
    from that private chat, optionally filtering out private-mode messages.

    Args:
        user_id: Telegram user ID; also serves as the private chat_id.
        exclude_private: If True, rows with is_private=1 are excluded.

    Returns:
        List of {"role": ..., "content": ...} dicts in chronological order.
    """
    query = "SELECT role, content FROM messages WHERE chat_id = ?"
    params: tuple = (user_id,)
    if exclude_private:
        query += " AND is_private = 0"
    query += " ORDER BY created_at ASC, id ASC"
    async with aiosqlite.connect(get_db_path()) as db:
        async with db.execute(query, params) as cursor:
            rows = await cursor.fetchall()
    return [{"role": row[0], "content": row[1]} for row in rows]

async def load_full_user_context(
    user_id: int,
    current_chat_id: int,
    exclude_private: bool = False,
) -> list[dict]:
    """
    Load bidirectional conversation context for a user across all chats.

    Implements bidirectional cross-pollination: returns all messages from the current
    chat (all participants) plus all messages from the user's private chat (both user
    and assistant turns), merged in chronological order. This eliminates amnesia when
    a user switches between private and group contexts.

    In Telegram, a user's private chat_id always equals their user_id. Three arms cover
    all cross-context cases:
      Arm 1 - chat_id = current_chat_id:
          All messages (user and bot) in the current chat. Always included.
      Arm 2 - chat_id = user_id AND chat_id != current_chat_id:
          All messages (both sides) from the user's private chat, included when the
          current chat is a group. Pulling both sides restores proper role alternation
          that was broken when only the user's own rows were fetched.
      Arm 3 - user_id = user_id AND chat_id != current_chat_id AND chat_id != user_id:
          The requesting user's own messages from any other group chats (not the private
          chat, not the current chat). Covers group-to-private and group-to-group context.
          Only the user's own rows are fetched from other groups (not bot replies).

    The exclude_private flag is controlled by the caller and applies regardless of chat type.
    Conversation.get_past_interaction always passes exclude_private=True so that is_private=1
    rows are never rehydrated into any session context.

    Args:
        user_id: Telegram user ID of the requesting user.
        current_chat_id: The chat_id of the current conversation (negative for groups, positive for private).
        exclude_private: If True, rows with is_private=1 are excluded from all arms of the query.
            Always passed as True by Conversation.get_past_interaction so that private-mode messages
            don't surface in shared contexts.

    Returns:
        List of {"role": ..., "content": ...} dicts in ascending chronological order (oldest first).
    """
    query = (
        "SELECT role, content FROM messages "
        "WHERE (chat_id = ? "
        "OR (chat_id = ? AND chat_id != ?) "
        "OR (user_id = ? AND chat_id != ? AND chat_id != ?))"
    )
    params: tuple = (current_chat_id, user_id, current_chat_id, user_id, current_chat_id, user_id)
    if exclude_private:
        query += " AND is_private = 0"
    query += " ORDER BY created_at ASC, id ASC"
    async with aiosqlite.connect(get_db_path()) as db:
        async with db.execute(query, params) as cursor:
            rows = await cursor.fetchall()
    return [{"role": row[0], "content": row[1]} for row in rows]

async def get_max_cross_chat_message_id(
    user_id: int,
    current_chat_id: int,
    exclude_private: bool = False,
) -> int:
    """
    Return the maximum message id from the cross-chat context for this user.

    Uses the same three-arm query as load_full_user_context, returning only the
    MAX(id) scalar. Used to establish a cursor for mid-session refresh detection.

    Returns:
        Maximum message id as an integer, or 0 if no rows exist.
    """
    query = (
        "SELECT MAX(id) FROM messages "
        "WHERE (chat_id = ? "
        "OR (chat_id = ? AND chat_id != ?) "
        "OR (user_id = ? AND chat_id != ? AND chat_id != ?))"
    )
    params: tuple = (current_chat_id, user_id, current_chat_id, user_id, current_chat_id, user_id)
    if exclude_private:
        query += " AND is_private = 0"
    async with aiosqlite.connect(get_db_path()) as db:
        async with db.execute(query, params) as cursor:
            row = await cursor.fetchone()
    return row[0] if row and row[0] is not None else 0

async def load_new_cross_chat_context(
    user_id: int,
    current_chat_id: int,
    since_id: int,
    exclude_private: bool = False,
) -> list[dict]:
    """
    Load cross-chat context rows with id strictly greater than since_id.

    Uses the same three-arm query as load_full_user_context but restricted to rows
    newer than the given cursor. Returns only the delta since the last context load.

    Args:
        user_id: Telegram user ID of the requesting user.
        current_chat_id: The chat_id of the current conversation.
        since_id: Exclusive lower bound - only rows with id > since_id are returned.
        exclude_private: If True, rows with is_private=1 are excluded.

    Returns:
        List of {"role": ..., "content": ...} dicts in ascending chronological order.
    """
    query = (
        "SELECT role, content FROM messages "
        "WHERE (chat_id = ? "
        "OR (chat_id = ? AND chat_id != ?) "
        "OR (user_id = ? AND chat_id != ? AND chat_id != ?)) "
        "AND id > ?"
    )
    params: tuple = (current_chat_id, user_id, current_chat_id, user_id, current_chat_id, user_id, since_id)
    if exclude_private:
        query += " AND is_private = 0"
    query += " ORDER BY created_at ASC, id ASC"
    async with aiosqlite.connect(get_db_path()) as db:
        async with db.execute(query, params) as cursor:
            rows = await cursor.fetchall()
    return [{"role": row[0], "content": row[1]} for row in rows]

async def get_shared_group_chat_ids(user_id: int, bot_id: int) -> list[int]:
    """
    Return chat_ids of groups where both the user and the bot have sent messages.

    Eligibility conditions:
    1. The requesting user has sent at least one message in the group.
    2. The bot has also sent at least one message in the group (bot is active there).

    Condition 2 prevents surfacing history from groups the bot was removed from
    or never properly joined.

    Args:
        user_id: Telegram user ID of the requesting user.
        bot_id: Telegram user ID of the bot account.

    Returns:
        List of negative chat_ids (Telegram group IDs).
    """
    query = """
        SELECT DISTINCT chat_id
        FROM messages
        WHERE user_id = ? AND chat_id < 0
        AND EXISTS (
            SELECT 1 FROM messages m2
            WHERE m2.chat_id = messages.chat_id AND m2.user_id = ?
        )
    """
    async with aiosqlite.connect(get_db_path()) as db:
        async with db.execute(query, (user_id, bot_id)) as cursor:
            rows = await cursor.fetchall()
    return [row[0] for row in rows]

async def load_shared_group_context(
    group_chat_ids: list[int],
    exclude_private: bool = True,
) -> list[dict]:
    """
    Load all messages from multiple group chats, merged in chronological order.

    Used to fill remaining token budget in a private chat with context from
    shared groups (full group context in private chat).

    Automatically batches the query if group_chat_ids exceeds _SQLITE_MAX_PARAMS
    (999, SQLite's bound parameter limit) to handle users who share many groups
    with the bot. Results are merged and sorted chronologically across batches.

    Args:
        group_chat_ids: List of group chat_ids (negative integers) to include.
        exclude_private: If True, rows with is_private=1 are excluded (default True).

    Returns:
        List of {"role": ..., "content": ...} dicts in ascending chronological order.
        Empty list if group_chat_ids is empty.
    """
    if not group_chat_ids:
        return []
    # SQLite has a max of 999 bound parameters per statement. Batch group_chat_ids to
    # avoid exceeding this limit for users who share many groups with the bot.
    all_rows: list[tuple] = []
    async with aiosqlite.connect(get_db_path()) as db:
        for i in range(0, len(group_chat_ids), _SQLITE_MAX_PARAMS):
            chunk = group_chat_ids[i:i + _SQLITE_MAX_PARAMS]
            placeholders = ','.join('?' * len(chunk))
            query = (
                f"SELECT role, content, created_at, id FROM messages "
                f"WHERE chat_id IN ({placeholders})"
            )
            if exclude_private:
                query += " AND is_private = 0"
            async with db.execute(query, tuple(chunk)) as cursor:
                all_rows.extend(await cursor.fetchall())
    all_rows.sort(key=lambda r: (r[2], r[3]))
    return [{"role": row[0], "content": row[1]} for row in all_rows]

async def delete_paired_bot_replies(chat_id: int, bot_id: int) -> None:
    """
    Delete bot reply rows in a group chat that are paired with deleted user messages or orphaned.

    Handles two cases in one pass:
    1. Paired bot replies: after delete_messages_for_user() removes a user's rows, the bot
       replies that immediately followed those messages become orphaned (no user row before
       them). These are caught by the orphan query below.
    2. Pre-existing orphans: assistant rows left behind by earlier /forget calls that deleted
       user rows but not their paired replies. Cleaned up opportunistically on every group
       /forget to heal historical debris.

    Both cases are unified by a single query: delete assistant rows whose immediately preceding
    message in the same chat is not a user message. Reliable because the bot processes messages
    sequentially - each reply is stored immediately after its user message, so the immediately
    preceding row unambiguously identifies the pairing. Other users' conversations are unaffected
    as long as their user rows remain (their user rows serve as the preceding message for their
    own bot replies).

    Args:
        chat_id: Telegram chat ID of the group chat to clean up.
        bot_id: Telegram user ID of the bot account (identifies assistant rows by user_id).
    """
    async with aiosqlite.connect(get_db_path()) as db:
        await db.execute(
            "DELETE FROM messages "
            "WHERE chat_id = ? AND role = 'assistant' AND user_id = ? "
            "AND NOT EXISTS ("
            "    SELECT 1 FROM messages prev "
            "    WHERE prev.chat_id = messages.chat_id "
            "      AND prev.id = ("
            "          SELECT MAX(m2.id) FROM messages m2 "
            "          WHERE m2.chat_id = messages.chat_id "
            "            AND m2.id < messages.id"
            "      ) "
            "      AND prev.role = 'user'"
            ")",
            (chat_id, bot_id),
        )
        await db.commit()

async def delete_messages_for_user(user_id: int) -> None:
    """
    Delete all message rows for a given user_id across all chats.
    Used by /forget in group chats - only the issuing user's rows are removed.
    """
    async with aiosqlite.connect(get_db_path()) as db:
        await db.execute("DELETE FROM messages WHERE user_id = ?", (user_id,))
        await db.commit()

async def delete_messages_for_chat(chat_id: int) -> None:
    """
    Delete all message rows for a given chat_id, including all participants and bot replies.
    Used by /forget in private chats to fully clear the conversation.
    """
    async with aiosqlite.connect(get_db_path()) as db:
        await db.execute("DELETE FROM messages WHERE chat_id = ?", (chat_id,))
        await db.commit()

async def delete_private_messages_for_user(user_id: int) -> None:
    """
    Delete all private-mode message rows for a user from their private chat.

    Removes rows where chat_id = user_id (the user's Telegram private chat) AND
    is_private = 1. This covers both the user's own messages and the bot's replies
    stored in the same private chat, since assistant turns share the same chat_id.

    Called when the user turns private mode off (/private off) to ensure that
    private-mode messages do not resurface on the next session reload.

    Args:
        user_id: Telegram user ID; also serves as the private chat_id.
    """
    async with aiosqlite.connect(get_db_path()) as db:
        await db.execute(
            "DELETE FROM messages WHERE chat_id = ? AND is_private = 1",
            (user_id,)
        )
        await db.commit()

async def get_private_mode(user_id: int) -> bool:
    """
    Return whether private mode is currently enabled for this user.
    Defaults to False if no row exists.
    """
    async with aiosqlite.connect(get_db_path()) as db:
        async with db.execute(
            "SELECT enabled FROM private_mode WHERE user_id = ?", (user_id,)
        ) as cursor:
            row = await cursor.fetchone()
    return bool(row[0]) if row else False

async def set_private_mode(user_id: int, enabled: bool) -> None:
    """
    Enable or disable private mode for a user.

    Upserts the private_mode row for this user_id. If the user already has a private_mode
    entry, it is updated; otherwise, a new row is created.

    Args:
        user_id: Telegram user ID.
        enabled: True to enable private mode, False to disable.
    """
    async with aiosqlite.connect(get_db_path()) as db:
        await db.execute(
            "INSERT INTO private_mode (user_id, enabled) VALUES (?, ?) "
            "ON CONFLICT(user_id) DO UPDATE SET enabled = excluded.enabled, "
            "updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')",
            (user_id, int(enabled)),
        )
        await db.commit()
