import requests
import json
from pydantic import BaseModel


class TestCase(BaseModel):
    """Represents a single test case that should be run."""

    id: str
    question: str


class TestCaseWithRubric(BaseModel):
    """Represents a single test case that should be run."""

    id: str | None
    question: str
    expected_response: str
    scenario_data: dict[str, object]


class TestCaseResponse(BaseModel):
    """Represents the results of running a single test case."""

    id: str
    response: str


class GradedTestCase(BaseModel):
    """Represents a single test case that has been graded."""

    id: str
    question: str
    response: str
    grade_accuracy: float
    grade_hallucination: float
    grade_reason: str


class RunInfo(BaseModel):
    """Represents a test run."""

    id: str
    completion_progress: float
    grade_accuracy: float | None
    grade_hallucination: float | None
    failed_cases: list[GradedTestCase] | None


class Dataset(BaseModel):
    """Represents a dataset."""

    id: str
    friendly_name: str | None


class EvalsClient:
    """Client for interacting with the Talc evals server."""

    def __init__(self, api_key: str, url: str, run_id: str | None = None) -> None:
        self.url: str = url
        self.run_id: str | None = run_id
        self.api_key: str = api_key

    def start_run(self) -> RunInfo:
        """Create a new test run."""

        response = requests.post(
            f"{self.url}/evals/run/create", headers={"X-TALC-API": self.api_key}
        )
        response.raise_for_status()

        info = RunInfo(**response.json())
        self.run_id = info.id
        return info

    def get_dataset(self, id: str) -> list[TestCase]:
        """Get all test cases in a single dataset from the server."""

        response = requests.get(
            f"{self.url}/evals/dataset/{id}", headers={"X-TALC-API": self.api_key}
        )
        response.raise_for_status()

        return [TestCase(**case) for case in response.json()]

    def submit_responses(self, results: list[TestCaseResponse]) -> None:
        """Submit responses generated by the system for grading."""
        if self.run_id is None:
            raise RuntimeError("No run started. Remember to call start_run() first.")

        data = [x.dict() for x in results if x.id != "id"]

        response = requests.post(
            f"{self.url}/evals/run/{self.run_id}/submit",
            json=data,
            headers={"X-TALC-API": self.api_key},
        )
        response.raise_for_status()

    def get_results(self) -> RunInfo:
        """Get graded results of a run."""
        if self.run_id is None:
            raise RuntimeError("No run started. Remember to call start_run() first.")

        response = requests.get(
            f"{self.url}/evals/run/{self.run_id}/results",
            headers={"X-TALC-API": self.api_key},
        )
        response.raise_for_status()

        return RunInfo(**response.json())

    def create_dataset(
        self, friendly_name: str, dataset: list[TestCaseWithRubric]
    ) -> Dataset:
        """Create a new dataset and save it to the database."""

        data = [x.dict() for x in dataset if x.id != "id"]

        response = requests.post(
            f"{self.url}/evals/dataset/create/{friendly_name}",
            json=data,
            headers={"X-TALC-API": self.api_key},
        )
        response.raise_for_status()

        return Dataset(**response.json())
