"""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
