"""
AgniPod SDK — Batch Processing
=================================

Usage::

    from agnipod import AgniPod

    client = AgniPod()

    # Create a JSON batch (up to 10,000 items, each can be chat or generate)
    result = client.batches.create(items=[
        {
            "custom_id": "req-001",
            "model": "qwen3:8b",
            "messages": [{"role": "user", "content": "What is 2+2?"}],
        },
        {
            "custom_id": "req-002",
            "model": "qwen3:8b",
            "prompt": "Translate to French: Hello",
        },
    ])
    print(result.batch_id)

    # Upload a JSONL file for large batches
    result = client.batches.create_jsonl(file_path="requests.jsonl")

    # Poll until complete
    batch = client.batches.wait(result.batch_id, on_progress=lambda p: print(p.status))

    # Download results (streamed to disk for large files)
    client.batches.results(batch.id, save_to="output.jsonl")

    # List all batches
    for batch in client.batches.list():
        print(batch.id, batch.status)
"""

from __future__ import annotations
import os
from typing import Any, Callable, Dict, List, Optional

from ._client import HttpClient
from ._types import Batch, BatchCreateResult, BatchList, BatchProgress, _APIObject
from ._exceptions import ValidationError, APITimeoutError

_PREFIX = "/batch"


class Batches:
    """Namespace for ``/batch`` operations."""

    def __init__(self, client: HttpClient):
        self._client = client

    def create(
        self,
        *,
        items: List[Dict[str, Any]],
        batch_name: Optional[str] = None,
    ) -> BatchCreateResult:
        """Create a synchronous JSON batch (up to 10,000 items).

        Each item should have ``custom_id``, ``model``, and either
        ``messages`` (chat) or ``prompt`` (generate).

        Parameters
        ----------
        items : list[dict]
            List of batch items.
        batch_name : str, optional
            Friendly name for the batch.

        Returns
        -------
        BatchCreateResult
            Contains ``batch_id``, ``total_items``, ``status``.
        """
        if not items or not isinstance(items, list):
            raise ValidationError("items must be a non-empty list")

        payload: Dict[str, Any] = {"items": items}
        if batch_name:
            payload["batch_name"] = batch_name

        data = self._client.post(_PREFIX, json_data=payload)
        return BatchCreateResult(data)

    def create_jsonl(
        self,
        *,
        file_path: str,
        batch_name: Optional[str] = None,
    ) -> BatchCreateResult:
        """Create a batch by uploading a JSONL file.

        Each line in the file must be a valid JSON object with ``custom_id``,
        ``model``, and either ``messages`` or ``prompt``.

        Parameters
        ----------
        file_path : str
            Path to a ``.jsonl`` file on disk.
        batch_name : str, optional
            Friendly name for the batch.

        Returns
        -------
        BatchCreateResult
        """
        if not os.path.isfile(file_path):
            raise ValidationError(f"File not found: {file_path}")

        with open(file_path, "rb") as fh:
            file_content = fh.read()

        form_data: Dict[str, str] = {}
        if batch_name:
            form_data["batch_name"] = batch_name

        data = self._client.upload_file(
            f"{_PREFIX}/upload",
            file_field="file",
            file_name=os.path.basename(file_path),
            file_content=file_content,
            content_type="application/jsonl",
            form_data=form_data,
        )
        return BatchCreateResult(data)

    def list(
        self,
        *,
        page: int = 1,
        limit: int = 20,
        status: Optional[str] = None,
    ) -> BatchList:
        """List batch jobs (paginated).

        Parameters
        ----------
        page : int
            Page number (default 1).
        limit : int
            Items per page (default 20, max 100).
        status : str, optional
            Filter by status (``validating``, ``validated``, ``processing``,
            ``completed``, ``failed``, ``cancelled``).

        Returns
        -------
        BatchList
            Iterable of ``Batch`` objects.
        """
        data = self._client.get(_PREFIX, params={
            "page": page,
            "limit": limit,
            "status": status,
        })
        return BatchList(data)

    def get(self, batch_id: int) -> Batch:
        """Get full details of a specific batch.

        Parameters
        ----------
        batch_id : int

        Returns
        -------
        Batch
        """
        data = self._client.get(f"{_PREFIX}/{batch_id}")
        batch_data = data.get("batch", data)
        return Batch(batch_data)

    def start(self, batch_id: int) -> _APIObject:
        """Start processing a validated JSONL batch.

        Parameters
        ----------
        batch_id : int

        Returns
        -------
        _APIObject
        """
        data = self._client.post(f"{_PREFIX}/{batch_id}/start")
        return _APIObject(data)

    def cancel(self, batch_id: int) -> _APIObject:
        """Cancel a batch in ``validating``, ``validated``, or ``processing`` state.

        Parameters
        ----------
        batch_id : int

        Returns
        -------
        _APIObject
        """
        data = self._client.post(f"{_PREFIX}/{batch_id}/cancel")
        return _APIObject(data)

    def progress(self, batch_id: int) -> BatchProgress:
        """Get batch processing progress.

        Parameters
        ----------
        batch_id : int

        Returns
        -------
        BatchProgress
            Contains ``status``, ``progress_percent``, ``processed_items``, etc.
        """
        data = self._client.get(f"{_PREFIX}/{batch_id}/progress")
        progress_data = data.get("progress", data)
        return BatchProgress(progress_data)

    def results(
        self,
        batch_id: int,
        *,
        save_to: Optional[str] = None,
        chunk_size: int = 8192,
    ) -> bytes:
        """Download batch results as JSONL.

        Streams the response to disk when ``save_to`` is provided (recommended
        for large batches to avoid loading all data into memory at once).

        Parameters
        ----------
        batch_id : int
        save_to : str, optional
            If provided, stream the JSONL file to this path on disk.
        chunk_size : int
            Stream chunk size in bytes (default 8 KB).

        Returns
        -------
        bytes
            Raw JSONL content, or ``b""`` when streaming to disk via ``save_to``.
        """
        resp = self._client.get(
            f"{_PREFIX}/{batch_id}/results",
            raw_response=True,
            stream=bool(save_to),
        )

        if save_to:
            with open(save_to, "wb") as fh:
                for data_chunk in resp.iter_content(chunk_size=chunk_size):
                    if data_chunk:
                        fh.write(data_chunk)
            return b""

        return resp.content

    def wait(
        self,
        batch_id: int,
        *,
        poll_interval: float = 5.0,
        timeout: Optional[float] = None,
        on_progress: Optional[Callable[[BatchProgress], None]] = None,
    ) -> Batch:
        """Poll until a batch reaches a terminal state.

        Parameters
        ----------
        batch_id : int
        poll_interval : float
            Seconds between polls (default 5).
        timeout : float, optional
            Max seconds to wait before raising ``APITimeoutError``.
        on_progress : callable, optional
            Called with ``BatchProgress`` on each poll.

        Returns
        -------
        Batch
            The final batch record.

        Raises
        ------
        APITimeoutError
            If *timeout* elapses before the batch finishes.
        """
        import time

        start = time.monotonic()
        terminal = {"completed", "failed", "cancelled"}

        while True:
            prog = self.progress(batch_id)

            if on_progress is not None:
                on_progress(prog)

            if prog.status in terminal:
                return self.get(batch_id)

            if timeout is not None and (time.monotonic() - start) >= timeout:
                raise APITimeoutError(
                    f"Batch {batch_id} did not finish within {timeout}s"
                )

            time.sleep(poll_interval)
