from typing import Any, Literal, Optional
import requests
from pydantic import BaseModel
from talc.grading import GradingSet


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

    id: str
    question: str
    expected_response: str | None


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]
    source_content: list[str]


class Grade(BaseModel):
    """The output of a single grader on a test case."""

    score: float
    reason: str | None
    grader: str


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

    id: str
    question: str
    response: str
    expected_response: str
    grading_complete: bool
    grades: list[Grade]


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

    id: str
    completion_progress: float
    grades: dict[str, float]
    test_cases: list[GradedTestCase] | None


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

    id: str
    friendly_name: str | None


class DatasetCreationJob(BaseModel):
    """Represents the status for a long-running job that creates a dataset in the DB."""

    id: str
    dataset_id: str
    status: str
    error: str | None


class Document(BaseModel):
    """Represents a document used as a ground truth when generating tests."""

    url: Optional[str] = None
    filepath: Optional[str] = None
    title: Optional[str] = None
    content: Optional[str] = None


class APIGenerationArgs(BaseModel):
    """Arguments for creating or getting an API Key."""

    access_token: str


class DatasetCreationArgs(BaseModel):
    """Arguments for creating a dataset."""

    friendly_name: str
    source_docs: list[Document]
    sampling_mode: (
        Literal["main_file"]
        | Literal["all_files"]
        | Literal["random"]
        | Literal["ranked"]
        | None
    )
    generation_modes: list[
        Literal["simple"]
        | Literal["complex"]
        | Literal["sweep"]
        | Literal["hallucination"]
        | str
    ]
    sample_count: int = 10


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, to_grade: GradingSet) -> 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 = to_grade.dict()

        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 upload_dataset(
        self, friendly_name: str, dataset: list[TestCaseWithRubric]
    ) -> Dataset:
        """Upload 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())

    def start_generate_dataset(
        self,
        friendly_name: str,
        source_docs: list[Document],
        sampling_mode: (
            Literal["main_file"]
            | Literal["all_files"]
            | Literal["random"]
            | Literal["ranked"]
        ),
        generation_modes: list[
            Literal["simple"]
            | Literal["complex"]
            | Literal["sweep"]
            | Literal["hallucination"]
            | str
        ],
        sample_count: int = 10,
    ) -> DatasetCreationJob:

        data = DatasetCreationArgs(
            friendly_name=friendly_name,
            source_docs=source_docs,
            sampling_mode=sampling_mode,
            generation_modes=generation_modes,
            sample_count=sample_count,
        ).dict()

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

        return DatasetCreationJob(**response.json())

    def get_dataset_creation_job(self, id: str) -> DatasetCreationJob:
        """Get the status of a dataset creation job."""
        response = requests.get(
            f"{self.url}/evals/dataset/generate/{id}",
            headers={"X-TALC-API": self.api_key},
        )
        response.raise_for_status()

        return DatasetCreationJob(**response.json())
