"""
Wrapper for JobData that adds eval_result functionality without modifying autogenerated code.
"""

from __future__ import annotations
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field

from adaptive_sdk.graphql_client import JobData, JobArtifactKind, JobDataDetailsArtifactsByproductsEvaluationByproducts

if TYPE_CHECKING:
    pass


class EvalMetricScore(BaseModel):
    """A single metric score for a model in an evaluation."""

    metric_key: str = Field(description="The key of the metric")
    metric_name: str = Field(description="The human-readable name of the metric")
    mean: float = Field(description="The mean score for this metric")
    min: float = Field(description="The minimum score for this metric")
    max: float = Field(description="The maximum score for this metric")
    stddev: float = Field(description="The standard deviation of scores")
    count: int = Field(description="The number of samples evaluated")
    feedback_count: int = Field(description="The number of feedbacks received")


class ModelEvalResult(BaseModel):
    """Evaluation results for a single model."""

    model_key: str = Field(description="The key of the model that was evaluated")
    model_name: str = Field(description="The human-readable name of the model")
    scores: list[EvalMetricScore] = Field(description="List of metric scores for this model", default_factory=list)


class JobDataPlus(JobData):
    """
    Extended JobData that adds eval_result() method.
    Inherits all attributes and methods from JobData.
    """

    @classmethod
    def from_job_data(cls, job_data: JobData) -> "JobDataPlus":
        """
        Create a JobDataWrapper from a JobData instance.

        Args:
            job_data: The JobData instance to wrap

        Returns:
            A new JobDataWrapper with all the same data
        """
        # Use model_validate to create a new instance with the same data
        return cls.model_validate(job_data.model_dump())

    @property
    def eval_results(self) -> list[ModelEvalResult] | None:
        """
        Extract and organize evaluation results from this job.

        Returns:
            A list of ModelEvalResult objects (one per model), or None if no eval results exist
        """
        if not self.details or not self.details.artifacts:
            return None

        # Dictionary to group scores by model
        model_results: dict[str, ModelEvalResult] = {}

        # Iterate through all artifacts looking for evaluation byproducts
        for artifact in self.details.artifacts:
            if not artifact.byproducts:
                continue

            # Check if this is an EvaluationByproducts with eval_results
            if artifact.kind == JobArtifactKind.EVALUATION:
                if artifact.byproducts is not None:
                    assert isinstance(artifact.byproducts, JobDataDetailsArtifactsByproductsEvaluationByproducts)

                    for eval_result in artifact.byproducts.eval_results:
                        model_key = eval_result.model_service.key
                        model_name = eval_result.model_service.name

                        # Create or get the model result entry
                        if model_key not in model_results:
                            model_results[model_key] = ModelEvalResult(
                                model_key=model_key, model_name=model_name, scores=[]
                            )

                        # Add the metric score to this model's results
                        metric_score = EvalMetricScore(
                            metric_key=eval_result.metric.key,
                            metric_name=eval_result.metric.name,
                            mean=eval_result.mean,
                            min=eval_result.min,
                            max=eval_result.max,
                            stddev=eval_result.stddev,
                            count=eval_result.count,
                            feedback_count=eval_result.feedback_count,
                        )
                        model_results[model_key].scores.append(metric_score)

        if not model_results:
            return None

        return list(model_results.values())
