"""Batch inference — run a prompt across a dataset."""
from __future__ import annotations

import base64
import csv
import io
import json
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from .client import Client


class Batch:
    def __init__(self, client: Client):
        self._client = client

    def create(
        self,
        name: str,
        model: str,
        prompt_template: str,
        *,
        data: list[dict[str, Any]] | None = None,
        csv_file: str | None = None,
        jsonl_file: str | None = None,
        system_prompt: str = "",
        max_tokens: int = 512,
        temperature: float = 0.7,
        concurrency: int = 5,
    ) -> dict:
        """Create a batch inference job.

        Provide exactly one data source:
        - ``data``: list of row dicts (sent as JSONL)
        - ``csv_file``: path to a CSV file
        - ``jsonl_file``: path to a JSONL file

        The ``prompt_template`` uses ``{column_name}`` placeholders that
        are filled from each row.  For example::

            client.batch.create(
                name="ner-clinical",
                model="mistralai/Mistral-7B-Instruct-v0.3",
                prompt_template="Extract named entities from: {text}",
                csv_file="clinical_notes.csv",
            )

        Returns:
            Batch summary dict with ``batch_id``, ``status``, ``total_rows``.
        """
        sources = sum([data is not None, csv_file is not None, jsonl_file is not None])
        if sources != 1:
            raise ValueError("Provide exactly one of: data, csv_file, jsonl_file")

        if data is not None:
            encoded = base64.b64encode(
                "\n".join(json.dumps(row, ensure_ascii=False) for row in data).encode()
            ).decode()
            data_format = "jsonl"
        elif csv_file is not None:
            raw = Path(csv_file).read_bytes()
            encoded = base64.b64encode(raw).decode()
            data_format = "csv"
        else:
            raw = Path(jsonl_file).read_bytes()  # type: ignore[arg-type]
            encoded = base64.b64encode(raw).decode()
            data_format = "jsonl"

        return self._client.post(
            "/v1/research/batch",
            json={
                "name": name,
                "model": model,
                "prompt_template": prompt_template,
                "data": encoded,
                "data_format": data_format,
                "system_prompt": system_prompt,
                "max_tokens": max_tokens,
                "temperature": temperature,
                "concurrency": concurrency,
            },
        )

    def get(self, batch_id: str) -> dict:
        """Get batch status and progress."""
        return self._client.get(f"/v1/research/batch/{batch_id}")

    def list(self) -> list[dict]:
        """List batch jobs for the authenticated user."""
        return self._client.get("/v1/research/batch")

    def results(self, batch_id: str) -> list[dict]:
        """Download batch results as a list of row dicts.

        Each dict has: ``row_index``, ``input``, ``output``, ``error``.
        """
        resp = self._client._session.get(
            f"{self._client.base_url}/v1/research/batch/{batch_id}/results",
            timeout=self._client.timeout,
        )
        if resp.status_code >= 400:
            from .exceptions import YarnError
            raise YarnError(
                f"Failed to fetch results: {resp.status_code}",
                status_code=resp.status_code,
            )
        rows = []
        for line in resp.text.strip().splitlines():
            if line.strip():
                rows.append(json.loads(line))
        return rows

    def cancel(self, batch_id: str) -> None:
        """Cancel a running batch job."""
        self._client.delete(f"/v1/research/batch/{batch_id}")

    def wait(
        self,
        batch_id: str,
        *,
        poll_interval: float = 5.0,
        timeout: float = 3600.0,
    ) -> dict:
        """Poll until a batch completes, fails, or is cancelled.

        Returns:
            Final batch status dict.

        Raises:
            TimeoutError: If the batch doesn't finish within ``timeout`` seconds.
        """
        deadline = time.monotonic() + timeout
        while True:
            status = self.get(batch_id)
            if status["status"] in ("completed", "failed", "cancelled"):
                return status
            if time.monotonic() > deadline:
                raise TimeoutError(
                    f"Batch {batch_id} did not complete within {timeout}s"
                )
            time.sleep(poll_interval)

    def run(
        self,
        name: str,
        model: str,
        prompt_template: str,
        *,
        data: list[dict[str, Any]] | None = None,
        csv_file: str | None = None,
        jsonl_file: str | None = None,
        system_prompt: str = "",
        max_tokens: int = 512,
        temperature: float = 0.7,
        concurrency: int = 5,
        poll_interval: float = 5.0,
        timeout: float = 3600.0,
    ) -> list[dict]:
        """Create a batch, wait for completion, and return results.

        Convenience method that combines ``create`` + ``wait`` + ``results``::

            results = client.batch.run(
                name="classify",
                model="mistralai/Mistral-7B-Instruct-v0.3",
                prompt_template="Classify: {text}",
                csv_file="data.csv",
            )
            for row in results:
                print(row["input"]["text"], "->", row["output"])
        """
        batch = self.create(
            name=name,
            model=model,
            prompt_template=prompt_template,
            data=data,
            csv_file=csv_file,
            jsonl_file=jsonl_file,
            system_prompt=system_prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            concurrency=concurrency,
        )
        self.wait(batch["batch_id"], poll_interval=poll_interval, timeout=timeout)
        return self.results(batch["batch_id"])
