"""
cortex/hippo/store/conversation_store.py — CRUD for conversations and messages.

All queries filter by owner_id for multi-tenant isolation.
Issue: #801, #802
"""

from __future__ import annotations

import logging
from typing import Optional

import aiosqlite

from cortex.hippo.models import (
    ConversationRecord,
    MessageRecord,
)

logger = logging.getLogger(__name__)


def _estimate_tokens(text: str) -> int:
    """Estimate token count (~4 chars per token)."""
    return max(1, len(text) // 4)


class ConversationStore:
    """Data access for hippo_conversations and hippo_messages tables."""

    def __init__(self, conn: aiosqlite.Connection) -> None:
        self._conn = conn

    # ── Conversations ──────────────────────────────────────────────────

    async def get_or_create_conversation(
        self, owner_id: str, session_id: str, session_key: Optional[str] = None
    ) -> ConversationRecord:
        """Get an existing conversation or create a new one.

        Lookup priority:
        1. If session_key provided → lookup by (owner_id, session_key)
           - If found: update session_id to current value, return existing
        2. Fallback to (owner_id, session_id) lookup (backward compat)
        3. If nothing found: create new with both session_id and session_key
        """
        # Priority 1: lookup by session_key (if provided)
        if session_key:
            async with self._conn.execute(
                """
                SELECT conversation_id, owner_id, session_id, title,
                       bootstrapped_at, created_at, updated_at, session_key
                FROM hippo_conversations
                WHERE owner_id = ? AND session_key = ?
                """,
                (owner_id, session_key),
            ) as cur:
                row = await cur.fetchone()
                if row:
                    conv_id = row[0]
                    old_session_id = row[2]
                    # Update session_id if it changed (session recycling)
                    if old_session_id != session_id:
                        await self._conn.execute(
                            """
                            UPDATE hippo_conversations
                            SET session_id = ?, updated_at = strftime('%Y-%m-%dT%H:%M:%SZ', 'now')
                            WHERE conversation_id = ?
                            """,
                            (session_id, conv_id),
                        )
                        await self._conn.commit()
                    return ConversationRecord(
                        conversation_id=conv_id,
                        owner_id=row[1],
                        session_id=session_id,
                        title=row[3],
                        bootstrapped_at=row[4],
                        created_at=row[5],
                        updated_at=row[6],
                        session_key=row[7],
                    )

        # Priority 2: lookup by session_id (backward compat)
        async with self._conn.execute(
            """
            SELECT conversation_id, owner_id, session_id, title,
                   bootstrapped_at, created_at, updated_at, session_key
            FROM hippo_conversations
            WHERE owner_id = ? AND session_id = ?
            """,
            (owner_id, session_id),
        ) as cur:
            row = await cur.fetchone()
            if row:
                return ConversationRecord(
                    conversation_id=row[0],
                    owner_id=row[1],
                    session_id=row[2],
                    title=row[3],
                    bootstrapped_at=row[4],
                    created_at=row[5],
                    updated_at=row[6],
                    session_key=row[7],
                )

        # Create new conversation with both session_id and session_key
        async with self._conn.execute(
            """
            INSERT INTO hippo_conversations (owner_id, session_id, session_key)
            VALUES (?, ?, ?)
            """,
            (owner_id, session_id, session_key),
        ) as cur:
            conv_id = cur.lastrowid

        await self._conn.commit()

        return ConversationRecord(
            conversation_id=conv_id,
            owner_id=owner_id,
            session_id=session_id,
            session_key=session_key,
        )

    async def get_conversation(
        self, owner_id: str, conversation_id: int
    ) -> Optional[ConversationRecord]:
        """Get a conversation by ID."""
        async with self._conn.execute(
            """
            SELECT conversation_id, owner_id, session_id, title,
                   bootstrapped_at, created_at, updated_at, session_key
            FROM hippo_conversations
            WHERE owner_id = ? AND conversation_id = ?
            """,
            (owner_id, conversation_id),
        ) as cur:
            row = await cur.fetchone()
            if not row:
                return None
            return ConversationRecord(
                conversation_id=row[0],
                owner_id=row[1],
                session_id=row[2],
                title=row[3],
                bootstrapped_at=row[4],
                created_at=row[5],
                updated_at=row[6],
                session_key=row[7],
            )

    # ── Messages ───────────────────────────────────────────────────────

    async def get_next_seq(self, conversation_id: int) -> int:
        """Get the next sequence number for a conversation."""
        async with self._conn.execute(
            "SELECT COALESCE(MAX(seq), 0) + 1 FROM hippo_messages WHERE conversation_id = ?",
            (conversation_id,),
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else 1

    async def create_message(
        self,
        conversation_id: int,
        owner_id: str,
        role: str,
        content: str,
        metadata: Optional[str] = None,
        token_count: Optional[int] = None,
    ) -> MessageRecord:
        """Insert a new message into the immutable store."""
        seq = await self.get_next_seq(conversation_id)
        if token_count is None:
            token_count = _estimate_tokens(content)

        async with self._conn.execute(
            """
            INSERT INTO hippo_messages
                (conversation_id, owner_id, seq, role, content, token_count, metadata)
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
            (conversation_id, owner_id, seq, role, content, token_count, metadata),
        ) as cur:
            msg_id = cur.lastrowid

        await self._conn.commit()

        return MessageRecord(
            message_id=msg_id,
            conversation_id=conversation_id,
            owner_id=owner_id,
            seq=seq,
            role=role,
            content=content,
            token_count=token_count,
        )

    async def get_messages(
        self,
        conversation_id: int,
        owner_id: str,
        limit: Optional[int] = None,
        offset: int = 0,
    ) -> list[MessageRecord]:
        """Get messages for a conversation, ordered by seq."""
        sql = """
            SELECT message_id, conversation_id, owner_id, seq, role,
                   content, token_count, created_at
            FROM hippo_messages
            WHERE conversation_id = ? AND owner_id = ?
            ORDER BY seq ASC
        """
        params: list = [conversation_id, owner_id]
        if limit is not None:
            sql += " LIMIT ? OFFSET ?"
            params.extend([limit, offset])

        async with self._conn.execute(sql, params) as cur:
            rows = await cur.fetchall()
            return [
                MessageRecord(
                    message_id=r[0],
                    conversation_id=r[1],
                    owner_id=r[2],
                    seq=r[3],
                    role=r[4],
                    content=r[5],
                    token_count=r[6],
                    created_at=r[7],
                )
                for r in rows
            ]

    async def get_tail_messages(
        self, conversation_id: int, owner_id: str, count: int
    ) -> list[MessageRecord]:
        """Get the last N messages (fresh tail)."""
        sql = """
            SELECT message_id, conversation_id, owner_id, seq, role,
                   content, token_count, created_at
            FROM hippo_messages
            WHERE conversation_id = ? AND owner_id = ?
            ORDER BY seq DESC
            LIMIT ?
        """
        async with self._conn.execute(sql, (conversation_id, owner_id, count)) as cur:
            rows = await cur.fetchall()
            # Reverse to get chronological order
            return [
                MessageRecord(
                    message_id=r[0],
                    conversation_id=r[1],
                    owner_id=r[2],
                    seq=r[3],
                    role=r[4],
                    content=r[5],
                    token_count=r[6],
                    created_at=r[7],
                )
                for r in reversed(rows)
            ]

    async def get_message_count(self, conversation_id: int, owner_id: str) -> int:
        """Get total message count for a conversation."""
        async with self._conn.execute(
            "SELECT COUNT(*) FROM hippo_messages WHERE conversation_id = ? AND owner_id = ?",
            (conversation_id, owner_id),
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else 0


    async def search_messages(
        self,
        owner_id: str,
        pattern: str,
        mode: str = "regex",
        conversation_id: Optional[int] = None,
        limit: int = 50,
    ) -> list[MessageRecord]:
        """Search messages via FTS5 (full_text) or LIKE (regex/fallback).

        Parameters
        ----------
        owner_id : str
            Tenant filter.
        pattern : str
            Search pattern.
        mode : str
            'full_text' for FTS5, 'regex' for LIKE.
        conversation_id : int, optional
            Restrict to a single conversation.
        limit : int
            Max results.
        """
        if mode == "full_text":
            return await self._search_messages_fts(owner_id, pattern, conversation_id, limit)
        return await self._search_messages_like(owner_id, pattern, conversation_id, limit)

    async def _search_messages_fts(
        self, owner_id: str, pattern: str,
        conversation_id: Optional[int], limit: int,
    ) -> list[MessageRecord]:
        """FTS5-based message search."""
        from cortex.hippo.store.fts_utils import sanitize_fts5_query

        fts_query = sanitize_fts5_query(pattern)
        if not fts_query:
            return await self._search_messages_like(owner_id, pattern, conversation_id, limit)

        conv_filter = "AND m.conversation_id = ?" if conversation_id else ""
        params: list = [fts_query, owner_id]
        if conversation_id:
            params.append(conversation_id)
        params.append(limit)

        try:
            async with self._conn.execute(
                f"""
                SELECT m.message_id, m.conversation_id, m.owner_id, m.seq,
                       m.role, m.content, m.token_count, m.created_at
                FROM hippo_fts_messages fts
                JOIN hippo_messages m ON m.message_id = fts.rowid
                WHERE hippo_fts_messages MATCH ? AND m.owner_id = ? {conv_filter}
                ORDER BY rank
                LIMIT ?
                """,
                params,
            ) as cur:
                rows = await cur.fetchall()
                return [
                    MessageRecord(
                        message_id=r[0], conversation_id=r[1], owner_id=r[2],
                        seq=r[3], role=r[4], content=r[5], token_count=r[6],
                        created_at=r[7],
                    )
                    for r in rows
                ]
        except Exception:
            logger.debug("FTS5 search failed, falling back to LIKE")
            return await self._search_messages_like(owner_id, pattern, conversation_id, limit)

    async def _search_messages_like(
        self, owner_id: str, pattern: str,
        conversation_id: Optional[int], limit: int,
    ) -> list[MessageRecord]:
        """LIKE-based message search fallback."""
        from cortex.hippo.store.utils import escape_like_pattern

        conv_filter = "AND conversation_id = ?" if conversation_id else ""
        escaped = escape_like_pattern(pattern)
        params: list = [owner_id, f"%{escaped}%"]
        if conversation_id:
            params.append(conversation_id)
        params.append(limit)

        async with self._conn.execute(
            f"""
            SELECT message_id, conversation_id, owner_id, seq, role,
                   content, token_count, created_at
            FROM hippo_messages
            WHERE owner_id = ? AND content LIKE ? ESCAPE '\\' {conv_filter}
            ORDER BY created_at DESC
            LIMIT ?
            """,
            params,
        ) as cur:
            rows = await cur.fetchall()
            return [
                MessageRecord(
                    message_id=r[0], conversation_id=r[1], owner_id=r[2],
                    seq=r[3], role=r[4], content=r[5], token_count=r[6],
                    created_at=r[7],
                )
                for r in rows
            ]

    async def get_message_by_id(self, message_id: int, owner_id: Optional[str] = None) -> Optional[MessageRecord]:
        """Get a single message by ID.

        When *owner_id* is provided, only returns the message if it belongs
        to that owner (tenant isolation).
        """
        if owner_id is not None:
            sql = """
                SELECT message_id, conversation_id, owner_id, seq, role,
                       content, token_count, created_at
                FROM hippo_messages WHERE message_id = ? AND owner_id = ?
            """
            params: tuple = (message_id, owner_id)
        else:
            sql = """
                SELECT message_id, conversation_id, owner_id, seq, role,
                       content, token_count, created_at
                FROM hippo_messages WHERE message_id = ?
            """
            params = (message_id,)
        async with self._conn.execute(sql, params) as cur:
            r = await cur.fetchone()
            if not r:
                return None
            return MessageRecord(
                message_id=r[0], conversation_id=r[1], owner_id=r[2],
                seq=r[3], role=r[4], content=r[5], token_count=r[6],
                created_at=r[7],
            )

    # ── Message Parts ────────────────────────────────────────────────

    async def create_message_parts(
        self, message_id: int, parts: list[dict]
    ) -> list[int]:
        """Insert message parts for a given message. Returns list of inserted IDs."""
        ids: list[int] = []
        for i, part in enumerate(parts):
            async with self._conn.execute(
                """
                INSERT INTO hippo_message_parts
                    (message_id, ordinal, part_type, text_content,
                     tool_name, tool_call_id, tool_input, tool_output, metadata)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    message_id,
                    part.get("ordinal", i),
                    part.get("part_type", "text"),
                    part.get("text_content"),
                    part.get("tool_name"),
                    part.get("tool_call_id"),
                    part.get("tool_input"),
                    part.get("tool_output"),
                    part.get("metadata"),
                ),
            ) as cur:
                ids.append(cur.lastrowid)
        await self._conn.commit()
        return ids

    async def get_message_parts(self, message_id: int) -> list[dict]:
        """Get all parts for a message, ordered by ordinal."""
        async with self._conn.execute(
            """
            SELECT id, message_id, ordinal, part_type, text_content,
                   tool_name, tool_call_id, tool_input, tool_output,
                   metadata, created_at
            FROM hippo_message_parts
            WHERE message_id = ?
            ORDER BY ordinal ASC
            """,
            (message_id,),
        ) as cur:
            rows = await cur.fetchall()
            return [
                {
                    "id": r[0],
                    "message_id": r[1],
                    "ordinal": r[2],
                    "part_type": r[3],
                    "text_content": r[4],
                    "tool_name": r[5],
                    "tool_call_id": r[6],
                    "tool_input": r[7],
                    "tool_output": r[8],
                    "metadata": r[9],
                    "created_at": r[10],
                }
                for r in rows
            ]

    async def get_message_parts_bulk(
        self, message_ids: list[int]
    ) -> dict[int, list[dict]]:
        """Batch-fetch parts for multiple messages. Returns {message_id: [parts]}."""
        if not message_ids:
            return {}

        placeholders = ",".join("?" for _ in message_ids)
        async with self._conn.execute(
            f"""
            SELECT id, message_id, ordinal, part_type, text_content,
                   tool_name, tool_call_id, tool_input, tool_output,
                   metadata, created_at
            FROM hippo_message_parts
            WHERE message_id IN ({placeholders})
            ORDER BY message_id, ordinal ASC
            """,
            message_ids,
        ) as cur:
            rows = await cur.fetchall()

        result: dict[int, list[dict]] = {mid: [] for mid in message_ids}
        for r in rows:
            result.setdefault(r[1], []).append(
                {
                    "id": r[0],
                    "message_id": r[1],
                    "ordinal": r[2],
                    "part_type": r[3],
                    "text_content": r[4],
                    "tool_name": r[5],
                    "tool_call_id": r[6],
                    "tool_input": r[7],
                    "tool_output": r[8],
                    "metadata": r[9],
                    "created_at": r[10],
                }
            )
        return result

    async def mark_conversation_bootstrapped(self, conversation_id: int) -> None:
        """Set bootstrapped_at to current timestamp."""
        await self._conn.execute(
            """
            UPDATE hippo_conversations
            SET bootstrapped_at = datetime('now')
            WHERE conversation_id = ?
            """,
            (conversation_id,),
        )
        await self._conn.commit()

    async def get_messages_by_ids(
        self,
        message_ids: list[int],
        owner_id: Optional[str] = None,
    ) -> dict[int, MessageRecord]:
        """Batch-fetch messages by IDs with optional tenant scoping.

        Args:
            message_ids: List of message IDs to fetch. Empty list returns {}.
            owner_id: When provided, restricts results to messages owned by
                this tenant. The IN-clause uses parameterized ``?``
                placeholders — the IDs are never interpolated as SQL text,
                so there is no SQL-injection risk.

        Returns:
            Mapping of message_id → MessageRecord for found messages only.
        """
        if not message_ids:
            return {}

        placeholders = ",".join("?" for _ in message_ids)
        owner_filter = "AND owner_id = ?" if owner_id is not None else ""
        params: list = list(message_ids)
        if owner_id is not None:
            params.append(owner_id)
        async with self._conn.execute(
            f"""
            SELECT message_id, conversation_id, owner_id, seq, role,
                   content, token_count, created_at
            FROM hippo_messages
            WHERE message_id IN ({placeholders}) {owner_filter}
            """,
            params,
        ) as cur:
            rows = await cur.fetchall()
            return {
                r[0]: MessageRecord(
                    message_id=r[0],
                    conversation_id=r[1],
                    owner_id=r[2],
                    seq=r[3],
                    role=r[4],
                    content=r[5],
                    token_count=r[6],
                    created_at=r[7],
                )
                for r in rows
            }

    async def get_total_message_tokens(self, conversation_id: int, owner_id: str) -> int:
        """Get total token count across all messages in a conversation."""
        async with self._conn.execute(
            """
            SELECT COALESCE(SUM(token_count), 0)
            FROM hippo_messages
            WHERE conversation_id = ? AND owner_id = ?
            """,
            (conversation_id, owner_id),
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else 0
