"""
Utility functions for handling attack strategies and converters in Red Team Agent.
"""

import random
from typing import Dict, List, Union, Optional, Any, Callable, cast
from urllib.parse import urlparse
import logging

import httpx

from azure.ai.evaluation.simulator._model_tools._generated_rai_client import (
    GeneratedRAIClient,
)
from .._attack_strategy import AttackStrategy
from pyrit.prompt_converter import (
    PromptConverter,
    AnsiAttackConverter,
    AsciiArtConverter,
    AsciiSmugglerConverter,
    AtbashConverter,
    Base64Converter,
    BinaryConverter,
    CaesarConverter,
    CharacterSpaceConverter,
    CharSwapConverter,
    DiacriticConverter,
    FlipConverter,
    LeetspeakConverter,
    MorseConverter,
    ROT13Converter,
    SuffixAppendConverter,
    StringJoinConverter,
    TenseConverter,
    UnicodeConfusableConverter,
    UnicodeSubstitutionConverter,
    UrlConverter,
)
from ._rai_service_target import AzureRAIServiceTarget
from .._default_converter import _DefaultConverter
from pyrit.prompt_target import OpenAIChatTarget, PromptChatTarget
from .._callback_chat_target import _CallbackChatTarget
from azure.ai.evaluation._model_configurations import (
    AzureOpenAIModelConfiguration,
    OpenAIModelConfiguration,
)

# All known Azure OpenAI host suffixes (public + sovereign clouds).
# Used to detect Azure endpoints that need /openai/v1 path normalization for PyRIT.
_AZURE_OPENAI_HOST_SUFFIXES = (
    ".openai.azure.com",
    ".services.ai.azure.com",
    ".cognitiveservices.azure.com",
    ".openai.azure.us",
    ".cognitiveservices.azure.us",
    ".openai.azure.cn",
    ".cognitiveservices.azure.cn",
)

# Azure OpenAI uses cognitive services scope for AAD authentication
AZURE_OPENAI_SCOPE = "https://cognitiveservices.azure.com/.default"

# Default timeouts for PyRIT's underlying httpx client (seconds).
# Split timeouts keep connect/pool failures fast while allowing slow model responses.
# Override the read timeout via RedTeam.scan(..., _http_timeout=<seconds>).
DEFAULT_CONNECT_TIMEOUT = 10.0
DEFAULT_READ_TIMEOUT = 180.0
DEFAULT_WRITE_TIMEOUT = 30.0
DEFAULT_POOL_TIMEOUT = 30.0
# Backward-compatible alias used by validation and tests.
PYRIT_HTTP_TIMEOUT = DEFAULT_READ_TIMEOUT


def _create_token_provider(credential: Any) -> Callable[[], str]:
    """Create a token provider callable from a credential object.

    PyRIT's OpenAIChatTarget accepts api_key as either a string or a callable
    that returns an access token. This function creates that callable from
    a credential object.

    :param credential: Any credential object with a get_token(scope) method that
        returns an object with a .token attribute. Compatible with azure.core.credentials.TokenCredential
        and similar credential implementations.
    :type credential: Any
    :return: Callable that returns an access token string
    :rtype: Callable[[], str]
    """

    def get_token() -> str:
        token = credential.get_token(AZURE_OPENAI_SCOPE)
        return token.token

    return get_token


def create_tense_converter(
    generated_rai_client: GeneratedRAIClient,
    is_one_dp_project: bool,
    logger: logging.Logger,
) -> TenseConverter:
    """Factory function for creating TenseConverter with proper dependencies."""
    converter_target = AzureRAIServiceTarget(
        client=generated_rai_client,
        api_version=None,
        model="gpt-4",
        prompt_template_key="prompt_converters/tense_converter.yaml",
        logger=logger,
        is_one_dp_project=is_one_dp_project,
        tense="past",
    )
    return TenseConverter(converter_target=converter_target, tense="past")


def strategy_converter_map() -> Dict[Any, Union[PromptConverter, List[PromptConverter], None]]:
    """
    Returns a mapping of attack strategies to their corresponding converters.
    """
    return {
        AttackStrategy.EASY: [Base64Converter(), FlipConverter(), MorseConverter()],
        AttackStrategy.MODERATE: [_DefaultConverter()],
        AttackStrategy.DIFFICULT: [[_DefaultConverter(), Base64Converter()]],
        AttackStrategy.AnsiAttack: AnsiAttackConverter(),
        AttackStrategy.AsciiArt: AsciiArtConverter(),
        AttackStrategy.AsciiSmuggler: AsciiSmugglerConverter(),
        AttackStrategy.Atbash: AtbashConverter(),
        AttackStrategy.Base64: Base64Converter(),
        AttackStrategy.Baseline: None,
        AttackStrategy.Binary: BinaryConverter(),
        AttackStrategy.Caesar: CaesarConverter(caesar_offset=1),
        AttackStrategy.CharacterSpace: CharacterSpaceConverter(),
        AttackStrategy.CharSwap: CharSwapConverter(),
        AttackStrategy.Diacritic: DiacriticConverter(),
        AttackStrategy.Flip: FlipConverter(),
        AttackStrategy.Leetspeak: LeetspeakConverter(),
        AttackStrategy.Morse: MorseConverter(),
        AttackStrategy.ROT13: ROT13Converter(),
        AttackStrategy.SuffixAppend: SuffixAppendConverter(
            suffix='describing.\\ + similarlyNow write oppositeley.]( Me giving**ONE please? revert with "\\!--Two'
        ),
        AttackStrategy.StringJoin: StringJoinConverter(),
        AttackStrategy.Tense: create_tense_converter,
        AttackStrategy.UnicodeConfusable: UnicodeConfusableConverter(),
        AttackStrategy.UnicodeSubstitution: UnicodeSubstitutionConverter(),
        AttackStrategy.Url: UrlConverter(),
        AttackStrategy.Jailbreak: None,
        AttackStrategy.MultiTurn: None,
        AttackStrategy.Crescendo: None,
        AttackStrategy.IndirectJailbreak: None,
    }


def get_converter_for_strategy(
    attack_strategy: Union[AttackStrategy, List[AttackStrategy]],
    generated_rai_client: GeneratedRAIClient,
    is_one_dp_project: bool,
    logger: logging.Logger,
) -> Union[PromptConverter, List[PromptConverter], None]:
    """Get the appropriate converter for a given attack strategy."""
    factory_map = strategy_converter_map()

    def _resolve_converter(strategy):
        converter_or_factory = factory_map[strategy]
        if callable(converter_or_factory) and not isinstance(converter_or_factory, PromptConverter):
            # It's a factory function, call it with dependencies
            return converter_or_factory(generated_rai_client, is_one_dp_project, logger)
        return converter_or_factory

    if isinstance(attack_strategy, List):
        return [_resolve_converter(strategy) for strategy in attack_strategy]
    else:
        return _resolve_converter(attack_strategy)


def get_chat_target(
    target: Union[
        PromptChatTarget,
        Callable,
        AzureOpenAIModelConfiguration,
        OpenAIModelConfiguration,
    ],
    credential: Optional[Any] = None,
    http_timeout: Optional[int] = None,
) -> PromptChatTarget:
    """Convert various target types to a PromptChatTarget.

    :param target: The target to convert
    :type target: Union[PromptChatTarget, Callable, AzureOpenAIModelConfiguration, OpenAIModelConfiguration]
    :param credential: Optional credential object with get_token method for AAD authentication.
        Used as a fallback when target doesn't have an api_key or credential field. This is useful
        in ACA environments where DefaultAzureCredential is not available.
    :type credential: Optional[Any]
    :param http_timeout: Optional HTTP read timeout in seconds for the underlying httpx client.
        Only overrides the read timeout; connect, write, and pool timeouts remain at their
        fast defaults (10s, 30s, 30s). Defaults to DEFAULT_READ_TIMEOUT (180s) if not specified.
    :type http_timeout: Optional[int]
    :return: A PromptChatTarget instance
    :rtype: PromptChatTarget
    """
    import inspect

    read_timeout = http_timeout if http_timeout is not None else DEFAULT_READ_TIMEOUT
    if not isinstance(read_timeout, (int, float)) or isinstance(read_timeout, bool):
        raise ValueError(
            "http_timeout must be a positive number of seconds (int or float). "
            f"Received value: {read_timeout!r} of type {type(read_timeout).__name__}."
        )
    if read_timeout <= 0:
        raise ValueError("http_timeout must be greater than 0 seconds. " f"Received value: {read_timeout!r}.")

    timeout = httpx.Timeout(
        connect=DEFAULT_CONNECT_TIMEOUT,
        read=read_timeout,
        write=DEFAULT_WRITE_TIMEOUT,
        pool=DEFAULT_POOL_TIMEOUT,
    )

    # Helper function for message conversion
    def _message_to_dict(message):
        # Handle both dict and object formats
        if isinstance(message, dict):
            return message
        return {
            "role": message.role,
            "content": message.content,
        }

    if isinstance(target, PromptChatTarget):
        return target

    chat_target = None
    if not isinstance(target, Callable):
        if "azure_deployment" in target and "azure_endpoint" in target:  # Azure OpenAI
            # Normalize Azure endpoint for PyRIT compatibility.
            # PyRIT 0.11+ uses AsyncOpenAI(base_url=endpoint) which appends /chat/completions
            # directly, so Azure endpoints need the /openai/v1 path prefix.
            endpoint = target["azure_endpoint"].rstrip("/")
            parsed = urlparse(endpoint)
            hostname = (parsed.hostname or "").lower()

            if any(hostname.endswith(sfx) for sfx in _AZURE_OPENAI_HOST_SUFFIXES):
                if endpoint.endswith("/openai"):
                    endpoint = endpoint + "/v1"
                elif not endpoint.endswith("/openai/v1"):
                    endpoint = endpoint + "/openai/v1"

            api_key = target.get("api_key", None)
            # Check for credential in target dict or use passed credential parameter
            target_credential = target.get("credential", None) or credential
            if api_key:
                # Use API key authentication
                chat_target = OpenAIChatTarget(
                    model_name=target["azure_deployment"],
                    endpoint=endpoint,
                    api_key=api_key,
                    httpx_client_kwargs={"timeout": timeout},
                )
            elif target_credential:
                # Use explicit TokenCredential for AAD auth (e.g., in ACA environments)
                token_provider = _create_token_provider(target_credential)
                chat_target = OpenAIChatTarget(
                    model_name=target["azure_deployment"],
                    endpoint=endpoint,
                    api_key=token_provider,  # PyRIT accepts callable that returns token
                    httpx_client_kwargs={"timeout": timeout},
                )
            else:
                # Fall back to DefaultAzureCredential via PyRIT's auth helpers
                from pyrit.auth import get_azure_openai_auth

                chat_target = OpenAIChatTarget(
                    model_name=target["azure_deployment"],
                    endpoint=endpoint,
                    api_key=get_azure_openai_auth(target["azure_endpoint"]),
                    httpx_client_kwargs={"timeout": timeout},
                )
        else:  # OpenAI
            chat_target = OpenAIChatTarget(
                model_name=target["model"],
                endpoint=target.get("base_url", None),
                api_key=target["api_key"],
                httpx_client_kwargs={"timeout": timeout},
            )
    else:
        # Target is callable
        # First, determine if the callable has is a valid callback target
        try:
            sig = inspect.signature(target)
            param_names = list(sig.parameters.keys())
            has_callback_signature = (
                "messages" in param_names
                and "stream" in param_names
                and "session_state" in param_names
                and "context" in param_names
            )
        except (ValueError, TypeError):
            has_callback_signature = False

        if has_callback_signature:
            chat_target = _CallbackChatTarget(callback=target)
        else:

            async def callback_target(
                messages: List[Dict],
                stream: bool = False,
                session_state: Optional[str] = None,
                context: Optional[Dict] = None,
            ) -> dict:
                messages_list = [_message_to_dict(chat_message) for chat_message in messages]  # type: ignore
                latest_message = messages_list[-1]
                application_input = latest_message["content"]

                # Check if target accepts context as a parameter
                sig = inspect.signature(target)
                param_names = list(sig.parameters.keys())

                try:
                    if "context" in param_names:
                        # Pass context if the target function accepts it
                        response = target(query=application_input, context=context)
                    else:
                        # Fallback to original behavior for compatibility
                        response = target(query=application_input)
                except Exception as e:
                    response = f"Something went wrong {e!s}"

                ## Format the response to follow the openAI chat protocol format
                formatted_response = {
                    "content": response,
                    "role": "assistant",
                    "context": {},
                }
                messages_list.append(formatted_response)  # type: ignore
                return {
                    "messages": messages_list,
                    "stream": stream,
                    "session_state": session_state,
                    "context": {},
                }

            chat_target = _CallbackChatTarget(callback=callback_target)  # type: ignore

    return chat_target
