"""Fact CRUD and search — vector + ILIKE fallback."""

import re
import uuid
from datetime import datetime, timezone

import asyncpg
import structlog

from navi.config import settings

from ._embeddings import _vector_to_str

log = structlog.get_logger()

_SEPARATORS = re.compile(r"[-_/\\.]")
_NOISE = re.compile(r"[^\w\s]", re.UNICODE)
_KEY_SPLIT = re.compile(r"[-_/.\s]+")
_AUTO_DUMP_THRESHOLD = 60
_VECTOR_CUTOFF_DISTANCE = 0.3

# Leading tokens that carry no meaning of their own: `preferred_language`,
# `ui_language` and `language` are one key, and the extractor writes all three.
_MERGE_PREFIX_TOKENS = frozenset({"preferred", "preference", "pref", "user", "my", "ui", "response"})


def _normalize_query(query: str) -> list[str]:
    q = _SEPARATORS.sub(" ", query)
    q = _NOISE.sub(" ", q)
    return [t for t in q.lower().split() if len(t) > 1]


def _key_tokens(key: str) -> frozenset[str]:
    """The tokens of a fact key, with *leading* noise words dropped.

    Leading only, and that is load-bearing: `ui_language` is the language key with a
    noise prefix, but `dark_ui` is a name of its own and must not collapse into
    `dark`. A rule that strips `ui` wherever it stands cannot tell those apart.
    """
    tokens = [t for t in _KEY_SPLIT.split(key.strip().lower()) if len(t) > 1]
    while tokens and tokens[0] in _MERGE_PREFIX_TOKENS:
        tokens.pop(0)
    return frozenset(tokens)


class FactMixin:
    """Fact storage operations.

    Expected on the composite class:
        _get_pool() -> asyncpg.Pool
        _generate_embedding(text) -> list[float] | None
        _has_pgvector() -> bool
    """

    async def upsert_fact(
        self,
        category: str,
        key: str,
        value: str,
        user_id: str | None = None,
        source_session_id: str | None = None,
        source: str = "conversation",
        confidence: int = 70,
        expires_at: datetime | None = None,
        source_context: str = "",
        dedupe: bool = False,
    ) -> None:
        """Write a fact, keyed on `(user_id, category, key)`.

        With `dedupe`, the same fact already stored under a different spelling of
        this key is hidden behind the row that survives — see `_merge_key_aliases`.
        Off by default, so callers that own their keys (the `memory` tool) write
        exactly the row they asked for.
        """
        now = datetime.now(timezone.utc)
        embedding = await self._generate_embedding(value)
        vec_str = _vector_to_str(embedding) if embedding else None
        pool = await self._get_pool()
        fact_id = str(uuid.uuid4())
        async with pool.acquire() as conn, conn.transaction():
            # The dedupe path marks losers before the row they point at exists, so a
            # failed insert would leave them hidden behind nothing. The plain path
            # needs no boundary of its own, and taking one anyway costs nothing.
            if dedupe:
                await self._merge_key_aliases(conn, category, key, user_id, fact_id)
            if vec_str:
                await conn.execute(
                    """INSERT INTO memory_facts
                        (id, user_id, category, key, value, created_at, updated_at, source_session_id,
                         embedding, source, confidence, expires_at, source_context)
                       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9::vector, $10, $11, $12, $13)
                       ON CONFLICT(user_id, category, key) DO UPDATE SET
                           value             = EXCLUDED.value,
                           updated_at        = EXCLUDED.updated_at,
                           source_session_id = EXCLUDED.source_session_id,
                           embedding         = EXCLUDED.embedding,
                           source            = EXCLUDED.source,
                           confidence        = EXCLUDED.confidence,
                           expires_at        = COALESCE(EXCLUDED.expires_at, memory_facts.expires_at),
                           source_context    = EXCLUDED.source_context""",
                    fact_id, user_id, category, key, value, now, now,
                    source_session_id, vec_str, source, confidence, expires_at, source_context,
                )
            else:
                await conn.execute(
                    """INSERT INTO memory_facts
                        (id, user_id, category, key, value, created_at, updated_at, source_session_id,
                         source, confidence, expires_at, source_context)
                       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
                       ON CONFLICT(user_id, category, key) DO UPDATE SET
                           value             = EXCLUDED.value,
                           updated_at        = EXCLUDED.updated_at,
                           source_session_id = EXCLUDED.source_session_id,
                           source            = EXCLUDED.source,
                           confidence        = EXCLUDED.confidence,
                           expires_at        = COALESCE(EXCLUDED.expires_at, memory_facts.expires_at),
                           source_context    = EXCLUDED.source_context""",
                    fact_id, user_id, category, key, value, now, now,
                    source_session_id, source, confidence, expires_at, source_context,
                )

    async def _merge_key_aliases(
        self, conn, category: str, key: str, user_id: str | None, new_id: str
    ) -> None:
        """Hide the same fact already stored under a different spelling of this key.

        The exact-key upsert cannot see `preferred_language` when `language` is what
        is already there, which is how the base grew 11 spellings of the same handful
        of keys. Compared on key tokens and **not** on value: `ui_language` and
        `language` are one key while `dark` and `dark_ui` are two, and that is set
        equality, not overlap. Values are deliberately left out of the decision: the
        survivor is the row that already holds the key the writer used (else the
        incoming one), and since the upsert writes the fresh value into it, the
        fresher value is what remains readable either way.

        Only *leading* noise tokens are dropped (see `_key_tokens`), so a trailing
        one survives: `dark_ui` keeps its `ui` and stays apart from `dark`, while
        `ui_dark` loses it and does not. That asymmetry is the price of reading `ui`
        as noise in `ui_language` and as part of the name in `dark_ui`; no rule reads
        it both ways.

        Marked, never deleted: a loser keeps its own key and its own value, so the
        merge is reversible with `UPDATE ... SET merged_into = NULL`.
        """
        tokens = _key_tokens(key)
        if not tokens:
            return
        # Hidden rows are read too: one of them may hold this exact key, and then the
        # upsert below will land on it through ON CONFLICT rather than inserting the
        # row we are about to point everything at.
        rows = await conn.fetch(
            "SELECT id, key, merged_into FROM memory_facts "
            "WHERE user_id IS NOT DISTINCT FROM $1 AND category = $2 AND id <> $3",
            user_id, category, new_id,
        )
        exact = [r["id"] for r in rows if r["key"] == key]
        # Only live rows are worth hiding; the rest are already behind something.
        aliases = [
            r["id"] for r in rows
            if r["key"] != key and r["merged_into"] is None and _key_tokens(r["key"]) == tokens
        ]
        winner = exact[0] if exact else new_id
        if exact:
            # A row with this exact key is what the ON CONFLICT will update, so it
            # becomes live again — and it has to be un-hidden for that, or the write
            # lands in a row no read can see and the losers get pointed at a row that
            # was never inserted.
            await conn.execute(
                "UPDATE memory_facts SET merged_into = NULL WHERE id = ANY($1::text[])", exact
            )
        if not aliases:
            return
        # The second clause re-points the losers of an earlier merge, so pointers
        # never form a chain through a row that is itself hidden.
        await conn.execute(
            "UPDATE memory_facts SET merged_into = $2 "
            "WHERE (id = ANY($1::text[]) OR merged_into = ANY($1::text[])) AND id <> $2",
            aliases, winner,
        )
        log.info(
            "memory.key_alias_merged",
            category=category, key=key[:60], merged=len(aliases), winner=winner,
        )

    async def search_facts(self, query: str, user_id: str | None = None, limit: int = 15) -> list[dict]:
        # 1. Try vector search if pgvector + embedding backend are available
        if self._embedding_backend and await self._has_pgvector():
            query_embedding = await self._generate_embedding(query)
            vec_str = _vector_to_str(query_embedding) if query_embedding else None
            if vec_str:
                try:
                    pool = await self._get_pool()
                    async with pool.acquire() as conn:
                        if user_id is None:
                            rows = await conn.fetch(
                                """SELECT id, category, key, value, updated_at,
                                          source, confidence, expires_at, source_context, merged_into,
                                          embedding <=> $1::vector AS distance
                                   FROM memory_facts
                                   WHERE (expires_at IS NULL OR expires_at > now())
                                     AND merged_into IS NULL
                                     AND user_id IS NULL
                                   ORDER BY embedding <=> $1::vector
                                   LIMIT $2""",
                                vec_str, limit,
                            )
                        else:
                            rows = await conn.fetch(
                                """SELECT id, category, key, value, updated_at,
                                          source, confidence, expires_at, source_context, merged_into,
                                          embedding <=> $1::vector AS distance
                                   FROM memory_facts
                                   WHERE (expires_at IS NULL OR expires_at > now())
                                     AND merged_into IS NULL
                                     AND user_id = $3
                                   ORDER BY embedding <=> $1::vector
                                   LIMIT $2""",
                                vec_str, limit, user_id,
                            )
                        results = [
                            _row_to_dict(r)
                            for r in rows
                            if r["distance"] < _VECTOR_CUTOFF_DISTANCE
                        ]
                        if results:
                            log.debug("memory.vector_search", hits=len(results), query=query[:40])
                            return results
                except Exception:
                    log.warning("memory.vector_search_failed", query=query[:40], exc_info=True)
            else:
                log.warning("memory.vector_fallback_to_ilike", reason="embed_failed", query=query[:40])
        else:
            reason = "no_embedding_backend" if not self._embedding_backend else "no_pgvector"
            log.warning("memory.vector_fallback_to_ilike", reason=reason, query=query[:40])

        # 2. Text fallback (original ILIKE logic)
        terms = _normalize_query(query)
        if not terms:
            return await self.get_all_facts(user_id=user_id, limit=limit)

        if await self.fact_count(user_id=user_id) <= _AUTO_DUMP_THRESHOLD:
            return await self.get_all_facts(user_id=user_id, limit=limit)

        base_params: list = []
        term_conds: list[str] = []
        for term in terms:
            like = f"%{term}%"
            i = len(base_params) + 1
            term_conds.append(
                f"(category ILIKE ${i} OR key ILIKE ${i + 1} OR value ILIKE ${i + 2})"
            )
            base_params.extend([like, like, like])

        user_filter = "user_id IS NULL" if user_id is None else f"user_id = ${len(base_params) + 1}"
        if user_id is not None:
            base_params.append(user_id)
        # This branch used to be the only read path that served expired facts, and
        # hidden merge losers used to be served by all of them.
        user_filter += " AND merged_into IS NULL AND (expires_at IS NULL OR expires_at > now())"

        limit_idx = len(base_params) + 1
        and_where = " AND ".join(term_conds)
        or_where = " OR ".join(term_conds)

        pool = await self._get_pool()
        async with pool.acquire() as conn:
            rows = await conn.fetch(
                f"SELECT id, category, key, value, updated_at, source, confidence, "
                f"expires_at, source_context, merged_into FROM memory_facts "
                f"WHERE {user_filter} AND {and_where} ORDER BY updated_at DESC LIMIT ${limit_idx}",
                *base_params, limit,
            )
            if rows:
                return [_row_to_dict(r) for r in rows]

            score_expr = " + ".join(
                f"CASE WHEN {c} THEN 1 ELSE 0 END" for c in term_conds
            )
            rows = await conn.fetch(
                f"SELECT id, category, key, value, updated_at, source, confidence, "
                f"expires_at, source_context, merged_into, ({score_expr}) AS score "
                f"FROM memory_facts WHERE {user_filter} AND ({or_where}) "
                f"ORDER BY score DESC, updated_at DESC LIMIT ${limit_idx}",
                *base_params, limit,
            )
        return [_row_to_dict(r) for r in rows]

    async def delete_fact(self, key: str, category: str | None = None, user_id: str | None = None) -> int:
        pool = await self._get_pool()
        async with pool.acquire() as conn, conn.transaction():
            await self._delete_merged_dependents(conn, key, category, user_id)
            if category:
                if user_id is None:
                    result = await conn.execute(
                        "DELETE FROM memory_facts WHERE LOWER(key)=LOWER($1) AND LOWER(category)=LOWER($2) AND user_id IS NULL",
                        key, category,
                    )
                else:
                    result = await conn.execute(
                        "DELETE FROM memory_facts WHERE LOWER(key)=LOWER($1) AND LOWER(category)=LOWER($2) AND user_id = $3",
                        key, category, user_id,
                    )
            else:
                if user_id is None:
                    result = await conn.execute(
                        "DELETE FROM memory_facts WHERE LOWER(key)=LOWER($1) AND user_id IS NULL", key
                    )
                else:
                    result = await conn.execute(
                        "DELETE FROM memory_facts WHERE LOWER(key)=LOWER($1) AND user_id = $2", key, user_id
                    )
        return int(result.split()[1])

    async def _delete_merged_dependents(
        self, conn, key: str, category: str | None, user_id: str | None
    ) -> None:
        """Forget drops the hidden aliases of the keys it is about to delete.

        A merge loser points at its winner; deleting the winner on its own would
        leave those rows hidden with nothing left to point at and no way back to
        them. They are deleted, not un-hidden — the user asked for the key gone.
        """
        conditions = ["LOWER(key) = LOWER($1)"]
        args: list = [key]
        if category:
            args.append(category)
            conditions.append(f"LOWER(category) = LOWER(${len(args)})")
        if user_id is None:
            conditions.append("user_id IS NULL")
        else:
            args.append(user_id)
            conditions.append(f"user_id = ${len(args)}")
        await conn.execute(
            "DELETE FROM memory_facts WHERE merged_into IN "
            f"(SELECT id FROM memory_facts WHERE {' AND '.join(conditions)})",
            *args,
        )

    async def get_all_facts(
        self,
        user_id: str | None = None,
        limit: int | None = None,
        offset: int = 0,
        search: str | None = None,
        sort_by: str = "category",
        sort_order: str = "desc",
        all_users: bool = False,
        include_merged: bool = False,
    ) -> list[dict]:
        q = (
            "SELECT id, category, key, value, updated_at, source, confidence, "
            "expires_at, source_context, merged_into FROM memory_facts "
        )
        params: list = []
        param_idx = 0

        def add_param(value):
            nonlocal param_idx
            param_idx += 1
            params.append(value)
            return f"${param_idx}"

        conditions = []
        if not all_users:
            if user_id is None:
                conditions.append("user_id IS NULL")
            else:
                conditions.append(f"user_id = {add_param(user_id)}")
        if not include_merged:
            conditions.append("merged_into IS NULL")
        if search:
            like = f"%{search}%"
            conditions.append(
                f"(key ILIKE {add_param(like)} OR value ILIKE {add_param(like)} OR category ILIKE {add_param(like)})"
            )

        if conditions:
            q += "WHERE " + " AND ".join(conditions) + " "

        allowed_cols = {"category", "key", "updated_at", "confidence", "source"}
        col = sort_by if sort_by in allowed_cols else "category"
        order = "DESC" if sort_order == "desc" else "ASC"
        if col == "category":
            q += f"ORDER BY {col} {order}, updated_at DESC"
        else:
            q += f"ORDER BY {col} {order}"

        if limit is not None:
            q += f" LIMIT {add_param(limit)}"
        if offset:
            q += f" OFFSET {add_param(offset)}"

        pool = await self._get_pool()
        async with pool.acquire() as conn:
            rows = await conn.fetch(q, *params)
        return [_row_to_dict(r) for r in rows]

    async def fact_count(
        self,
        user_id: str | None = None,
        all_users: bool = False,
        search: str | None = None,
        include_merged: bool = False,
    ) -> int:
        q = "SELECT COUNT(*) FROM memory_facts"
        params: list = []
        conditions = []
        if not all_users:
            if user_id is None:
                conditions.append("user_id IS NULL")
            else:
                conditions.append(f"user_id = ${len(params) + 1}")
                params.append(user_id)
        if not include_merged:
            conditions.append("merged_into IS NULL")
        if search:
            like = f"%{search}%"
            conditions.append(
                f"(key ILIKE ${len(params) + 1} OR value ILIKE ${len(params) + 2} OR category ILIKE ${len(params) + 3})"
            )
            params.extend([like, like, like])
        if conditions:
            q += " WHERE " + " AND ".join(conditions)
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            return await conn.fetchval(q, *params) or 0


    async def get_categories(self, user_id: str | None = None) -> list[str]:
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            if user_id is None:
                rows = await conn.fetch(
                    "SELECT DISTINCT category FROM memory_facts "
                    "WHERE user_id IS NULL AND merged_into IS NULL ORDER BY category"
                )
            else:
                rows = await conn.fetch(
                    "SELECT DISTINCT category FROM memory_facts "
                    "WHERE user_id = $1 AND merged_into IS NULL ORDER BY category",
                    user_id,
                )
        return [r["category"] for r in rows]


def _row_to_dict(row: asyncpg.Record) -> dict:
    val = row["updated_at"]
    updated_at = val.isoformat() if hasattr(val, "isoformat") else str(val)
    expires_at = row.get("expires_at")
    expires_str = expires_at.isoformat() if hasattr(expires_at, "isoformat") else None
    return {
        "id": row["id"],
        "category": row["category"],
        "key": row["key"],
        "value": row["value"],
        "updated_at": updated_at,
        "source": row.get("source", "conversation"),
        "confidence": row.get("confidence", 70),
        "expires_at": expires_str,
        "source_context": row.get("source_context", ""),
        "merged_into": row.get("merged_into"),
    }