from __future__ import annotations

from typing import Any

from rlm_core.protocols.data_source import DataSource, DataSourceCapabilities, SearchResult

from ..message_database import MessageDatabase
from ..prompts import SCHEMA_CONTACTS


_BASE_SCHEMA = """messages: guid, chat_id, chat_name, sender, text, timestamp, is_from_me, is_group
messages_fts: FTS5 on (text, chat_name, sender)
events: id, title, start_date, end_date, location, notes, calendar_name, is_all_day, participants, recurrence_rule
events_fts: FTS5 on (title, location, notes)
file_cache: path, content_md, file_hash, parsed_at, parse_model, parse_method, page_count, token_cost
file_cache_fts: FTS5 on (path, content_md)"""


class MessageDataSource(DataSource):
    """DataSource for the Recall shadow.db database."""

    def __init__(self, db: MessageDatabase, *, has_contacts: bool = False) -> None:
        self._db = db
        self._has_contacts = has_contacts

    def name(self) -> str:
        return "messages"

    def capabilities(self) -> DataSourceCapabilities:
        return {
            "sql": True,
            "search": True,
            "get": True,
            "recent": True,
            "schema": True,
            "stats": True,
        }

    def get_schema(self) -> str:
        if self._has_contacts:
            return _BASE_SCHEMA + "\n" + SCHEMA_CONTACTS
        return _BASE_SCHEMA

    def get_stats(self) -> dict[str, Any]:
        return self._db.get_stats()

    def search(
        self,
        query: str,
        limit: int = 1000,
        sender: str | None = None,
        chat: str | None = None,
        days: int | None = None,
        table: str = "messages",
        filters: dict[str, Any] | None = None,
    ) -> list[SearchResult]:
        if filters:
            sender = filters.get("sender", sender)
            chat = filters.get("chat", chat)
            days = filters.get("days", days)
            table = filters.get("table", table)
        results = self._db.search(query, limit=limit, sender=sender, chat=chat, days=days, table=table)
        return [
            {
                **row,
                "id": row.get("guid") or row.get("id") or "",
                "snippet": row.get("text") or "",
                "metadata": row,
                "score": None,
            }
            for row in results
        ]

    def query(self, sql: str) -> list[dict[str, Any]]:
        return self._db.query(sql)

    def get(self, id: str) -> dict[str, Any] | None:
        rows = self._db.query("SELECT * FROM messages WHERE guid = '" + id + "' LIMIT 1")
        if rows:
            return rows[0]
        return None

    def get_recent(self, days: int = 30, limit: int = 10000) -> list[dict[str, Any]]:
        return self._db.get_recent(days, limit)

    def get_messages(self, limit: int | None = 10000) -> list[dict[str, Any]]:
        return self._db.get_messages(limit)

    def get_chats(self) -> list[dict[str, Any]]:
        return self._db.get_chats()

    def get_chat_messages(self, chat_id: str, limit: int = 1000) -> list[dict[str, Any]]:
        return self._db.get_chat_messages(chat_id, limit)

    def sample(self, n: int = 1000) -> list[dict[str, Any]]:
        return self._db.sample(n)

    def get_retrieved_guids(self) -> list[str]:
        return self._db.get_retrieved_guids()
