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

Usage::

    from agnipod import AgniPod

    client = AgniPod()

    # Create a batch from a list of items
    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",
            "messages": [{"role": "user", "content": "What is 3+3?"}],
        },
    ])
    print(result.batch_id)

    # Start processing
    client.batches.start(result.batch_id)

    # Check progress
    progress = client.batches.progress(result.batch_id)
    print(f"{progress.progress_percent}% complete")

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

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

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


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 batch job from a JSON list of items.

        Each item should have ``custom_id``, ``model``, and ``messages`` keys.

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

        Returns
        -------
        BatchCreateResult
            Contains ``batch_id``, ``total_items``, ``is_valid``.
        """
        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("/batch", 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.

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

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

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

        # Use multipart upload for JSONL
        url = f"{self._client.base_url}/batch/upload"
        files = {"file": (os.path.basename(file_path), file_content, "application/jsonl")}
        form_data = {}
        if batch_name:
            form_data["batch_name"] = batch_name

        resp = self._client._session.post(url, files=files, data=form_data, timeout=self._client.timeout)

        from ._exceptions import raise_for_status
        from ._client import _safe_json
        body = _safe_json(resp)
        raise_for_status(resp.status_code, body)
        return BatchCreateResult(body)

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

        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 with ``.pagination``.
        """
        data = self._client.get("/batch", params={
            "page": page,
            "limit": limit,
            "status": status,
        })
        return BatchList(data)

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

        Parameters
        ----------
        batch_id : int
            The batch ID.

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

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

        Parameters
        ----------
        batch_id : int
            The batch ID.

        Returns
        -------
        _APIObject
            Contains ``success``, ``message``, ``batch_id``.
        """
        data = self._client.post(f"/batch/{batch_id}/start")
        return _APIObject(data)

    def cancel(self, batch_id: int) -> _APIObject:
        """Cancel a batch.

        Parameters
        ----------
        batch_id : int
            The batch ID.

        Returns
        -------
        _APIObject
            Contains ``success``, ``message``, ``batch_id``.
        """
        data = self._client.post(f"/batch/{batch_id}/cancel")
        return _APIObject(data)

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

        Parameters
        ----------
        batch_id : int
            The batch ID.

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

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

        Parameters
        ----------
        batch_id : int
            The batch ID.
        save_to : str, optional
            If provided, save the file to this path on disk.

        Returns
        -------
        bytes
            Raw JSONL content.
        """
        resp = self._client.get(f"/batch/{batch_id}/results", raw_response=True)

        content = resp.content

        if save_to:
            with open(save_to, "wb") as f:
                f.write(content)

        return content

    def wait(
        self,
        batch_id: int,
        *,
        poll_interval: float = 5.0,
        timeout: Optional[float] = None,
        on_progress: Any = None,
    ) -> Batch:
        """Poll until a batch completes (or fails/is cancelled).

        Parameters
        ----------
        batch_id : int
            The batch ID.
        poll_interval : float
            Seconds between polls (default 5).
        timeout : float, optional
            Max seconds to wait (``None`` = forever).
        on_progress : callable, optional
            Called with a ``BatchProgress`` on each poll.

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

        Raises
        ------
        TimeoutError
            If *timeout* is reached.
        """
        import time
        from ._exceptions import APITimeoutError

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

        while True:
            prog = self.progress(batch_id)

            if on_progress:
                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)
