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

from __future__ import annotations

import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Annotated, Any, Literal
from uuid import uuid4

from pydantic import (
    AwareDatetime,
    BaseModel,
    BeforeValidator,
    ConfigDict,
    Field,
    PlainSerializer,
    field_validator,
    model_validator,
)

from pyrit.models.identifiers.component_identifier import ComponentIdentifier

ScoreType = Literal["true_false", "float_scale", "unknown"]


# Annotated alias that round-trips ``ComponentIdentifier`` fields through the flat
# dict storage shape. ``ComponentIdentifier`` is a Pydantic model with a custom flat
# serializer. Defined here (rather than in ``message_piece.py``) because ``score.py``
# sits lower in the import graph; ``message_piece`` imports it from here.
ComponentIdentifierField = Annotated[
    ComponentIdentifier,
    BeforeValidator(lambda v: ComponentIdentifier.model_validate(v) if isinstance(v, dict) else v),
    PlainSerializer(lambda v: v.model_dump() if v is not None else None, return_type=dict | None),
]


class Score(BaseModel):
    """Represents a normalized score generated by a scorer component."""

    model_config = ConfigDict(
        arbitrary_types_allowed=True,
        extra="forbid",
        validate_assignment=False,
    )

    id: uuid.UUID | str = Field(default_factory=uuid4)

    # The value the scorer ended up with; e.g. "true" (if true_false) or "0.5" (if float_scale)
    score_value: str

    # Value that can include a description of the score value
    score_value_description: str | None = None

    # The type of the scorer; e.g. "true_false" or "float_scale"
    score_type: ScoreType

    # The harms categories (e.g. ["hate", "violence"]) – can be multiple
    score_category: list[str] | None = None

    # Extra data the scorer provides around the rationale of the score
    score_rationale: str | None = None

    # Custom metadata a scorer might use. This can vary by scorer.
    score_metadata: dict[str, str | int | float] | None = Field(default_factory=dict)

    # The identifier of the scorer class, including relevant information
    scorer_class_identifier: ComponentIdentifierField | None = None

    # This is the ID of the MessagePiece that the score is scoring. Note a scorer can
    # generate an additional request. This is NOT that, but the ID associated with what
    # we're scoring.
    message_piece_id: uuid.UUID | str

    # Timestamp of when the score was created
    timestamp: AwareDatetime = Field(default_factory=lambda: datetime.now(tz=timezone.utc))

    # The task based on which the text is scored (the original attacker model's objective).
    objective: str | None = None

    # ------------------------------------------------------------------ #
    # Validators
    # ------------------------------------------------------------------ #
    @field_validator("score_metadata", mode="before")
    @classmethod
    def _default_metadata(cls, value: Any) -> Any:
        """
        Coerce ``None`` metadata to an empty dict (preserves legacy ``score_metadata or {}``).

        Returns:
            ``{}`` when ``value`` is ``None``, otherwise ``value`` unchanged.
        """
        return {} if value is None else value

    @model_validator(mode="after")
    def _validate_score_value(self) -> Score:
        """
        Enforce that ``score_value`` is consistent with ``score_type``.

        Returns:
            ``self`` when validation passes.

        Raises:
            ValueError: If the value is incompatible with the score-type constraints.
        """
        self._check_score_value()
        return self

    def _check_score_value(self) -> None:
        """
        Validate ``score_value`` against ``score_type`` constraints.

        Raises:
            ValueError: If the value is incompatible with the score-type constraints.
        """
        if self.score_type == "true_false" and str(self.score_value).lower() not in ("true", "false"):
            raise ValueError(f"True False scorers must have a score value of 'true' or 'false' not {self.score_value}")
        if self.score_type == "float_scale":
            try:
                numeric = float(self.score_value)
            except ValueError as e:
                raise ValueError(f"Float scale scorers require a numeric score value. Got {self.score_value}") from e
            if not (0 <= numeric <= 1):
                raise ValueError(f"Float scale scorers must have a score value between 0 and 1. Got {self.score_value}")

    # ------------------------------------------------------------------ #
    # Public API
    # ------------------------------------------------------------------ #
    def get_value(self) -> bool | float:
        """
        Return the value of the score based on its type.

        If the score type is "true_false", it returns True if the score value is "true" (case-insensitive),
        otherwise it returns False.

        If the score type is "float_scale", it returns the score value as a float.

        Returns:
            bool | float: Parsed score value.

        Raises:
            ValueError: If the score type is unknown.
        """
        if self.score_type == "true_false":
            return self.score_value.lower() == "true"
        if self.score_type == "float_scale":
            return float(self.score_value)

        raise ValueError(f"Unknown scorer type: {self.score_type}")

    def __str__(self) -> str:
        """
        Return a concise text representation of this score.

        Returns:
            str: Human-readable score summary.
        """
        category_str = f": {', '.join(self.score_category) if self.score_category else ''}"
        if self.scorer_class_identifier:
            scorer_type = self.scorer_class_identifier.class_name or "Unknown"
            return f"{scorer_type}{category_str}: {self.score_value}"
        return f"{category_str}: {self.score_value}"

    __repr__ = __str__


@dataclass
class UnvalidatedScore:
    """
    Score is an object that validates all the fields. However, we need a common
    data class that can be used to store the raw score value before it is normalized and validated.
    """

    # The raw score value; has no scale. E.g. in likert could be 1-5
    raw_score_value: str

    score_value_description: str
    score_category: list[str] | None
    score_rationale: str
    score_metadata: dict[str, str | int | float] | None
    scorer_class_identifier: ComponentIdentifier
    message_piece_id: uuid.UUID | str
    objective: str | None
    id: uuid.UUID | str | None = None
    timestamp: datetime | None = None

    def to_score(self, *, score_value: str, score_type: ScoreType) -> Score:
        """
        Convert this unvalidated score into a validated Score.

        Args:
            score_value (str): Normalized score value.
            score_type (ScoreType): Score type.

        Returns:
            Score: Validated score object.

        """
        return Score(
            id=self.id if self.id else uuid4(),
            score_value=score_value,
            score_value_description=self.score_value_description,
            score_type=score_type,
            score_category=self.score_category,
            score_rationale=self.score_rationale,
            score_metadata=self.score_metadata,
            scorer_class_identifier=self.scorer_class_identifier,
            message_piece_id=self.message_piece_id,
            timestamp=self.timestamp if self.timestamp else datetime.now(tz=timezone.utc),
            objective=self.objective,
        )