# Copyright (C) 2026 GoodData Corporation

"""Text chunking for rules retrieval.

We chunk by markdown heading sections (human navigation) rather than arbitrary character windows.
This improves both embedding retrieval and lexical matching.
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field

from aida_mcp.tools.rules.store import Rule
from aida_mcp.tools.rules.vocab import HarvestedTokens, harvest_section_tokens

logger = logging.getLogger(__name__)

# Rules smaller than this are not considered "large" for output formatting.
MIN_CHUNK_THRESHOLD = 1500


@dataclass
class RuleChunk:
    """A chunk of a rule document."""

    rule_path: str
    chunk_index: int
    total_chunks: int
    text: str
    is_description: bool = False  # True if this is the description-only chunk
    heading_path: str | None = None  # Populated for markdown section chunks

    # Tokens for lexical scoring (BM25). Computed during chunking.
    tokens: set[str] = field(default_factory=lambda: set[str]())
    boost_tokens: set[str] = field(default_factory=lambda: set[str]())

    @property
    def is_whole_rule(self) -> bool:
        """Whether this chunk represents the entire rule (not chunked)."""
        return self.total_chunks == 1 and not self.is_description


class RuleChunker:
    """
    Chunks rules into markdown heading sections for retrieval.

    We keep an explicit description chunk (if present) and then produce section chunks.
    """

    _HEADING_RE = re.compile(r"^(#{1,6})\s+(.+?)\s*$")

    def __init__(self, *, min_chunk_threshold: int = MIN_CHUNK_THRESHOLD):
        # This threshold is used by formatting to decide "large rule" behavior; chunking is always section-based.
        self.min_chunk_threshold = min_chunk_threshold

    def _split_into_sections(self, content: str) -> list[tuple[str | None, str]]:
        """Split markdown content into heading-based sections.

        Returns list of (heading_path, section_text). If no headings are present, returns one section.
        """
        lines = content.splitlines()
        headings: list[tuple[int, str, int]] = []  # (level, title, line_idx)
        for i, line in enumerate(lines):
            m = self._HEADING_RE.match(line)
            if not m:
                continue
            level = len(m.group(1))
            title = m.group(2).strip()
            headings.append((level, title, i))

        if not headings:
            return [(None, content)]

        # Walk headings and create sections; maintain a stack for heading_path.
        sections: list[tuple[str | None, str]] = []
        stack: list[tuple[int, str]] = []

        def build_path() -> str:
            return " > ".join(t for _, t in stack)

        # Preserve any preamble content before the first heading (otherwise it would be lost).
        first_heading_line = headings[0][2]
        if first_heading_line > 0:
            preamble = "\n".join(lines[:first_heading_line]).strip()
            if preamble:
                sections.append((None, preamble))

        for idx, (level, title, start_line) in enumerate(headings):
            # Determine end line: next heading start
            end_line = headings[idx + 1][2] if idx + 1 < len(headings) else len(lines)

            # Update stack for this heading
            while stack and stack[-1][0] >= level:
                stack.pop()
            stack.append((level, title))

            section_text = "\n".join(lines[start_line:end_line]).strip()
            if section_text:
                sections.append((build_path(), section_text))

        return sections

    def chunk_rule(self, rule: Rule) -> list[RuleChunk]:
        """
        Convert a rule into chunks for embedding.

        Args:
            rule: The rule to chunk

        Returns:
            List of RuleChunk objects
        """
        chunks: list[RuleChunk] = []

        # Always create a description chunk for semantic matching on "what is this about"
        if rule.description:
            desc_tokens = harvest_section_tokens(heading_path=None, text=rule.description)
            chunks.append(
                RuleChunk(
                    rule_path=rule.path,
                    chunk_index=0,
                    total_chunks=1,  # Will be updated below
                    text=rule.description,
                    is_description=True,
                    heading_path=None,
                    tokens=desc_tokens.tokens,
                    boost_tokens=desc_tokens.boost_tokens,
                )
            )

        content = rule.content
        sections = self._split_into_sections(content)
        logger.debug(f"Split {rule.path} into {len(sections)} markdown section(s)")

        for heading_path, section_text in sections:
            harvested: HarvestedTokens = harvest_section_tokens(heading_path=heading_path, text=section_text)
            chunks.append(
                RuleChunk(
                    rule_path=rule.path,
                    chunk_index=len(chunks),
                    total_chunks=1,  # Will be updated below
                    text=section_text,
                    is_description=False,
                    heading_path=heading_path,
                    tokens=harvested.tokens,
                    boost_tokens=harvested.boost_tokens,
                )
            )

        # Update total_chunks for all chunks
        total = len(chunks)
        for chunk in chunks:
            chunk.total_chunks = total

        return chunks

    def chunk_rules(self, rules: list[Rule]) -> list[RuleChunk]:
        """
        Convert multiple rules into chunks.

        Args:
            rules: List of rules to chunk

        Returns:
            List of all RuleChunk objects
        """
        all_chunks: list[RuleChunk] = []
        for rule in rules:
            all_chunks.extend(self.chunk_rule(rule))

        logger.info(f"Chunked {len(rules)} rules into {len(all_chunks)} chunks")
        return all_chunks
