from pydantic import BaseModel, Field
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union

from evalscope.utils.logger import get_logger

logger = get_logger()

Value = Dict[str, Union[int, float, bool]]


class Score(BaseModel):
    """Score generated by a scorer."""

    value: Value = Field(default_factory=dict)
    """Score value as a dictionary. Key is the score name, value is the score value.
    The first key is considered the main score by default."""

    extracted_prediction: Optional[str] = Field(default=None)
    """Answer extracted from model output (optional)"""

    prediction: Optional[str] = Field(default=None)
    """Original prediction text from the model (optional)"""

    explanation: Optional[str] = Field(default=None)
    """Explanation of score (optional)."""

    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
    """Additional metadata related to the score"""

    main_score_name: Optional[str] = Field(default=None)
    """Main score name, if applicable. This is used to indicate which score is the primary score in a multi-score scenario."""  # noqa: E501

    @property
    def main_value(self) -> Union[int, float, bool]:
        """Main score value."""
        if self.main_score_name and self.main_score_name in self.value:
            return self.value[self.main_score_name]
        elif self.value:
            # If main_score_name is not set or not found, use the first value and update main_score_name
            first_key = next(iter(self.value))
            self.main_score_name = first_key
            return self.value[first_key]
        return None

    @main_value.setter
    def main_value(self, value: Union[int, float, bool]):
        """Set the main score value."""
        if self.main_score_name:
            # If main_score_name is already set, use it
            self.value[self.main_score_name] = value
        elif self.value:
            # If no main_score_name but value dict exists, use the first key
            first_key = next(iter(self.value))
            self.main_score_name = first_key
            self.value[first_key] = value
        else:
            # If neither main_score_name nor value dict exists, initialize both
            self.main_score_name = 'default'
            self.value[self.main_score_name] = value


class SampleScore(BaseModel):
    """Score for a Sample."""

    score: Score
    """A score"""

    sample_id: Optional[Union[str, int]] = Field(default=None)
    """A sample id"""

    group_id: Optional[Union[str, int]] = Field(default=None)
    """A group id for the sample, used for grouping k repeated samples."""

    sample_metadata: Optional[Dict[str, Any]] = Field(default=None)
    """Metadata from the sample"""


class AggScore(BaseModel):
    """Output of an aggregation operation."""

    score: float = Field(default=0.0)
    """Aggregated value as a float."""

    metric_name: str = Field(default='')
    """Name of the metric being aggregated."""

    aggregation_name: str = Field(default='')
    """Name of the aggregation methods"""

    num: int = Field(default=0)
    """Number of samples used in the aggregation."""

    ids: Optional[List[Union[str, int]]] = Field(default=None)
    """List of sample IDs used in the aggregation, if applicable."""

    metadata: Optional[Dict[str, Any]] = Field(default=None)
    """Additional metadata related to the aggregation."""


class Aggregator:

    name = 'default'

    def __call__(self, scores: List[SampleScore]) -> List[AggScore]:
        r"""Aggregate a metric on a list of scores.

        Args:
          scores: List of scores.

        Returns:
          List[AggregatOutput]: List of aggregated outputs.
        """
        ...
