# Copyright 2024 Klira AI
#
# 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.

"""Policy augmentation — guideline extraction and prompt injection.

Guidelines are extracted during FastRulesEngine evaluation and passed
explicitly via GuardrailProcessingResult.guidelines — NEVER via OTel context.
(Learning #7)

Augmentation verification: augmentation_applied is set ONLY after
confirming guidelines are present in the actual LLM payload.
(Learning #8)
"""

from __future__ import annotations

import logging
from typing import Any

logger = logging.getLogger(__name__)


def build_augmented_messages(
    messages: list[dict[str, Any]],
    guidelines: list[str],
) -> list[dict[str, Any]]:
    """Inject guidelines into LLM messages as a system message.

    Creates a new system message with the guidelines prepended
    to the message list. Does NOT modify the original messages list.

    Args:
        messages: Original LLM messages.
        guidelines: List of guideline strings to inject.

    Returns:
        New messages list with guidelines prepended as system message.
    """
    if not guidelines:
        return messages

    guidelines_text = "\n".join(f"- {g}" for g in guidelines)
    system_message = {
        "role": "system",
        "content": (
            "IMPORTANT COMPLIANCE GUIDELINES:\n"
            f"{guidelines_text}\n\n"
            "You MUST follow these guidelines in your response."
        ),
    }

    # Prepend guidelines system message, after any existing system message
    result: list[dict[str, Any]] = []
    inserted = False
    for msg in messages:
        result.append(msg)
        if not inserted and msg.get("role") == "system":
            result.append(system_message)
            inserted = True

    # If no system message exists, prepend guidelines at the start
    if not inserted:
        result.insert(0, system_message)

    return result


def verify_augmentation(
    original_messages: list[dict[str, Any]],
    augmented_messages: list[dict[str, Any]],
    guidelines: list[str],
) -> bool:
    """Verify that guidelines were actually injected into the messages.

    Returns True only if the augmented messages contain the guideline
    content. This is used to set augmentation_applied accurately.
    (Learning #8: never set augmentation_applied before verification)
    """
    if not guidelines:
        return False

    # Check that augmented messages have more content than original
    if len(augmented_messages) <= len(original_messages):
        return False

    # Check that at least one guideline appears in the augmented messages
    all_content = " ".join(str(msg.get("content", "")) for msg in augmented_messages)
    return any(g in all_content for g in guidelines)
