import copy
import re
from collections import deque
from typing import AsyncGenerator, Deque, List, Optional, Set

import structlog

from rasa.builder.copilot.exceptions import (
    CopilotFinalBufferReached,
    CopilotStreamEndedEarly,
)
from rasa.builder.copilot.models import (
    CopilotOutput,
    GeneratedContent,
    ReferenceEntry,
    ReferenceSection,
    ResponseCategory,
    ResponseCompleteness,
)
from rasa.builder.document_retrieval.models import Document

structlogger = structlog.get_logger()

# Controlled prediction markers for detecting roleplay and out-of-scope requests
ROLEPLAY_PREDICTION = "[ROLEPLAY_REQUEST_DETECTED]"
OUT_OF_SCOPE_PREDICTION = "[OUT_OF_SCOPE_REQUEST_DETECTED]"

# Predefined responses for controlled predictions
ROLEPLAY_RESPONSE = """I'm here to help you build and develop your Rasa chatbot, not to
roleplay as the assistant you're creating.

If you'd like to test how your bot would respond, you can:
- Share your bot's configuration files and I can help you understand its behavior
- Ask me to help you modify responses or flows
- Get guidance on how to improve your bot's conversation flow

What would you like to work on with your Rasa chatbot today?"""

OUT_OF_SCOPE_RESPONSE = """I'm specifically designed to help with Rasa chatbot
development and customization.

If you have questions about Rasa development or need help with your chatbot, I'm here to
help! What would you like to work on?"""

# Response for guardrail policy violations
GUARDRAIL_POLICY_VIOLATION_RESPONSE = """I'm here to help you with Rasa chatbot
development in a safe and respectful environment. I cannot assist with requests that
involve harmful content, inappropriate material, or attempts to manipulate my responses.

Please note that continued attempts to use inappropriate content or manipulate
responses will result in project termination and account blocking.
"""

# Redacted message for copilot responses to guardrail policy violations.
COPILOT_REDACTED_MESSAGE = "[Copilot response redacted for guardrails check.]"

# Common LLM response prefixes and suffixes before the actual content. These are removed
# from the content.
LLM_PREFIXES_TO_SUFFIX_REMOVE = {
    "```": "```",
    "```markdown": "```",
    '"""': '"""',
}

# Regex pattern to match inline citations in the generated markdown.
INLINE_CITATION_PATTERN = r"\[([^\]]+)\]\(([^)]+)\)"


class CopilotResponseHandler:
    """Handles the controlled responses from Copilot.

    This handler manages two types of data:
    - llm_stream_buffer: A list of tokens streamed from the LLM during processing.
    - generated_responses: A list of cleaned responses.

    Parameters:
        rolling_buffer_size: Size of the rolling buffer for prefix/suffix handling.
    """

    def __init__(
        self,
        rolling_buffer_size: int = 20,
    ):
        self._rolling_buffer_size = rolling_buffer_size

        # Rolling buffer for handling special tokens and prefix/suffix removal.
        self._rolling_buffer: Deque[str] = deque(maxlen=self._rolling_buffer_size)

        # A list of tokens streamed from the LLM during processing. Tokens are added
        # to the buffer when the rolling buffer is full.
        self._llm_stream_buffer: List[str] = []

        # A list of cleaned and generated responses.
        self._generated_responses: List[GeneratedContent] = []

        # Flag for checking if the stream has ended (raised by the async generator).
        self._stream_ended: bool = False
        # Flag for checking if the stream ended before the max expected special response
        # tokens were reached.
        self._stream_ended_early: bool = False

        # Maximum number of tokens to check for special responses (e.g. roleplay,
        # out-of-scope, etc.).
        self._max_expected_special_response_tokens: int = 20

        # Prefix/suffix tracking
        self._stream_started_with_prefix: bool = False
        self._prefix_found: Optional[str] = None
        self._suffix_found: Optional[str] = None

    @property
    def llm_stream_buffer(self) -> List[str]:
        return copy.deepcopy(self._llm_stream_buffer)

    @property
    def generated_responses(self) -> List[GeneratedContent]:
        return copy.deepcopy(self._generated_responses)

    @property
    def llm_stream_buffer_content(self) -> str:
        return "".join(self._llm_stream_buffer)

    @property
    def token_count(self) -> int:
        return len(self._llm_stream_buffer)

    def clear_buffers(self) -> None:
        """Clear all buffers and reset the handler."""
        self._rolling_buffer.clear()
        self._llm_stream_buffer.clear()
        self._generated_responses.clear()
        self._stream_started_with_prefix = False
        self._prefix_found = None
        self._suffix_found = None
        self._stream_ended = False

    async def _buffer_stream(
        self, response_stream: AsyncGenerator[str, None]
    ) -> AsyncGenerator[str, None]:
        """Wrapper that performs rolling buffer handling and automatically saves chunks.

        Args:
            response_stream: The original streaming response from the LLM.

        Yields:
            The same chunks from the original stream, but persisted in the buffer.

        Raises:
            StopAsyncIteration: If the stream has ended.
        """
        try:
            while True:
                # Get the next chunk from LLM stream and add it to the rolling buffer
                # and the LLM stream buffer.
                chunk = await response_stream.__anext__()
                self._rolling_buffer.append(chunk)

                # Only yield when buffer is full to maintain the rolling buffer
                # behavior
                if len(self._rolling_buffer) == self._rolling_buffer_size:
                    # Yield the oldest element.
                    oldest_element = self._rolling_buffer.popleft()
                    self._llm_stream_buffer.append(oldest_element)
                    yield oldest_element

        except StopAsyncIteration:
            self._llm_stream_buffer.extend(self._rolling_buffer)
            self._stream_ended = True
            # Need to raise a new instance of StopAsyncIteration to avoid the
            # RuntimeError: generator raised StopIteration
            raise CopilotFinalBufferReached()

        except Exception as e:
            structlogger.exception(
                "copilot_response_handler._buffer_stream.unexpected_error",
            )
            raise e

    async def handle_response(
        self, response_stream: AsyncGenerator[str, None]
    ) -> AsyncGenerator[CopilotOutput, None]:
        """Intercept a streaming response and handle special responses from the Copilot.

        Args:
            response_stream: The original streaming response from the LLM.

        Yields:
            ResponseEvent objects representing either generated tokens, default
            responses, or reference sections.
        """
        # Clear the stream buffer and reference buffer at the start
        self.clear_buffers()

        try:
            # Exhaust the buffer early to check for controlled predictions and prefix
            # detection.
            await self._exhaust_buffer_for_early_detection(response_stream)

            # Check for controlled predictions in the collected tokens
            controlled_response = self._check_for_controlled_predictions(
                self.llm_stream_buffer_content,
            )
            if controlled_response is not None:
                self._generated_responses.append(controlled_response)
                yield controlled_response
                return

            # At this point, no controlled predictions were found. Check if the stream
            # started with a prefix and if present, remove it. Yield the clean content.
            initial_content = self._remove_prexif(self.llm_stream_buffer_content)
            if initial_content:
                generated_content = GeneratedContent(
                    content=initial_content,
                    response_category=ResponseCategory.COPILOT,
                    response_completeness=ResponseCompleteness.TOKEN,
                )
                self._generated_responses.append(generated_content)
                yield generated_content

            # Continue streaming remaining chunks with rolling buffer handling
            async for chunk in self._buffer_stream(response_stream):
                generated_content = GeneratedContent(
                    content=chunk,
                    response_category=ResponseCategory.COPILOT,
                    response_completeness=ResponseCompleteness.TOKEN,
                )
                self._generated_responses.append(generated_content)
                yield generated_content

        # Stream ended early
        except CopilotStreamEndedEarly:
            # Check for controlled predictions in the collected tokens
            controlled_response = self._check_for_controlled_predictions(
                self.llm_stream_buffer_content,
            )
            if controlled_response is not None:
                self._generated_responses.append(controlled_response)
                yield controlled_response
                return

            # At this point, no controlled predictions were found. Clean the content
            # from the prefix and suffix if present.
            final_content = self._remove_prefix_and_suffix(
                self.llm_stream_buffer_content,
            )
            if final_content:
                generated_content = GeneratedContent(
                    content=final_content,
                    response_category=ResponseCategory.COPILOT,
                    response_completeness=ResponseCompleteness.COMPLETE,
                )
                self._generated_responses.append(generated_content)
                yield generated_content

        # Stream has ended, process the final rolling buffer. Remove the suffix if
        # present.
        except CopilotFinalBufferReached:
            final_content = self._remove_suffix(self.llm_stream_buffer_content)
            if final_content:
                generated_content = GeneratedContent(
                    content=final_content,
                    response_category=ResponseCategory.COPILOT,
                    response_completeness=ResponseCompleteness.COMPLETE,
                )
                self._generated_responses.append(generated_content)
                yield generated_content

        # Unexpected error occurred.
        except Exception as e:
            structlogger.exception(
                "copilot_response_handler.handle_response.unexpected_error",
            )
            raise e

    async def _exhaust_buffer_for_early_detection(
        self, response_stream: AsyncGenerator[str, None]
    ) -> None:
        """Exhaust the buffer for early detection.

        Args:
            response_stream: The original streaming response from the LLM.
        """
        try:
            async for _ in self._buffer_stream(response_stream):
                if self.token_count >= self._max_expected_special_response_tokens:
                    break
        except CopilotFinalBufferReached:
            # The stream ended before the max expected special response tokens
            # were reached.
            raise CopilotStreamEndedEarly()

    def _check_for_controlled_predictions(
        self, content: str
    ) -> Optional[CopilotOutput]:
        """Check for controlled predictions in the collected tokens.

        Returns:
            Controlled response if found, None otherwise.
        """
        # If the Copilot detected a roleplay request, return the default response.
        if ROLEPLAY_PREDICTION in content:
            structlogger.info("response_interceptor.roleplay_request_detected")
            return GeneratedContent(
                response_category=ResponseCategory.ROLEPLAY_DETECTION,
                content=ROLEPLAY_RESPONSE,
                response_completeness=ResponseCompleteness.COMPLETE,
            )

        # If the Copilot detected an out-of-scope request, return the default response.
        elif OUT_OF_SCOPE_PREDICTION in content:
            structlogger.info("response_interceptor.out_of_scope_request_detected")
            return GeneratedContent(
                response_category=ResponseCategory.OUT_OF_SCOPE_DETECTION,
                content=OUT_OF_SCOPE_RESPONSE,
                response_completeness=ResponseCompleteness.COMPLETE,
            )

        return None

    def _remove_prexif(self, content: str) -> str:
        """Process the initial content from the buffer, handling prefix removal.

        Returns:
            Processed content with prefix removed if applicable.
        """
        # Check if content starts with any of the known prefixes
        for prefix in LLM_PREFIXES_TO_SUFFIX_REMOVE.keys():
            if content.startswith(prefix):
                self._stream_started_with_prefix = True
                self._prefix_found = prefix
                self._code_block_depth = 1
                structlogger.debug(
                    "copilot_response_handler.handle_response.prefix_detected",
                    prefix=prefix,
                    content_length=len(content),
                )
                return content[len(prefix) :]

        return content

    def _remove_suffix(self, content: str) -> str:
        """Process the rolling buffer content, handling suffix removal.

        Returns:
            Processed content with suffix removed if applicable.
        """
        # Check if content ends with any of the known suffixes
        for suffix in LLM_PREFIXES_TO_SUFFIX_REMOVE.values():
            if content.endswith(suffix):
                self._stream_ended_with_suffix = True
                self._suffix_found = suffix
                structlogger.debug(
                    "copilot_response_handler.handle_response.suffix_detected",
                    suffix=suffix,
                    content_length=len(content),
                )
                return content[: -len(suffix)]

        return content

    def _remove_prefix_and_suffix(self, content: str) -> str:
        """Remove the prefix and suffix from the content.

        Returns:
            Processed content with prefix and suffix removed if applicable.
        """
        content = self._remove_prexif(content)
        content = self._remove_suffix(content)
        return content

    @staticmethod
    def respond_to_guardrail_policy_violations() -> GeneratedContent:
        """Respond to guardrail policy violations.

        Returns:
            CopilotOutput object with the response.
        """
        return GeneratedContent(
            response_category=ResponseCategory.GUARDRAILS_POLICY_VIOLATION,
            content=GUARDRAIL_POLICY_VIOLATION_RESPONSE,
            response_completeness=ResponseCompleteness.COMPLETE,
        )

    @staticmethod
    def get_copilot_redacted_message() -> str:
        """Get the redacted message for copilot responses.

        Returns:
            Redacted message for copilot responses.
        """
        return COPILOT_REDACTED_MESSAGE

    def extract_references(self, documents: List[Document]) -> ReferenceSection:
        """Extract references from the LLM stream buffer content.

        This method performs regex matching to find markdown links in the format:
        [text](url).

        The matched links are validated against the provided documents, and a
        ReferenceSection is returned with valid references.

        Note:
            This method can only be called after `handle_response` has been called,
            as it relies on the LLM stream buffer content that gets populated
            during the response handling process.

        Args:
            documents: List of Document objects to match URLs against

        Returns:
            ReferenceSection containing valid reference entries.

        Raises:
            RuntimeError: If called before `handle_response` has populated the buffer.
        """
        if not self.llm_stream_buffer_content:
            message = (
                "`extract_references` can only be called after `handle_response` "
                "has been called to populate the LLM stream buffer."
            )
            structlogger.error(
                "copilot_response_handler.extract_references.buffer_not_populated",
                event_info=message,
                llm_stream_buffer_content=self.llm_stream_buffer_content,
            )
            raise RuntimeError(message)

        # Find all matches in the buffered content
        matches = re.findall(INLINE_CITATION_PATTERN, self.llm_stream_buffer_content)

        valid_references: List[ReferenceEntry] = []
        document_urls: Set[str] = {
            document.url for document in documents if document.url
        }

        for reference_text, reference_url in matches:
            # Check if link_text is a number
            if not reference_text.isdigit():
                structlogger.warning(
                    "copilot_response_handler.extract_references"
                    ".link_text_not_a_number",
                    event_info="Link text is not a number.",
                    link_text=reference_text,
                    link_url=reference_url,
                )

            # Check if URL exists in the documents
            if reference_url not in document_urls:
                structlogger.warning(
                    "copilot_response_handler.extract_references.url_not_found",
                    event_info=(
                        "URL not found in provided documents. Omitted from reference "
                        "section.",
                    ),
                    link_url=reference_url,
                    document_urls=document_urls,
                )
                continue

            # Find the corresponding document for this URL and get its index in the list
            for document_index, document in enumerate(documents, start=1):
                if document.url == reference_url:
                    reference_entry = ReferenceEntry(
                        index=document_index,
                        title=document.title or f"Reference {document_index}",
                        url=reference_url,
                    )
                    valid_references.append(reference_entry)
                    break

        return ReferenceSection(references=valid_references)
