import datetime
import pathlib
from typing import Any, NamedTuple, Optional
import httpx
import requests
import aiofiles
import xxhash
import urllib
from pydantic import BaseModel, RootModel
from talc.knowledge_base import IngestionJobSpec
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
    created_at: Optional[datetime.datetime] = 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"
            and "detail" in response.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:
        """Starts the process of generating a dataset."""
        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(
        self,
        kb: list[Document],
        name: str,
        verify: bool = False,
        auto_link: bool = False,
        root_url: str | None = None,
    ) -> 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)
        ):
            print([doc.title for doc in kb])
            print([doc.filepath for doc in kb])
            files = [
                (
                    "files",
                    (
                        urllib.parse.quote_plus(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,
                "root_url": urllib.parse.quote_plus(root_url) if root_url else None,
            },
            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()

    class UploadBlobResult(NamedTuple):
        blob_hash: bytes
        cached: bool

    async def upload_knowledge_base_blob(
        self, content: pathlib.Path | bytes
    ) -> UploadBlobResult:
        client = httpx.AsyncClient()
        hasher = xxhash.xxh128()
        if isinstance(content, bytes):
            hasher.update(content)
        else:
            buff = bytearray(65536)
            async with aiofiles.open(content, "rb") as f:
                bytes_read = 1
                while bytes_read:
                    bytes_read = await f.readinto(buff)
                    hasher.update(buff[:bytes_read])
        blob_hash = hasher.digest()

        exists_response = await client.get(
            f"{self.url}/kb/blobs/{blob_hash.hex()}/exists",
            headers={"X-TALC-API": self.api_key},
        )

        if exists_response.json()["exists"]:
            return self.UploadBlobResult(blob_hash, True)

        if isinstance(content, bytes):
            upload_response = await client.post(
                f"{self.url}/kb/blobs",
                headers={"X-TALC-API": self.api_key},
                content=content,
            )
        else:
            async with aiofiles.open(content, "rb") as f:
                upload_response = await client.post(
                    f"{self.url}/kb/blobs",
                    headers={"X-TALC-API": self.api_key},
                    content=f,
                )

        actual_blob_hash = bytes.fromhex(upload_response.json()["blob_hash"])
        if actual_blob_hash != blob_hash:
            raise ValueError("Blob hash mismatch")

        return self.UploadBlobResult(blob_hash, False)

    async def ingest_knowledge_base(self, spec: IngestionJobSpec) -> str:
        client = httpx.AsyncClient()
        response = await client.post(
            f"{self.url}/kb/ingest",
            content=spec.model_dump_json(),
            headers={"X-TALC-API": self.api_key},
            timeout=None,
        )
        return response.json()["job_id"]

    async def poll_ingestion_job(self, job_id: str) -> Optional[KnowledgeBase]:
        client = httpx.AsyncClient()
        response = await client.get(
            f"{self.url}/kb/ingestion_jobs/{job_id}",
            headers={"X-TALC-API": self.api_key},
            timeout=None,
        )
        if response.status_code == 404:
            return None

        handle_fastAPI_errors(response)

        return (
            RootModel[Optional[KnowledgeBase]].model_validate_json(response.text).root
        )
