"""Concrete implementations of guardrails."""

import asyncio
import json
import re
from typing import Any, Callable, Dict, List, Optional, Union

from .base import GuardrailAction, GuardrailResult
from .types import InputGuardrail, OutputGuardrail


class KeywordGuardrail(InputGuardrail):
    """Blocks or modifies messages containing prohibited keywords."""

    def __init__(
        self,
        name: str,
        blocked_keywords: List[str],
        action: GuardrailAction = GuardrailAction.BLOCK,
        replacement: str = "[REDACTED]",
        case_sensitive: bool = False,
        enabled: bool = True,
        blocked_message: Optional[str] = None,
    ):
        """
        Initialize a keyword guardrail.

        Args:
            name: Name of the guardrail
            blocked_keywords: List of keywords to block
            action: Action to take (BLOCK or MODIFY)
            replacement: Text to replace blocked keywords with (if action is MODIFY)
            case_sensitive: Whether keyword matching is case sensitive
            enabled: Whether the guardrail is active
            blocked_message: Custom message when content is blocked
        """
        super().__init__(name, enabled, blocked_message)
        self.blocked_keywords = (
            blocked_keywords
            if case_sensitive
            else [kw.lower() for kw in blocked_keywords]
        )
        self.action = action
        self.replacement = replacement
        self.case_sensitive = case_sensitive

    async def check(self, content: str, context: Dict[str, Any]) -> GuardrailResult:
        """Check for blocked keywords."""
        check_content = content if self.case_sensitive else content.lower()

        for keyword in self.blocked_keywords:
            if keyword in check_content:
                if self.action == GuardrailAction.BLOCK:
                    return GuardrailResult(
                        action=GuardrailAction.BLOCK,
                        reason=f"Blocked keyword found: {keyword}",
                        custom_message=self.blocked_message,
                    )
                elif self.action == GuardrailAction.MODIFY:
                    # Replace the keyword
                    if self.case_sensitive:
                        modified = content.replace(keyword, self.replacement)
                    else:
                        pattern = re.compile(re.escape(keyword), re.IGNORECASE)
                        modified = pattern.sub(self.replacement, content)

                    return GuardrailResult(
                        action=GuardrailAction.MODIFY,
                        content=modified,
                        reason=f"Keyword '{keyword}' replaced",
                    )

        return GuardrailResult(action=GuardrailAction.PASS, content=content)


class LLMGuardrail(OutputGuardrail):
    """Uses a smaller LLM model to verify content safety."""

    def __init__(
        self,
        name: str,
        provider: "LLMProvider",
        safety_prompt: Optional[str] = None,
        enabled: bool = True,
        blocked_message: Optional[str] = None,
    ):
        """
        Initialize an LLM-based guardrail.

        Args:
            name: Name of the guardrail
            provider: LLM provider to use for safety checks
            safety_prompt: Custom prompt for safety checking
            enabled: Whether the guardrail is active
            blocked_message: Custom message when content is blocked
        """
        super().__init__(name, enabled, blocked_message)
        self.provider = provider
        self.safety_prompt = (
            safety_prompt
            or """
        Analyze the following text for safety and appropriateness.
        Respond with JSON: {"safe": true/false, "reason": "explanation if unsafe"}
        
        Text: {content}
        """
        )

    async def check(self, content: str, context: Dict[str, Any]) -> GuardrailResult:
        """Check content safety using LLM."""
        # Create a temporary context for the safety model
        from ..context import ContextManager

        temp_context = ContextManager()
        temp_context.add_message_dict(
            {"role": "user", "content": self.safety_prompt.format(content=content)},
            "safety_check",
        )

        try:
            response = await self.provider.get_response(temp_context)

            # Parse JSON response
            result = json.loads(response)

            if not result.get("safe", True):
                return GuardrailResult(
                    action=GuardrailAction.BLOCK,
                    reason=result.get("reason", "Content deemed unsafe"),
                    custom_message=self.blocked_message,
                )
        except Exception as e:
            self._logger.warning(
                f"Error in LLM safety check: {e}. Assuming content is safe."
            )

        return GuardrailResult(action=GuardrailAction.PASS, content=content)


class RegexGuardrail(InputGuardrail):
    """Applies regex rules to detect and modify patterns."""

    def __init__(
        self,
        name: str,
        patterns: Dict[str, Union[str, GuardrailAction]],
        enabled: bool = True,
        blocked_message: Optional[str] = None,
    ):
        """
        Initialize a regex guardrail.

        Args:
            name: Name of the guardrail
            patterns: Dictionary of {regex_pattern: replacement_string or GuardrailAction}
            enabled: Whether the guardrail is active
            blocked_message: Custom message when content is blocked
        """
        super().__init__(name, enabled, blocked_message)
        self.patterns = patterns

    async def check(self, content: str, context: Dict[str, Any]) -> GuardrailResult:
        """Apply regex patterns to check content."""
        modified_content = content
        modifications_made = False

        for pattern, action_or_replacement in self.patterns.items():
            if re.search(pattern, content):
                if isinstance(action_or_replacement, GuardrailAction):
                    if action_or_replacement == GuardrailAction.BLOCK:
                        return GuardrailResult(
                            action=GuardrailAction.BLOCK,
                            reason=f"Pattern '{pattern}' detected",
                            custom_message=self.blocked_message,
                        )
                else:
                    # It's a replacement string
                    modified_content = re.sub(
                        pattern, action_or_replacement, modified_content
                    )
                    modifications_made = True

        if modifications_made:
            return GuardrailResult(
                action=GuardrailAction.MODIFY,
                content=modified_content,
                reason="Patterns replaced",
            )

        return GuardrailResult(action=GuardrailAction.PASS, content=content)


class CustomFunctionGuardrail(InputGuardrail):
    """Allows using a custom function as a guardrail."""

    def __init__(
        self,
        name: str,
        check_function: Callable[[str, Dict[str, Any]], GuardrailResult],
        enabled: bool = True,
        blocked_message: Optional[str] = None,
    ):
        """
        Initialize a custom function guardrail.

        Args:
            name: Name of the guardrail
            check_function: Function that checks content and returns GuardrailResult
            enabled: Whether the guardrail is active
            blocked_message: Custom message when content is blocked
        """
        super().__init__(name, enabled, blocked_message)
        self.check_function = check_function

    async def check(self, content: str, context: Dict[str, Any]) -> GuardrailResult:
        """Execute the custom check function."""
        if asyncio.iscoroutinefunction(self.check_function):
            result = await self.check_function(content, context)
        else:
            result = await asyncio.to_thread(self.check_function, content, context)

        # Apply custom message if blocking and not already set
        if (
            result.action == GuardrailAction.BLOCK
            and self.blocked_message
            and not result.custom_message
        ):
            result.custom_message = self.blocked_message

        return result
