import datetime
from typing import Any, Literal, Optional
import requests
from pydantic import BaseModel
from talc.synthetic import (
    Document,
    QuestionGenerationConfig,
)
from talc.grading import GradingSet
from requests.exceptions import HTTPError


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."""

    reason: str | None
    score: float
    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
    last_updated: datetime.datetime
    status: str
    error: str | None
    diagnostics: Any
    name: Optional[str] = None


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

    access_token: str


class KnowledgeBase(BaseModel):
    """A collection of documents, with relationships between them."""

    id: Optional[str] = None
    friendly_name: Optional[str] = None
    documents: list[str] = []  # List of document IDs
    errors: list[str] | None = None

    # FastAPI returns exception details in field .detail, but
    # Python's standard Response does not check this field.


def handle_fastAPI_errors(response):
    if response.status_code > 400:
        if response.headers["content-type"] == "application/json":
            raise HTTPError(response.json()["detail"], response=response)
        else:
            raise HTTPError(response.text, response=response)


class TalcClient:
    """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}/run/create", headers={"X-TALC-API": self.api_key}
        )
        handle_fastAPI_errors(response)

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

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

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

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

    def list_kbs(self) -> list[KnowledgeBase]:
        """Get all datasets from the server."""

        response = requests.get(
            f"{self.url}/kb/list", headers={"X-TALC-API": self.api_key}
        )
        handle_fastAPI_errors(response)

        return [KnowledgeBase(**kb) for kb 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}/run/{self.run_id}/submit",
            json=data,
            headers={"X-TALC-API": self.api_key},
        )
        handle_fastAPI_errors(response)

    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}/run/{self.run_id}/results",
            headers={"X-TALC-API": self.api_key},
        )
        handle_fastAPI_errors(response)

        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}/dataset/create/{friendly_name}",
            json=data,
            headers={"X-TALC-API": self.api_key},
        )
        handle_fastAPI_errors(response)

        return Dataset(**response.json())

    def start_generate_dataset(
        self,
        config: QuestionGenerationConfig,
    ) -> DatasetCreationJob:
        response = requests.post(
            f"{self.url}/dataset/generate/",
            json=config.dict(),
            headers={"X-TALC-API": self.api_key},
        )
        handle_fastAPI_errors(response)

        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}/dataset/generate/{id}",
            headers={"X-TALC-API": self.api_key},
            timeout=10,
        )
        handle_fastAPI_errors(response)

        return DatasetCreationJob(**response.json())

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

        return DatasetCreationJob(**response.json())

    def get_recent_datasets(self) -> list[DatasetCreationJob]:
        """Get the status of the most recent dataset creation jobs."""
        response = requests.get(
            f"{self.url}/dataset/recent",
            headers={"X-TALC-API": self.api_key},
            timeout=10,
        )
        response.raise_for_status()

        return [DatasetCreationJob(**job) for job in response.json()]

    def upload_knowledge_base_asbytes(
        self,
        kb: list[Document],
        name: str,
        verify: bool = False,
        auto_link: bool = False,
    ) -> KnowledgeBase:
        """Upload a knowledge base to the server. The content of each document is uploaded as raw bytes."""
        if (
            isinstance(kb, list)
            and isinstance(kb[0], Document)
            and isinstance(kb[0].content, bytes)
        ):
            files = [
                ("files", (doc.filepath, doc.content, doc.content_type))
                for doc in kb
                if doc.content is not None
            ]
        else:
            raise ValueError(
                "Invalid knowledge base format. Only a list of non-null Documents with byte content supported."
            )

        response = requests.post(
            f"{self.url}/kb/create/{name}",
            params={"verify": verify, "auto_link": auto_link},
            files=files,  # type: ignore
            headers={"X-TALC-API": self.api_key},
        )
        handle_fastAPI_errors(response)
        return KnowledgeBase(**response.json())

    def create_analytics_event(
        self,
        event_name: str,
        event_data: dict[str, Any],
    ) -> None:
        """Create an analytics event."""
        response = requests.post(
            f"{self.url}/analytics",
            json={
                "event_name": event_name,
                "event_data": event_data,
            },
            headers={"X-TALC-API": self.api_key},
        )
        response.raise_for_status()

    # def get_knowledge_base(self, knowledgebase_id: str) -> list[DownloadDocument]:
    #     """Get a knowledge base from the server."""
    #     response = requests.get(
    #         f"{self.url}/kb/{knowledgebase_id}",
    #         headers={"X-TALC-API": self.api_key},
    #     )
    #     handle_fastAPI_errors(response)
