"""Context generator service for creating context from various sources."""

import re
from typing import List

from rhesis.sdk.utils import count_tokens


class ContextGenerator:
    """Service for generating context from various sources using intelligent semantic chunking."""

    def __init__(
        self,
        max_context_tokens: int = 1500,
    ):
        """
        Initialize the context generator.

        Args:
            max_context_tokens: Maximum tokens per context (user preference)
        """
        self.max_context_tokens = min(max_context_tokens, 3000)

        if max_context_tokens > 3000:
            print(f"⚠️  Context size capped at 3000 tokens (you requested {max_context_tokens})")

    def generate_contexts(self, text: str) -> List[str]:
        """
        Generate contexts using intelligent semantic chunking with hard size limits.

        Strategy:
        1. Identify semantic boundaries (headers, sections, paragraphs)
        2. Create contexts that respect these boundaries
        3. Enforce hard token limit; if a semantic span exceeds the context limit, split it abruptly
        4. If there are no internal boundaries, slice the text into token-capped windows
        """
        if not text:
            raise ValueError("Cannot generate contexts from empty text")

        semantic_boundaries = self._identify_semantic_boundaries(text)

        # If no internal boundaries (just [0, len(text)]), slice linearly
        if len(semantic_boundaries) <= 2:
            contexts: List[str] = []
            start_pos = 0
            while start_pos < len(text):
                end_pos = self._find_token_capped_end(text, start_pos, len(text))
                chunk = text[start_pos:end_pos].strip()
                if chunk:
                    contexts.append(chunk)
                if end_pos <= start_pos:
                    break
                start_pos = end_pos
            return contexts

        # Create contexts from semantic boundaries with abrupt splits when needed
        contexts = self._create_contexts_from_boundaries(text, semantic_boundaries)

        return contexts

    def _identify_semantic_boundaries(self, text: str) -> List[int]:
        """Identify semantic boundaries in the text."""
        boundaries = [0]  # Start of text

        lines = text.split("\n")
        current_pos = 0

        for line in lines:
            line_length = len(line) + 1  # +1 for newline

            # Check for markdown headers
            if re.match(r"^#{1,6}\s+", line):
                boundaries.append(current_pos)

            # Check for section separators
            elif re.match(r"^[-*_]{3,}$", line):
                boundaries.append(current_pos)

            # Check for major paragraph breaks (double newlines)
            elif line.strip() == "" and current_pos > 0:
                # Look ahead to see if this is a paragraph break
                next_non_empty = current_pos + line_length
                while (
                    next_non_empty < len(text)
                    and text[next_non_empty : next_non_empty + 1].isspace()
                ):
                    next_non_empty += 1

                if next_non_empty < len(text) and text[next_non_empty : next_non_empty + 1] not in [
                    "#",
                    "-",
                    "*",
                    "_",
                ]:
                    boundaries.append(current_pos)

            current_pos += line_length

        # Add end of text
        boundaries.append(len(text))

        return boundaries

    def _create_contexts_from_boundaries(self, text: str, boundaries: List[int]) -> List[str]:
        """Create contexts from semantic boundaries."""
        contexts: List[str] = []

        # Start from the first boundary and create contexts sequentially
        start_idx = 0
        while start_idx < len(boundaries) - 1:
            # Find the best end boundary for this context
            end_idx = self._find_best_end_boundary(boundaries, start_idx, text)

            if end_idx <= start_idx:
                break

            # Extract context text
            start_pos = boundaries[start_idx]
            end_pos = boundaries[end_idx]
            context_text = text[start_pos:end_pos].strip()

            if context_text:
                # If the single span between adjacent boundaries exceeds token limit,
                # split it abruptly into token-capped windows
                span_tokens = count_tokens(text[start_pos:end_pos])
                if span_tokens is None:
                    raise ValueError("Failed to count tokens - text may be malformed or invalid")
                if end_idx == start_idx + 1 and span_tokens > self.max_context_tokens:
                    local_start = start_pos
                    while local_start < end_pos:
                        local_end = self._find_token_capped_end(text, local_start, end_pos)
                        piece = text[local_start:local_end].strip()
                        if piece:
                            contexts.append(piece)
                        if local_end <= local_start:
                            break
                        local_start = local_end
                else:
                    contexts.append(context_text)

            # Move to the next boundary
            start_idx = end_idx

        return contexts

    def _find_best_end_boundary(self, boundaries: List[int], start_idx: int, text: str) -> int:
        """Find the best end boundary for a context."""
        start_pos = boundaries[start_idx]

        # Find the furthest boundary within size limit
        for i in range(start_idx + 1, len(boundaries)):
            end_pos = boundaries[i]
            token_len = count_tokens(text[start_pos:end_pos])
            if token_len is None:
                raise ValueError("Failed to count tokens - text may be malformed or invalid")

            if token_len > self.max_context_tokens:
                # We've exceeded the limit, go back one
                if i > start_idx + 1:
                    return i - 1
                else:
                    # Even the smallest context is too big, use it anyway
                    return i

        # Use the last boundary
        return len(boundaries) - 1

    def _find_token_capped_end(self, text: str, start_pos: int, hard_end: int) -> int:
        """Find the furthest end index within hard_end that stays under the token limit."""
        low = start_pos + 1
        high = hard_end
        best = None
        while low <= high:
            mid = (low + high) // 2
            tokens = count_tokens(text[start_pos:mid])
            if tokens is None:
                raise ValueError("Failed to count tokens - text may be malformed or invalid")
            if tokens <= self.max_context_tokens:
                best = mid
                low = mid + 1
            else:
                high = mid - 1
        if best is None:
            return min(start_pos + 1, hard_end)
        return best
