import asyncpg

from utils.search.models import (
    MAX_LIMIT,
    ChatKind,
    SearchFilters,
    SearchHit,
    SearchPage,
    decode_cursor,
    encode_cursor,
)

_PARSE = """
SELECT (regexp_replace(websearch_to_tsquery('russian', q)::text,
                       '''((?:[^'']|''''){3,})''', '''\\1'':*', 'g')::tsquery
     || regexp_replace(websearch_to_tsquery('simple', q)::text,
                       '''((?:[^'']|''''){3,})''', '''\\1'':*', 'g')::tsquery
       )::text AS tsquery,
       tsvector_to_array(search_vector(q)) AS terms
FROM translate($1, 'Ёё', 'Ее') AS q
"""

_CHANNELS = """
SELECT coalesce(array_agg(chat_id), '{}'::bigint[]) FROM (
    SELECT DISTINCT ON (chat_id) chat_id,
           coalesce(raw->'chat'->>'type', raw->>'type') AS type
    FROM chat_history
    WHERE account_id = $1
      AND coalesce(raw->'chat'->>'type', raw->>'type') IS NOT NULL
    ORDER BY chat_id, ts DESC
) typed WHERE type = 'ChatType.CHANNEL'
"""

_KEY_ORDER = "s.date DESC, s.chat_id DESC, s.message_id DESC"

_COLS = (
    "s.chat_id, s.message_id, s.date, s.sender_id, s.text, s.has_media, "
    "s.edited_at, s.deleted_at"
)

_TEXT = (
    f"SELECT {_COLS}, 0 AS rank, NULL::text AS extracted_text, "  # noqa: S608
    "NULL::text AS version_text, NULL::timestamptz AS version_at "
    "FROM message_search s WHERE {where}"
)

_PROBE = (
    "CROSS JOIN LATERAL (SELECT * FROM message_search s "
    "WHERE s.chat_id = {source}.chat_id AND s.message_id = {source}.message_id "
    "AND {where} LIMIT 1) s "
)

_STT = (
    f"SELECT {_COLS}, 1 AS rank, md.extracted_text, NULL::text, "  # noqa: S608
    "NULL::timestamptz FROM media md "
    f"{_PROBE.replace('{source}', 'md')}"
    "WHERE md.account_id = {account} "
    "AND search_vector(md.extracted_text) @@ {tsquery}"
)

_VERSION = (
    "SELECT DISTINCT ON (s.date, s.chat_id, s.message_id) "  # noqa: S608
    f"{_COLS}, 2 AS rank, NULL::text, v.text, v.observed_at "
    "FROM message_versions v "
    f"{_PROBE.replace('{source}', 'v')}"
    "WHERE v.account_id = {account} "
    "AND search_vector(v.text) @@ {tsquery}"
)

_WITH_MEDIA = (
    "SELECT h.*, md.kind AS media_kind FROM ({hits}) h "
    "LEFT JOIN media md ON md.account_id = {account} "
    "AND md.chat_id = h.chat_id AND md.message_id = h.message_id "
    "ORDER BY h.date DESC, h.chat_id DESC, h.message_id DESC"
)

_SOURCES = ("text", "stt", "version")


class _Binder:
    def __init__(self) -> None:
        self.params: list[object] = []

    def __call__(self, value: object) -> str:
        self.params.append(value)
        return f"${len(self.params)}"


def _kind(kind: ChatKind, channels: list[int], bind: _Binder) -> str:
    if kind == "private":
        return "s.chat_id > 0"
    member = f"s.chat_id = ANY({bind(channels)}::bigint[])"
    return member if kind == "channel" else f"s.chat_id < 0 AND NOT {member}"


def _conditions(
    filters: SearchFilters, channels: list[int], bind: _Binder
) -> list[str]:
    conds = [f"s.account_id = {bind(filters.account_id)}"]
    compared = (
        ("s.chat_id =", filters.chat_id),
        ("s.sender_id =", filters.sender_id),
        ("s.has_media =", filters.has_media),
        ("s.date >=", filters.date_from),
        ("s.date <=", filters.date_to),
        ("s.text ~", filters.regex),
    )
    conds += [f"{left} {bind(value)}" for left, value in compared if value is not None]
    flags = (("s.deleted_at", filters.deleted), ("s.edited_at", filters.edited))
    conds += [
        f"{column} IS {'NOT ' if flag else ''}NULL"
        for column, flag in flags
        if flag is not None
    ]
    if filters.kind is not None:
        conds.append(_kind(filters.kind, channels, bind))
    if filters.cursor is not None:
        date, chat_id, message_id = decode_cursor(filters.cursor)
        conds.append(
            "(s.date, s.chat_id, s.message_id) < "
            f"({bind(date)}, {bind(chat_id)}, {bind(message_id)})"
        )
    return conds


def _branch(sql: str, limit: str, order: str = _KEY_ORDER, **parts: str) -> str:
    return f"({sql.format(**parts)} ORDER BY {order} LIMIT {limit})"


async def search_messages(pool: asyncpg.Pool, filters: SearchFilters) -> SearchPage:
    async with pool.acquire() as conn, conn.transaction():
        await conn.execute("SET LOCAL jit = off")
        return await _search(conn, filters)


async def _search(conn: asyncpg.Connection, filters: SearchFilters) -> SearchPage:
    query = (filters.query or "").strip()
    terms: list[str] = []
    tsquery = None
    if query:
        parsed = await conn.fetchrow(_PARSE, query)
        terms = list(parsed["terms"])
        if not parsed["tsquery"]:
            return SearchPage(hits=[], next_cursor=None, terms=terms)
        tsquery = parsed["tsquery"]

    channels: list[int] = []
    if filters.kind in {"group", "channel"}:
        channels = await conn.fetchval(_CHANNELS, filters.account_id)
    bind = _Binder()
    where = " AND ".join(_conditions(filters, channels, bind))
    size = min(filters.limit, MAX_LIMIT)
    limit = bind(size + 1)

    account = bind(filters.account_id)
    if tsquery is None:
        sql = _branch(_TEXT, limit, where=where)
    else:
        match = f"{bind(tsquery)}::tsquery"
        branches = (
            _branch(_TEXT, limit, where=f"{where} AND s.tsv @@ {match}"),
            _branch(_STT, limit, where=where, account=account, tsquery=match),
            _branch(
                _VERSION,
                limit,
                f"{_KEY_ORDER}, v.observed_at DESC",
                where=where,
                account=account,
                tsquery=match,
            ),
        )
        sql = (
            "SELECT DISTINCT ON (date, chat_id, message_id) * FROM ("  # noqa: S608
            f"{' UNION ALL '.join(branches)}) s "
            f"ORDER BY {_KEY_ORDER}, rank LIMIT {limit}"
        )

    rows = await conn.fetch(_WITH_MEDIA.format(hits=sql, account=account), *bind.params)
    hits = [
        SearchHit(**{**dict(row), "source": _SOURCES[row["rank"]]})
        for row in rows[:size]
    ]
    next_cursor = None
    if len(rows) > size:
        last = hits[-1]
        next_cursor = encode_cursor(last.date, last.chat_id, last.message_id)
    return SearchPage(hits=hits, next_cursor=next_cursor, terms=terms)