"""Raw conversation chunks — the append-only record of what compression dropped.

Long-term memory used to have three shapes: curated facts, per-session sketches,
and the live context. When the context was compressed, everything that did not
fit the summary was simply gone — the message rows stay in the database, but
nothing searched them by meaning, and after a reload the session carried only the
summary. "We discussed this earlier" was unanswerable.

This table is the missing store: every message that compression demotes is
written here once, verbatim (a very large single message is elided in the
middle), with the turn range and timestamps it came from and the breadcrumb keys
extracted for it. Three properties are load-bearing:

* **Append-only and immutable.** A row is never updated. A correction is a new
  chunk written later; `UNIQUE(session_id, content_hash)` only prevents the same
  dropped run from being indexed twice, which is what makes a compression replay
  idempotent.
* **Read together with facts and sketches**, through `UnifiedMemorySearch`, so an
  agent has one place to ask. Writing stays separate: chunks are written only by
  the compressor, never by a tool.
* **Provenance is part of the record** — session, turn range, timestamps — because
  an answer built on a raw excerpt must say when it was said, and a curated fact
  from last week must win over a raw line from last month.
"""

from __future__ import annotations

import hashlib
import uuid
from datetime import datetime, timezone
from typing import TYPE_CHECKING

import structlog

from ._embeddings import _vector_to_str
from ._facts import _normalize_query

if TYPE_CHECKING:
    from navi.llm.base import Message

log = structlog.get_logger()

# One chunk is one searchable unit. 4000 chars is roughly a long tool result or
# a dozen chat messages — small enough that a hit is about one topic, large
# enough that the answer's context is not split across three rows.
CHUNK_MAX_CHARS = 4000
# A single message can be megabytes (a full file read). Keep its head and tail:
# the middle of a dump is the part nobody asks for later, while the command line
# at the top and the error at the bottom are usually the breadcrumbs.
CHUNK_MESSAGE_MAX_CHARS = 12_000

# Cosine distance cutoff for a chunk match. Facts use 0.3 because a fact value is
# one short sentence; a 4 KB passage against a short query naturally sits farther
# away, and 0.45 is where the pilot's paraphrased queries still landed. Anything
# past it goes to the lexical fallback instead — a wrong vector hit is worse than
# a keyword hit for a long passage.
VECTOR_CUTOFF_DISTANCE = 0.45

_ROLE_PREFIX = {
    "user": "User",
    "assistant": "Assistant",
    "tool": "Tool result",
}


def build_chunks(
    messages: list[Message],
    breadcrumbs: list[str] | None = None,
    max_chars: int = CHUNK_MAX_CHARS,
) -> list[dict]:
    """Group dropped messages into immutable chunks.

    Consecutive messages are packed together until `max_chars`, so a turn's user
    message, tool calls and the assistant reply land in one chunk. Summary and
    compression marker messages are skipped — a summary is not raw history, and
    indexing old summaries again would make each compression index its own past.
    """
    chunks: list[dict] = []
    current: list[str] = []
    current_len = 0
    start_seq: int | None = None
    end_seq: int | None = None
    started_at: datetime | None = None
    ended_at: datetime | None = None
    keys = list(breadcrumbs or [])

    def flush() -> None:
        nonlocal current, current_len, start_seq, end_seq, started_at, ended_at
        if not current:
            return
        text = "\n\n".join(current).strip()
        if text:
            chunks.append({
                "content": text,
                "content_hash": hashlib.sha256(text.encode("utf-8", "replace")).hexdigest(),
                "start_seq": start_seq if start_seq is not None else -1,
                "end_seq": end_seq if end_seq is not None else -1,
                "started_at": started_at,
                "ended_at": ended_at,
                "breadcrumbs": keys,
            })
        current, current_len = [], 0
        start_seq = end_seq = None
        started_at = ended_at = None

    for msg in messages:
        if msg.role == "system" or msg.is_summary or msg.is_compression:
            continue
        text = _format_message(msg)
        if not text:
            continue
        if current and current_len + len(text) > max_chars:
            flush()
        current.append(text)
        current_len += len(text) + 2
        if start_seq is None:
            start_seq = msg.sequence_number
            started_at = msg.created_at
        if msg.sequence_number >= 0:
            end_seq = msg.sequence_number
        if msg.created_at is not None:
            ended_at = msg.created_at
    flush()
    return chunks


def _format_message(msg: Message) -> str:
    content = (msg.content or "").strip()
    if not content:
        # An assistant message that only requested tool calls still records what
        # was asked for; the result follows in its own message.
        if msg.tool_calls:
            names = ", ".join(tc.name for tc in msg.tool_calls if getattr(tc, "name", None))
            return f"Assistant requested tools: {names}" if names else ""
        return ""
    if len(content) > CHUNK_MESSAGE_MAX_CHARS:
        head = CHUNK_MESSAGE_MAX_CHARS * 3 // 4
        tail = CHUNK_MESSAGE_MAX_CHARS - head - 40
        content = f"{content[:head]}\n…[middle elided]…\n{content[-tail:]}"
    prefix = _ROLE_PREFIX.get(msg.role, msg.role.title())
    if msg.role == "tool" and msg.name:
        prefix = f"Tool result ({msg.name})"
    return f"{prefix}: {content}"


class ChunkMixin:
    """Storage and search for dropped-raw-history chunks.

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

    async def add_chunks(
        self,
        chunks: list[dict],
        user_id: str | None = None,
        session_id: str | None = None,
    ) -> int:
        """Write chunks, skipping ones already indexed for this session."""
        if not chunks or not session_id:
            return 0
        # The embedding column exists only where pgvector does; on a plain Postgres
        # the `::vector` cast in the INSERT would fail the whole write, so the row
        # is written text-only and found through FTS.
        use_vectors = await self._has_pgvector()
        embeddings = (
            await self._generate_embeddings([c["content"] for c in chunks])
            if use_vectors else [None] * len(chunks)
        )
        now = datetime.now(timezone.utc)
        pool = await self._get_pool()
        written = 0
        async with pool.acquire() as conn:
            for chunk, embedding in zip(chunks, embeddings):
                vec_str = _vector_to_str(embedding) if embedding else None
                if use_vectors:
                    result = await conn.execute(
                        """INSERT INTO memory_chunks
                            (id, user_id, session_id, start_seq, end_seq, started_at, ended_at,
                             content, breadcrumbs, content_hash, created_at, embedding)
                           VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12::vector)
                           ON CONFLICT (session_id, content_hash) DO NOTHING""",
                        str(uuid.uuid4()), user_id, session_id,
                        chunk["start_seq"], chunk["end_seq"],
                        chunk.get("started_at"), chunk.get("ended_at"),
                        chunk["content"], list(chunk.get("breadcrumbs") or []),
                        chunk["content_hash"], now, vec_str,
                    )
                else:
                    result = await conn.execute(
                        """INSERT INTO memory_chunks
                            (id, user_id, session_id, start_seq, end_seq, started_at, ended_at,
                             content, breadcrumbs, content_hash, created_at)
                           VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
                           ON CONFLICT (session_id, content_hash) DO NOTHING""",
                        str(uuid.uuid4()), user_id, session_id,
                        chunk["start_seq"], chunk["end_seq"],
                        chunk.get("started_at"), chunk.get("ended_at"),
                        chunk["content"], list(chunk.get("breadcrumbs") or []),
                        chunk["content_hash"], now,
                    )
                if result == "INSERT 0 1":
                    written += 1
        if written:
            log.info(
                "memory.chunks_written",
                session_id=session_id,
                written=written,
                skipped=len(chunks) - written,
            )
        return written

    async def search_chunks(
        self,
        query: str,
        user_id: str | None = None,
        session_id: str | None = None,
        limit: int = 8,
    ) -> list[dict]:
        """Search raw chunks: vector first, lexical fallback when that misses.

        `session_id` narrows to one conversation; omitting it searches every
        session of the user. Anonymous rows (`user_id IS NULL`) are their own
        scope, exactly like facts.
        """
        if limit <= 0:
            return []
        rows: list[dict] = []
        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:
                    rows = await self._vector_chunks(
                        vec_str, user_id, session_id, limit,
                    )
                except Exception:
                    log.warning("memory.chunk_vector_failed", query=query[:40], exc_info=True)
        if rows:
            log.debug("memory.chunk_vector_search", hits=len(rows), query=query[:40])
            return rows
        rows = await self._lexical_chunks(query, user_id, session_id, limit)
        if rows:
            log.debug("memory.chunk_lexical_search", hits=len(rows), query=query[:40])
        return rows

    async def _vector_chunks(
        self,
        vec_str: str,
        user_id: str | None,
        session_id: str | None,
        limit: int,
    ) -> list[dict]:
        where, params = _chunk_filter(user_id, session_id)
        params.append(vec_str)
        vec_idx = len(params)
        params.append(limit * 3)
        lim_idx = len(params)
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            rows = await conn.fetch(
                f"""SELECT id, session_id, start_seq, end_seq, started_at, ended_at,
                           content, breadcrumbs, created_at,
                           embedding <=> ${vec_idx}::vector AS distance
                    FROM memory_chunks
                    WHERE {where}
                    ORDER BY embedding <=> ${vec_idx}::vector
                    LIMIT ${lim_idx}""",
                *params,
            )
        hits: list[dict] = []
        for row in rows:
            distance = row["distance"]
            if distance is None or distance >= VECTOR_CUTOFF_DISTANCE:
                continue
            hit = _row_to_chunk(row)
            hit["score"] = max(0.0, 1.0 - float(distance))
            hits.append(hit)
            if len(hits) >= limit:
                break
        return hits

    async def _lexical_chunks(
        self,
        query: str,
        user_id: str | None,
        session_id: str | None,
        limit: int,
    ) -> list[dict]:
        terms = _normalize_query(query)
        where, params = _chunk_filter(user_id, session_id)
        if not terms:
            return await self._recent_chunks(where, params, limit)
        ts_query = " ".join(terms)
        like_patterns = [f"%{term}%" for term in terms]
        params.append(ts_query)
        ts_idx = len(params)
        params.append(like_patterns)
        like_idx = len(params)
        params.append(limit * 4)
        lim_idx = len(params)
        match_clause = (
            f"(to_tsvector('simple', content) @@ plainto_tsquery('simple', ${ts_idx}) "
            f"OR content ILIKE ANY(${like_idx}::text[]) "
            f"OR array_to_string(breadcrumbs, ' ') ILIKE ANY(${like_idx}::text[]))"
        )
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            rows = await conn.fetch(
                f"""SELECT id, session_id, start_seq, end_seq, started_at, ended_at,
                           content, breadcrumbs, created_at,
                           ts_rank(to_tsvector('simple', content),
                                   plainto_tsquery('simple', ${ts_idx})) AS rank
                    FROM memory_chunks
                    WHERE {where} AND {match_clause}
                    ORDER BY rank DESC, ended_at DESC NULLS LAST
                    LIMIT ${lim_idx}""",
                *params,
            )
        hits = []
        normalized_terms = {t.lower() for t in terms if len(t) > 2}
        for row in rows:
            hit = _row_to_chunk(row)
            haystack = f"{hit['content']} {' '.join(hit['breadcrumbs'])}".lower()
            hit["terms_hit"] = sum(1 for term in normalized_terms if term in haystack)
            hit["score"] = None
            hits.append(hit)
        # A row matched one long term by accident; a row matching several query
        # terms is the one the caller meant. Rank on that before the DB order.
        hits.sort(key=lambda h: (-h["terms_hit"], -_ts(h.get("rank"))))
        return hits[:limit]
    async def _recent_chunks(self, where: str, params: list, limit: int) -> list[dict]:
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            rows = await conn.fetch(
                f"""SELECT id, session_id, start_seq, end_seq, started_at, ended_at,
                           content, breadcrumbs, created_at
                    FROM memory_chunks
                    WHERE {where}
                    ORDER BY ended_at DESC NULLS LAST
                    LIMIT ${len(params) + 1}""",
                *params, limit,
            )
        hits = []
        for index, row in enumerate(rows):
            hit = _row_to_chunk(row)
            hit["score"] = 1.0 - index * 0.05
            hit["terms_hit"] = 0
            hits.append(hit)
        return hits

    async def latest_summary_at(self, session_id: str) -> datetime | None:
        """When this session was last compressed — raw chunks before it are verbatim
        history that the current summary has already curated."""
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            try:
                return await conn.fetchval(
                    "SELECT max(created_at) FROM session_messages "
                    "WHERE session_id = $1 AND is_summary = TRUE",
                    session_id,
                )
            except Exception:
                log.debug("memory.latest_summary_failed", session_id=session_id, exc_info=True)
                return None

    async def chunk_count(self, session_id: str | None = None, user_id: str | None = None) -> int:
        where, params = _chunk_filter(user_id, session_id)
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            return int(await conn.fetchval(
                f"SELECT count(*) FROM memory_chunks WHERE {where}", *params
            ) or 0)


def _chunk_filter(user_id: str | None, session_id: str | None) -> tuple[str, list]:
    """WHERE clause for the (user, session) scope. Anonymous rows are NULL-owned."""
    clauses = []
    params: list = []
    if user_id is None:
        clauses.append("user_id IS NULL")
    else:
        params.append(user_id)
        clauses.append(f"user_id = ${len(params)}")
    if session_id is not None:
        params.append(session_id)
        clauses.append(f"session_id = ${len(params)}")
    return " AND ".join(clauses) or "TRUE", params


def _row_to_chunk(row) -> dict:
    return {
        "rank": row["rank"] if "rank" in row.keys() else None,
        "id": row["id"],
        "session_id": row["session_id"],
        "start_seq": row["start_seq"],
        "end_seq": row["end_seq"],
        "started_at": row["started_at"],
        "ended_at": row["ended_at"],
        "content": row["content"],
        "breadcrumbs": list(row["breadcrumbs"] or []),
        "created_at": row["created_at"],
        "terms_hit": 0,
    }


def _ts(rank) -> float:
    try:
        return float(rank or 0.0)
    except (TypeError, ValueError):
        return 0.0