# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

"""Parsing for Meta LlamaGuard safety-classifier responses."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from pyrit.exceptions import InvalidJsonException

if TYPE_CHECKING:
    from collections.abc import Collection

LLAMAGUARD_3_CATEGORY_CODES: tuple[str, ...] = tuple(f"S{index}" for index in range(1, 15))


def parse_llamaguard_response(
    text: str,
    *,
    allowed_categories: Collection[str] = LLAMAGUARD_3_CATEGORY_CODES,
) -> dict[str, Any]:
    """
    Parse a LlamaGuard response for ``CallableResponseHandler``.

    Args:
        text (str): Raw text returned by a LlamaGuard endpoint.
        allowed_categories (Collection[str]): Category codes accepted in an unsafe response.
            Defaults to the Llama Guard 3 S1-S14 taxonomy.

    Returns:
        dict[str, Any]: A true/false score dictionary with rationale and classifier metadata.

    Raises:
        InvalidJsonException: If the response does not match the configured LlamaGuard contract.
    """
    raw = text.strip()
    if not raw:
        raise InvalidJsonException(message="LlamaGuard returned an empty response.")

    lines = raw.splitlines()
    verdict = lines[0].strip().lower()
    if verdict == "safe":
        if len(lines) != 1:
            raise InvalidJsonException(message="A safe LlamaGuard response must contain exactly one line.")
        return _build_safe_response(raw)

    if verdict == "unsafe":
        if len(lines) != 2:
            raise InvalidJsonException(
                message="An unsafe LlamaGuard response must contain a verdict and one category line."
            )
        categories = _parse_unsafe_categories(lines[1], allowed_categories=allowed_categories)
        return _build_unsafe_response(raw=raw, categories=categories)

    raise InvalidJsonException(message=f"LlamaGuard response has an unknown verdict: {lines[0]!r}.")


def _parse_unsafe_categories(category_line: str, *, allowed_categories: Collection[str]) -> list[str]:
    categories = [category.strip() for category in category_line.split(",")]
    if not category_line.strip() or any(not category for category in categories):
        raise InvalidJsonException(message="An unsafe LlamaGuard response must report at least one category.")
    if len(set(categories)) != len(categories):
        raise InvalidJsonException(message="A LlamaGuard response must not report duplicate categories.")

    unknown_categories = sorted(set(categories) - set(allowed_categories))
    if unknown_categories:
        unknown = ", ".join(unknown_categories)
        raise InvalidJsonException(message=f"LlamaGuard reported categories outside the configured policy: {unknown}.")
    return categories


def _build_safe_response(raw: str) -> dict[str, Any]:
    return {
        "score_value": "False",
        "description": "Content classified as safe by LlamaGuard.",
        "rationale": "LlamaGuard returned 'safe'; no configured safety categories were violated.",
        "metadata": {"raw_classifier_output": raw},
    }


def _build_unsafe_response(*, raw: str, categories: list[str]) -> dict[str, Any]:
    category_text = ", ".join(categories)
    return {
        "score_value": "True",
        "description": "Content classified as unsafe by LlamaGuard.",
        "rationale": f"LlamaGuard returned 'unsafe'; violated categories: {category_text}.",
        "metadata": {
            "violated_categories": ",".join(categories),
            "raw_classifier_output": raw,
        },
    }