from elluminate.resources.base import BaseResource
from elluminate.schemas.agent_results import (
    AgentResultsUploadRequest,
    AgentResultsUploadResponse,
    AgentTrialResult,
)


class AgentResultsResource(BaseResource):
    def upload(
        self,
        experiment_id: int,
        results: list[AgentTrialResult],
        task_name_column: str,
        epoch: int = 1,
        evaluate: bool = True,
    ) -> AgentResultsUploadResponse:
        """Upload agent trial results to an experiment."""
        payload = AgentResultsUploadRequest(
            results=results,
            epoch=epoch,
            task_name_column=task_name_column,
            evaluate=evaluate,
        )
        response = self._post(
            f"experiments/{experiment_id}/agent-results",
            json=payload.model_dump(),
        )
        return AgentResultsUploadResponse.model_validate(response.json())

    async def aupload(
        self,
        experiment_id: int,
        results: list[AgentTrialResult],
        task_name_column: str,
        epoch: int = 1,
        evaluate: bool = True,
    ) -> AgentResultsUploadResponse:
        """Upload agent trial results to an experiment (async)."""
        payload = AgentResultsUploadRequest(
            results=results,
            epoch=epoch,
            task_name_column=task_name_column,
            evaluate=evaluate,
        )
        response = await self._apost(
            f"experiments/{experiment_id}/agent-results",
            json=payload.model_dump(),
        )
        return AgentResultsUploadResponse.model_validate(response.json())
