"""Prompt Injection Mitigation sanitization throughput benchmarks.

Measures the performance of:
- PromptSanitizer.sanitize_user_input (clean and dirty text)
- PromptSanitizer.wrap_user_content
- PromptSanitizer.augment_system_prompt
- PromptSanitizer.sanitize_and_wrap
"""

from __future__ import annotations

import sys
from pathlib import Path

try:
    from cleveragents.application.services.prompt_sanitizer import PromptSanitizer
except ModuleNotFoundError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
    from cleveragents.application.services.prompt_sanitizer import PromptSanitizer


# ---------------------------------------------------------------------------
# Clean input suite
# ---------------------------------------------------------------------------


class TimeSanitizeCleanInput:
    """Benchmark sanitization of clean (benign) input."""

    timeout = 30

    def setup(self) -> None:
        """Prepare sanitizer and clean text."""
        self.sanitizer = PromptSanitizer()
        self.clean_text = (
            "Please refactor the authentication module to use OAuth2 " * 100
        )

    def time_sanitize_clean_input(self) -> None:
        """Benchmark sanitize_user_input on clean text."""
        self.sanitizer.sanitize_user_input(self.clean_text)


# ---------------------------------------------------------------------------
# Dirty input suite
# ---------------------------------------------------------------------------


class TimeSanitizeDirtyInput:
    """Benchmark sanitization of dirty (control chars + HTML entities) input."""

    timeout = 30

    def setup(self) -> None:
        """Prepare sanitizer and dirty text."""
        self.sanitizer = PromptSanitizer()
        self.dirty_text = (
            "Hello <world> & 'friends' with \x00\x01\x02 control chars " * 100
        )

    def time_sanitize_dirty_input(self) -> None:
        """Benchmark sanitize_user_input on dirty text."""
        self.sanitizer.sanitize_user_input(self.dirty_text)


# ---------------------------------------------------------------------------
# Wrap content suite
# ---------------------------------------------------------------------------


class TimeWrapContent:
    """Benchmark wrap_user_content throughput."""

    timeout = 30

    def setup(self) -> None:
        """Prepare sanitizer and content."""
        self.sanitizer = PromptSanitizer()
        self.text = "Some user content " * 500

    def time_wrap_user_content(self) -> None:
        """Benchmark wrap_user_content."""
        self.sanitizer.wrap_user_content(self.text)


# ---------------------------------------------------------------------------
# Augment system prompt suite
# ---------------------------------------------------------------------------


class TimeAugmentSystemPrompt:
    """Benchmark augment_system_prompt throughput."""

    timeout = 30

    def setup(self) -> None:
        """Prepare sanitizer and prompt."""
        self.sanitizer = PromptSanitizer()
        self.prompt = "You are a helpful coding assistant. " * 50

    def time_augment_system_prompt(self) -> None:
        """Benchmark augment_system_prompt."""
        self.sanitizer.augment_system_prompt(self.prompt)


# ---------------------------------------------------------------------------
# Sanitize-and-wrap suite
# ---------------------------------------------------------------------------


class TimeSanitizeAndWrap:
    """Benchmark sanitize_and_wrap (combined convenience method)."""

    timeout = 30

    def setup(self) -> None:
        """Prepare sanitizer and text."""
        self.sanitizer = PromptSanitizer()
        self.text = "Refactor authentication to use <OAuth2> & JWT tokens " * 50

    def time_sanitize_and_wrap(self) -> None:
        """Benchmark sanitize_and_wrap."""
        self.sanitizer.sanitize_and_wrap(self.text)