"""
Copyright 2024-2026 ChatterMate

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Splitting a knowledge page into embedder-sized chunks.

The search embedder only reads ~512 tokens, so a page stored as one row is
searchable by its first screen only, and a single hit can be tens of thousands
of chars — enough to overflow the model's context when a few are replayed. Each
page is therefore stored as several small rows that together hold every byte of
it.

Stored chunks are disjoint verbatim slices: concatenating a page's chunks gives
the page back exactly, which is what the dashboard, the FAQ generator and page
edits all rely on. Context across a boundary is handled at embedding time (the
tail of the previous chunk is prepended to the text that is embedded, never to
the text that is stored) and at search time (neighbouring chunks are returned
with a hit).

Why not agno's chunkers: ``ChunkingStrategy.clean_text`` collapses whitespace,
which flattens tables, lists and headings into one line and would persist that
way; ``MarkdownChunking`` needs the ``unstructured`` package.

Chunk ids are ``<page_id>::<n>`` (1-based). The ``_<n>`` form agno uses for PDFs
is kept readable by the page-grouping helpers, but it is not used for new chunks
because real page URLs end in ``_<n>`` too (``/blog/post_2``) and an upsert on
such an id would silently overwrite another page.
"""

import re
from dataclasses import dataclass
from typing import List, Optional

from agno.document import Document
from agno.document.chunking.strategy import ChunkingStrategy
from agno.embedder import Embedder

from app.core.config import settings

CHUNK_ID_SEPARATOR = "::"
# Legacy agno PDF chunk suffix. Stripped when grouping, never generated.
_LEGACY_CHUNK_ID_RE = re.compile(r"_\d+$")
_CHUNK_ID_RE = re.compile(re.escape(CHUNK_ID_SEPARATOR) + r"\d+$")

# Break preference, best first. A chunk ends just after the chosen separator so
# the separator itself is never lost.
_BREAKS = ("\n\n", "\n", ". ", " ")
# Don't take a break that would leave a chunk under this share of the target
# size; a hard cut at full size is better than a 40-char chunk.
_MIN_FILL = 0.3
_TITLE_MAX_CHARS = 120
# Metadata that describes a row's place in its page. Set by chunk_page; must be
# dropped before a page's content is re-chunked.
CHUNK_META_KEYS = frozenset({"page_id", "chunk", "chunk_count", "chunk_size"})


def page_id_of(chunk_id: str) -> str:
    """Map a chunk id back to its page id. Mirrors ``PAGE_ID_EXPR`` in SQL and
    ``basePageId`` in the frontend — keep the three in sync."""
    stripped = _CHUNK_ID_RE.sub("", chunk_id)
    if stripped != chunk_id:
        return stripped
    return _LEGACY_CHUNK_ID_RE.sub("", chunk_id)


def chunk_id(page_id: str, index: int) -> str:
    return f"{page_id}{CHUNK_ID_SEPARATOR}{index}"


def split_text(text: str, size: Optional[int] = None) -> List[str]:
    """Split ``text`` into consecutive pieces of at most ``size`` chars,
    preferring to break on paragraph, line, sentence then word boundaries.
    The pieces are disjoint verbatim slices: ``"".join(pieces) == text``."""
    size = settings.KNOWLEDGE_CHUNK_SIZE if size is None else size
    if size <= 0:
        raise ValueError("chunk size must be positive")

    pieces: List[str] = []
    start = 0
    length = len(text)
    while start < length:
        if length - start <= size:
            pieces.append(text[start:])
            break
        end = _break_point(text, start, size)
        pieces.append(text[start:end])
        start = end
    return pieces


def _break_point(text: str, start: int, size: int) -> int:
    """End of the piece starting at ``start``: just after the best separator in
    the window, else a hard cut at ``size``. Always past ``start``."""
    window = text[start:start + size]
    floor = int(size * _MIN_FILL)
    for sep in _BREAKS:
        idx = window.rfind(sep)
        if idx >= floor:
            return start + idx + len(sep)
    return start + size


@dataclass
class ChunkDocument(Document):
    """A chunk whose embedding is computed from more than its stored text: the
    page title and the tail of the previous chunk are prepended, so a mid-page
    chunk about pricing still matches "pricing" queries that only name the
    page, and a sentence cut by the boundary is embedded whole on one side."""

    embedding_text: Optional[str] = None

    def embed(self, embedder: Optional[Embedder] = None) -> None:
        _embedder = embedder or self.embedder
        if _embedder is None:
            raise ValueError("No embedder provided")
        self.embedding, self.usage = _embedder.get_embedding_and_usage(
            self.embedding_text or self.content
        )


def given_title(document: Document) -> Optional[str]:
    """The title the source recorded for the page, if any."""
    title = (document.meta_data or {}).get("title")
    if isinstance(title, str) and title.strip():
        return title.strip()[:_TITLE_MAX_CHARS]
    return None


def page_title(document: Document) -> str:
    """What to call the page when embedding: its recorded title, else its first
    non-empty line. The fallback is for the embedding only — it is never stored
    as the page's title, since a crawled first line is as likely to be a cookie
    banner as a heading."""
    title = given_title(document)
    if title:
        return title
    for line in (document.content or "").splitlines():
        if line.strip():
            return line.strip()[:_TITLE_MAX_CHARS]
    return document.id or ""


def needs_explicit_chunk_id(page_id: Optional[str]) -> bool:
    """A page id that ends in ``_<n>`` reads as a legacy chunk suffix to the
    grouping rule, which would file ``/blog/post_2`` under ``/blog/post`` and
    let an edit of the sibling delete it. Such pages get a ``::1`` id even
    when they fit in one chunk."""
    return bool(page_id) and bool(_LEGACY_CHUNK_ID_RE.search(page_id))


def chunk_page(
    document: Document,
    size: Optional[int] = None,
    overlap: Optional[int] = None,
) -> List[Document]:
    """Split one page into chunk documents. A page that fits in one chunk is
    returned unchanged (same id, same metadata), so small pages are stored
    exactly as before. ``overlap`` is how much of the previous chunk is
    embedded with each chunk; it never changes what is stored."""
    size = settings.KNOWLEDGE_CHUNK_SIZE if size is None else size
    overlap = settings.KNOWLEDGE_CHUNK_OVERLAP if overlap is None else overlap
    content = document.content or ""
    if len(content) <= size and not needs_explicit_chunk_id(document.id):
        return [document]

    pieces = split_text(content, size)
    title = page_title(document)
    page_id = document.id
    base_meta = {k: v for k, v in (document.meta_data or {}).items() if k not in CHUNK_META_KEYS}
    recorded_title = given_title(document)
    if recorded_title:
        base_meta["title"] = recorded_title
    chunks: List[Document] = []
    previous = ""
    for index, piece in enumerate(pieces, start=1):
        meta = {
            **base_meta,
            "page_id": page_id,
            "chunk": index,
            "chunk_count": len(pieces),
            "chunk_size": len(piece),
        }
        context = previous[-overlap:] if overlap > 0 else ""
        chunks.append(
            ChunkDocument(
                id=chunk_id(page_id, index),
                name=document.name,
                content=piece,
                meta_data=meta,
                embedder=document.embedder,
                embedding_text=f"{title}\n{context}{piece}",
            )
        )
        previous = piece
    return chunks


class PageChunking(ChunkingStrategy):
    """agno ``ChunkingStrategy`` adapter so the PDF readers chunk the same way
    as the website crawler."""

    def __init__(self, size: Optional[int] = None, overlap: Optional[int] = None):
        self.size = size
        self.overlap = overlap

    def chunk(self, document: Document) -> List[Document]:
        return chunk_page(document, self.size, self.overlap)