import functools
import os
from collections.abc import Callable
from typing import Any, ClassVar, TypeVar

try:
    from azure.ai.contentsafety import BlocklistClient, ContentSafetyClient
    from azure.ai.contentsafety.models import (
        AddOrUpdateTextBlocklistItemsOptions,
        AnalyzeImageOptions,
        AnalyzeTextOptions,
        ImageData,
        RemoveTextBlocklistItemsOptions,
        TextBlocklist,
        TextBlocklistItem,
        TextCategory,
    )
    from azure.core.credentials import AzureKeyCredential
    from azure.core.exceptions import HttpResponseError
except ImportError as e:
    msg = (
        "azure-ai-contentsafety package is not installed. "
        "Please install it with `pip install 'any-guardrail[azure-content-safety]'` to use AzureContentSafety guardrail."
    )
    raise ImportError(msg) from e

from any_guardrail.base import GuardrailOutput, ThreeStageGuardrail
from any_guardrail.types import GuardrailInferenceOutput, GuardrailPreprocessOutput

# Type aliases for Azure Content Safety
AzureAnalyzeInput = AnalyzeTextOptions | AnalyzeImageOptions
AzureAnalyzeOutput = Any  # Azure API response type

# TypeVar for preserving function signatures in decorator
F = TypeVar("F", bound=Callable[..., Any])


def error_message(message: str) -> Callable[[F], F]:
    """Handle exceptions for Azure Content Safety operations."""

    def error_handler_decorator(func: F) -> F:
        """Handle exceptions for the wrapped function."""

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            try:
                return func(*args, **kwargs)
            except HttpResponseError as e:
                raise RuntimeError(message + f" Details: {e!s}") from e

        return wrapper  # type: ignore[return-value]

    return error_handler_decorator


class AzureContentSafety(
    ThreeStageGuardrail[AzureAnalyzeInput, AzureAnalyzeOutput, bool, dict[str, int | list[str] | None], float]
):
    """Guardrail implementation using Azure Content Safety service.

    Azure Content Safety provides content moderation capabilities for text and images. To learn more about Azure
    Content Safety, visit the [official documentation](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/contentsafety/azure-ai-contentsafety`).

    """

    SUPPORTED_MODELS: ClassVar = ["azure-content-safety"]

    def __init__(
        self,
        endpoint: str | None = None,
        api_key: str | None = None,
        threshold: int = 2,
        score_type: str = "max",
        blocklist_names: list[str] | None = None,
    ) -> None:
        """Initialize Azure Content Safety client.

        Args:
            endpoint (str): The endpoint URL for the Azure Content Safety service.
            api_key (str): The API key for authenticating with the service.
            threshold (int): The threshold for determining if content is unsafe.
            score_type (str): The type of score to use ("max" or "avg").
            blocklist_names (List[str] | None): List of blocklist names to use for content evaluation.

        """
        if api_key:
            credential = AzureKeyCredential(api_key)
        else:
            try:
                credential = AzureKeyCredential(os.environ["CONTENT_SAFETY_KEY"])
            except KeyError as e:
                msg = "CONTENT_SAFETY_KEY environment variable is not set. Either provide an api_key or set the environment variable."
                raise KeyError(msg) from e
        if not endpoint:
            try:
                endpoint = os.environ["CONTENT_SAFETY_ENDPOINT"]
            except KeyError as e:
                msg = "CONTENT_SAFETY_ENDPOINT environment variable is not set. Either provide an endpoint or set the environment variable."
                raise KeyError(msg) from e

        self.client = ContentSafetyClient(endpoint=endpoint, credential=credential)
        self.blocklist_client = BlocklistClient(endpoint=endpoint, credential=credential)
        self.threshold = threshold

        if score_type not in ["max", "avg"]:
            msg = "score_type must be either 'max' or 'avg'"
            raise ValueError(msg)
        self.score_type = score_type

        if blocklist_names:
            if not isinstance(blocklist_names, list):
                msg = "blocklist_names must be a list of strings"
                raise ValueError(msg)
            for name in blocklist_names:
                if not isinstance(name, str):
                    msg = "blocklist_names must be a list of strings"
                    raise ValueError(msg)
        self.blocklist_names = blocklist_names

    def validate(self, content: str) -> GuardrailOutput[bool, dict[str, int | list[str] | None], float]:
        """Validate content using Azure Content Safety.

        Args:
            content (str): The content to be evaluated.

        Returns:
            GuardrailOutput: The result of the guardrail evaluation.

        """
        model_inputs = self._pre_processing(content)
        model_outputs = self._inference(model_inputs)
        return self._post_processing(model_outputs)

    @error_message("Was unable to create or update blocklist.")
    def create_or_update_blocklist(
        self, blocklist_name: str, blocklist_description: str, add_to_blocklist_names: bool = True
    ) -> None:
        """Create or update a blocklist in Azure Content Safety.

        Args:
            blocklist_name (str): The name of the blocklist.
            blocklist_description (str): The description of the blocklist.
            add_to_blocklist_names (bool): Whether to add the blocklist name to the guardrail's blocklist_names.

        """
        self.blocklist_client.create_or_update_text_blocklist(
            blocklist_name=blocklist_name,
            options=TextBlocklist(blocklist_name=blocklist_name, description=blocklist_description),
        )
        if add_to_blocklist_names:
            self.blocklist_names.append(blocklist_name) if self.blocklist_names else setattr(
                self, "blocklist_names", [blocklist_name]
            )

    @error_message("Was unable to add blocklist items.")
    def add_blocklist_items(self, blocklist_name: str, blocklist_terms: list[str]) -> None:
        """Add items to a blocklist.

        Args:
            blocklist_name (str): The name of the blocklist.
            blocklist_terms (List[str]): The terms to add to the blocklist.

        """
        blocklist_items = []
        for term in blocklist_terms:
            blocklist_items.append(TextBlocklistItem(text=term))

        self.blocklist_client.add_or_update_blocklist_items(
            blocklist_name=blocklist_name,
            options=AddOrUpdateTextBlocklistItemsOptions(blocklist_items=blocklist_items),
        )

    @error_message("Was unable to list blocklists.")
    def list_blocklists(self) -> list[dict[str, str | None]]:
        """List all blocklists in Azure Content Safety.

        Returns:
            List[Dict[str, str]]: A list of blocklist details.

        """
        blocklists = self.blocklist_client.list_text_blocklists()
        return [{"name": blocklist.blocklist_name, "description": blocklist.description} for blocklist in blocklists]

    @error_message("Was unable to list blocklist items.")
    def list_blocklist_items(self, blocklist_name: str) -> list[dict[str, str | None]]:
        """List items in a blocklist.

        Args:
            blocklist_name (str): The name of the blocklist.

        Returns:
            List[Dict[str, str]]: The list of blocklist items.

        """
        blocklist_items = self.blocklist_client.list_text_blocklist_items(blocklist_name=blocklist_name)
        return [
            {"id": item.blocklist_item_id, "text": item.text, "description": item.description}
            for item in blocklist_items
        ]

    @error_message("Was unable to get blocklist.")
    def get_blocklist(self, blocklist_name: str) -> dict[str, str | None]:
        """Get a blocklist by name.

        Args:
            blocklist_name (str): The name of the blocklist.

        Returns:
            dict[str, str]: The blocklist details.

        """
        blocklist = self.blocklist_client.get_text_blocklist(blocklist_name=blocklist_name)
        return {"name": blocklist.blocklist_name, "description": blocklist.description}

    @error_message("Was unable to get blocklist item.")
    def get_blocklist_item(self, blocklist_name: str, item_id: str) -> dict[str, str | None]:
        """Get a blocklist item by ID.

        Args:
            blocklist_name (str): The name of the blocklist.
            item_id (str): The ID of the blocklist item.

        Returns:
            dict[str, str]: The blocklist item details.

        """
        item = self.blocklist_client.get_text_blocklist_item(blocklist_name=blocklist_name, blocklist_item_id=item_id)
        return {"id": item.blocklist_item_id, "text": item.text, "description": item.description}

    @error_message("Was unable to delete blocklist.")
    def delete_blocklist(self, blocklist_name: str) -> None:
        """Delete a blocklist by name.

        Args:
            blocklist_name (str): The name of the blocklist.

        """
        self.blocklist_client.delete_text_blocklist(blocklist_name=blocklist_name)
        self.blocklist_names.remove(
            blocklist_name
        ) if self.blocklist_names and blocklist_name in self.blocklist_names else None

    @error_message("Was unable to delete blocklist item.")
    def delete_blocklist_items(self, blocklist_name: str, item_ids: list[str]) -> None:
        """Delete a blocklist item by ID.

        Args:
            blocklist_name (str): The name of the blocklist.
            item_ids (List[str]): The IDs of the blocklist items.

        """
        self.blocklist_client.remove_blocklist_items(  # type: ignore [call-overload]
            blocklist_name=blocklist_name,
            blocklist_item_id=RemoveTextBlocklistItemsOptions(blocklist_item_ids=item_ids),
        )

    def _pre_processing(self, text: str) -> GuardrailPreprocessOutput[AzureAnalyzeInput]:
        if self._is_existing_path(text):
            try:
                with open(text, "rb") as file:
                    options: AzureAnalyzeInput = AnalyzeImageOptions(image=ImageData(content=file.read()))
            except ValueError as e:
                msg = "Must provide a file path to an image file."
                raise ValueError(msg) from e
        else:
            if self.blocklist_names:
                options = AnalyzeTextOptions(
                    text=text, blocklist_names=self.blocklist_names, halt_on_blocklist_hit=False
                )
            else:
                options = AnalyzeTextOptions(text=text)
        return GuardrailPreprocessOutput(data=options)

    @error_message("Was unable to analyze text or image.")
    def _inference(
        self, model_inputs: GuardrailPreprocessOutput[AzureAnalyzeInput]
    ) -> GuardrailInferenceOutput[AzureAnalyzeOutput]:
        if isinstance(model_inputs.data, AnalyzeTextOptions):
            response = self.client.analyze_text(model_inputs.data)
        else:
            response = self.client.analyze_image(model_inputs.data)  # type: ignore [assignment]
        return GuardrailInferenceOutput(data=response)

    def _post_processing(
        self, model_outputs: GuardrailInferenceOutput[AzureAnalyzeOutput]
    ) -> GuardrailOutput[bool, dict[str, int | list[str] | None], float]:
        results_dict = {
            "hate": next(item for item in model_outputs.data.categories_analysis if item.category == TextCategory.HATE),
            "self_harm": next(
                item for item in model_outputs.data.categories_analysis if item.category == TextCategory.SELF_HARM
            ),
            "sexual": next(
                item for item in model_outputs.data.categories_analysis if item.category == TextCategory.SEXUAL
            ),
            "violence": next(
                item for item in model_outputs.data.categories_analysis if item.category == TextCategory.VIOLENCE
            ),
        }

        explanation = {key: result.severity for key, result in results_dict.items() if result is not None}

        if self.score_type == "max":
            score = max(
                explanation_score for explanation_score in explanation.values() if explanation_score is not None
            )
        else:
            score = sum(
                explanation_score for explanation_score in explanation.values() if explanation_score is not None
            ) / sum(1 for explanation_score in explanation.values() if explanation_score is not None)

        explanation["blocklist"] = model_outputs.data.blocklists_match if self.blocklist_names else None

        valid = score < self.threshold
        if valid and explanation.get("blocklist"):
            valid = False

        return GuardrailOutput(valid=valid, explanation=explanation, score=score)

    def _is_existing_path(self, text: str) -> bool:
        return os.path.exists(text)
