# client.py
import json
import os
import warnings
import httpx
import time
from importlib.metadata import version as _pkg_version
from pathlib import Path
from typing import Any

from poma.exceptions import ApiError, InvalidInputError, InvalidResponseError, PomaSDKError
from poma.models import PomaResult
from poma.retrieval import generate_cheatsheets
from poma.utils import unpack_poma_archive


try:
    USER_AGENT = f"poma-ai-sdk/{_pkg_version('poma')}"
except Exception:
    USER_AGENT = "poma-ai-sdk/unknown"

API_BASE_URL = "https://api.poma-ai.com/v2"


def _format_debug_dump(data: Any) -> str:
    try:
        return json.dumps(data, indent=2, sort_keys=True, ensure_ascii=False)
    except TypeError:
        return repr(data)


def _build_content_disposition(filename: str) -> str:
    safe_filename = (
        filename.replace("\\", "\\\\")
        .replace('"', '\\"')
        .replace("\r", "")
        .replace("\n", "")
    )
    return f'attachment; filename="{safe_filename}"'


def _extract_status_error_details(status_data: dict[str, Any]) -> str:
    direct_error = status_data.get("error")
    if direct_error:
        return str(direct_error)

    detail = status_data.get("detail")
    if isinstance(detail, dict):
        nested_error = detail.get("error")
        if nested_error:
            return str(nested_error)
        return _format_debug_dump(detail)
    if detail:
        return str(detail)

    details = status_data.get("details")
    if isinstance(details, dict):
        nested_error = details.get("error")
        if nested_error:
            return str(nested_error)
        return _format_debug_dump(details)
    if details:
        return str(details)

    return "No details provided."


def _should_retry_with_form_upload(
    error: httpx.HTTPStatusError,
    form_endpoint: str,
) -> bool:
    body = (error.response.text or "").lower()
    form_hint = form_endpoint.lower()
    return "multipart form data is not allowed" in body or form_hint in body


def _post_ingest_request(
    client: httpx.Client,
    base_api_url: str,
    path: Path,
    payload: dict[str, Any],
    *,
    raw_endpoint: str = "ingest",
    form_endpoint: str = "ingestForm",
    show_progress: bool = False,
) -> httpx.Response:
    file_bytes = path.read_bytes()
    request_headers = {
        "Content-Type": "application/octet-stream",
        "Content-Disposition": _build_content_disposition(path.name),
    }

    try:
        response = client.post(
            f"{base_api_url}/{raw_endpoint}",
            content=file_bytes,
            headers=request_headers,
            params=payload,
        )
        response.raise_for_status()
        return response
    except httpx.HTTPStatusError as error:
        if not _should_retry_with_form_upload(error, form_endpoint):
            raise

    response = client.post(
        f"{base_api_url}/{form_endpoint}",
        data=payload,
        files={"file": (path.name, file_bytes)},
    )
    response.raise_for_status()
    return response


def _submit_ingest_job(
    client: httpx.Client,
    base_api_url: str,
    file_path: str | os.PathLike[str],
    *,
    base_url: str | None = None,
    raw_endpoint: str = "ingest",
    form_endpoint: str = "ingestForm",
    show_progress: bool = False,
) -> dict[str, Any]:
    if file_path is None or (isinstance(file_path, str) and not file_path.strip()):
        raise ValueError("file_path must be a non-empty string or path-like.")
    path = Path(file_path)
    payload: dict[str, Any] = {"is_sdk": True}
    if base_url:
        payload["base_url"] = base_url
    try:
        response = _post_ingest_request(
            client,
            base_api_url,
            path,
            payload,
            raw_endpoint=raw_endpoint,
            form_endpoint=form_endpoint,
            show_progress=show_progress,
        )
    except httpx.HTTPStatusError as error:
        status = error.response.status_code
        body = (error.response.text or "").strip()
        msg = f"Failed to submit file '{path}'"
        if body:
            msg = f"{msg}: {body}"
        raise ApiError(
            msg, status_code=status, response_body=body, file_path=str(path)
        ) from error
    try:
        return response.json()
    except ValueError as error:
        raise InvalidResponseError(
            f"Server returned non-JSON or empty body (file: {path})"
        ) from error


class Poma:
    """
    Client for interacting with the POMA API.
    """

    def __init__(
        self,
        api_key: str | None = None,
        *,
        timeout: float = 600.0,
        client: httpx.Client | None = None,
    ):
        """
        Initialize the POMA client.
        Args:
            api_key (str, optional):
                API key for authenticating with POMA. If not provided,
                the value is read from the environment variable `POMA_API_KEY`.
            timeout (float, default=600.0):
                Timeout (in seconds) for all HTTP requests.
            client (httpx.Client, optional):
                Custom HTTP client. If not provided, a default client is created.
        """
        api_base_url = API_BASE_URL
        # Override API base URL if environment variable is set
        if os.environ.get("API_BASE_URL"):
            api_base_url = os.environ.get("API_BASE_URL")
        if not api_base_url:
            raise ValueError("API base URL cannot be empty.")

        self.base_api_url = api_base_url.rstrip("/")
        self._client = client or httpx.Client(
            timeout=timeout, headers={"user-agent": USER_AGENT}
        )
        if not (api_key := api_key or os.environ.get("POMA_API_KEY", "")):
            raise Exception("POMA_API_KEY environment variable not set.")
        self._client.headers.update({"Authorization": f"Bearer {api_key}"})
        warnings.warn(
            "Poma is deprecated and will be removed in a future version. "
            "Use PrimeCut instead: from poma import PrimeCut",
            DeprecationWarning,
            stacklevel=2,
        )

    def start_chunk_file(
        self,
        file_path: str | os.PathLike[str],
        *,
        base_url: str | None = None,
        show_progress: bool = False,
    ) -> dict[str, Any]:
        """
        Submit a file with text to POMA for ingestion.

        Uses ``POST /ingest`` with a raw ``application/octet-stream``
        body and automatically retries on ``POST /ingestForm`` when the
        primary v2 upload path is rejected. On HTTP error, raises
        :exc:`ApiError` with
        :attr:`ApiError.status_code` and :attr:`ApiError.response_body` set from
        the API response. Typical status codes: 400 (bad request / failed to create job),
        401 (missing or invalid authorization), 403 (not found, quota exceeded, too many jobs),
        500 (internal error, failed to upload file).

        Args:
            file_path (str | os.PathLike[str]):
                Path to the input file (string or path-like). Must have an allowed file extension.
            base_url (str, optional):
                Optional base URL to resolve relative links within the file.
            show_progress (bool, default=False):
                Retained for compatibility. Upload requests do not log verbose debug output.
        Returns:
            A dictionary containing a unique job identifier for the submitted job.

        Raises:
            ApiError: On API HTTP error (see :attr:`ApiError.status_code` and
                :attr:`ApiError.response_body`).
            InvalidResponseError: If the response body is not valid JSON.
        """
        return _submit_ingest_job(
            self._client,
            self.base_api_url,
            file_path,
            base_url=base_url,
            raw_endpoint="ingest",
            form_endpoint="ingestForm",
            show_progress=show_progress,
        )

    def start_chunk_file_eco(
        self,
        file_path: str | os.PathLike[str],
        *,
        base_url: str | None = None,
        show_progress: bool = False,
    ) -> dict[str, Any]:
        """
        Submit a file with text to POMA eco ingestion.

        Uses ``POST /ingestEco`` with a raw ``application/octet-stream``
        body and automatically retries on ``POST /ingestEcoForm`` when
        the primary v2 upload path is rejected.
        """
        return _submit_ingest_job(
            self._client,
            self.base_api_url,
            file_path,
            base_url=base_url,
            raw_endpoint="ingestEco",
            form_endpoint="ingestEcoForm",
            show_progress=show_progress,
        )

    def get_chunk_result(
        self,
        job_id: str,
        *,
        initial_delay: float = 5.0,
        poll_interval: float = 3.0,
        max_interval: float = 15.0,
        show_progress: bool = False,
        download_dir: str | os.PathLike[str] | None = None,
        filename: str | None = None,
    ) -> dict[str, Any]:
        """
        Poll POMA for the result of an ingestion job until completion.

        Uses ``GET /jobs/{job_id}/status``. On HTTP error, raises :exc:`ApiError` with
        :attr:`ApiError.status_code` and :attr:`ApiError.response_body` set from the API.
        Typical status codes: 401 (authorization), 403 (account not found), 404 (job not found
        or no access), 500 (failed to get job result). If the job completes with status
        ``"failed"``, also raises :exc:`ApiError` (without status_code/response_body).

        Args:
            job_id (str):
                The unique identifier of the submitted job.
            initial_delay (float, default=5.0):
                Initial delay (in seconds) before the first poll request.
            poll_interval (float, default=1.0):
                Starting interval (in seconds) between polling requests.
            max_interval (float, default=15.0):
                Maximum interval (in seconds) between polling requests.
            show_progress (bool, default=False):
                If True, logs progress messages during polling.
            download_dir (str | os.PathLike[str], optional):
                Directory to save the downloaded file in. If neither download_dir nor
                filename is set, the result is returned in memory (no file saved). If
                filename is set but download_dir is not, the file is saved in the
                current directory.
            filename (str, optional):
                Name for the saved .poma file. If it does not end with ``.poma``, that
                suffix is appended. If not set when saving to disk, uses the server
                filename when provided, otherwise ``{job_id}.poma``.
        Returns:
            The JSON result containing at least the keys `chunks` and `chunksets`.

        Raises:
            ApiError: On API HTTP error or when job status is ``"failed"`` (see
                :attr:`ApiError.status_code` and :attr:`ApiError.response_body` where set).
            InvalidResponseError: On unexpected job status in the response.
        """
        time.sleep(initial_delay)
        current_interval = poll_interval
        last_status = None

        while True:
            time.sleep(current_interval)
            try:
                response = self._client.get(f"{self.base_api_url}/jobs/{job_id}/status")
                response.raise_for_status()
                data = response.json()
                status = data.get("status", "")
                if status == "done":
                    download = data.get("download", {})
                    download_url = download.get("download_url", "")
                    if not download_url:
                        raise RuntimeError(
                            "Failed to receive download URL from server."
                        )

                    if download_dir is None and filename is None:
                        # Return bytes content instead of saving to file
                        file_bytes = self.download_bytes(download_url)
                        return self.unpack_poma_archive(
                            poma_archive_data=file_bytes
                        )
                    else:
                        # Save downloaded file (to download_dir or current dir if only filename set)
                        save_filename = (
                            filename or download.get("filename") or f"{job_id}.poma"
                        )
                        if not save_filename.endswith(".poma"):
                            save_filename = f"{save_filename}.poma"
                        save_dir = (
                            download_dir
                            if download_dir not in (None, "")
                            else "."
                        )
                        downloaded_file_path = self.download_file(
                            download_url, save_filename, save_directory=save_dir
                        )
                        return self.unpack_poma_archive(
                            poma_archive_path=downloaded_file_path
                        )
                elif status == "failed":
                    error_code = data.get("code", "unknown")
                    error_details = data.get("error", "No details provided.")
                    error_message = (
                        f"Job failed with code {error_code}: {error_details}"
                    )
                    raise ApiError(
                        f"Job {job_id} failed: {data.get('error', error_message)}",
                        job_id=job_id,
                    )
                elif status == "processing":
                    if show_progress:
                        print(f"Job {job_id} is still processing...")
                    if last_status == "pending":
                        current_interval = poll_interval
                    current_interval = min(current_interval * 1.5, max_interval)
                elif status == "pending":
                    if show_progress:
                        print(
                            f"Job {job_id} is pending (queued due to rate limiting, sequential processing - common on demo accounts)..."
                        )
                    current_interval = min(current_interval * 1.5, max_interval)
                else:
                    raise InvalidResponseError(
                        f"Unexpected job status for job {job_id}: {status}"
                    )
                last_status = status
            except httpx.HTTPStatusError as error:
                status = error.response.status_code
                body = (error.response.text or "").strip()
                msg = f"Failed to get job status for '{job_id}'"
                if body:
                    msg = f"{msg}: {body}"
                raise ApiError(
                    msg,
                    status_code=status,
                    response_body=body,
                    job_id=job_id,
                ) from error
            except PomaSDKError:
                raise
            except Exception as error:
                raise RuntimeError(f"POMA-AI job polling failed: {error}") from error

    def unpack_poma_archive(
        self,
        poma_archive_data: bytes | None = None,
        poma_archive_path: str | os.PathLike[str] | None = None,
    ) -> dict[str, Any]:
        """
        Extract chunks and chunksets from a POMA archive; delegates to module-level function.
        POMA archive is a zip file containing chunks.json and chunksets.json.
        Args:
            poma_archive_data: The POMA archive as bytes.
            poma_archive_path: Path to the POMA archive file.
        Returns:
            dict: A dictionary with ``chunks`` and ``chunksets`` keys.
        """
        return unpack_poma_archive(
            poma_archive_data=poma_archive_data,
            poma_archive_path=poma_archive_path,
        )

    # Backwards-compatible alias
    extract_chunks_and_chunksets_from_poma_archive = unpack_poma_archive

    def create_cheatsheet(
        self,
        relevant_chunksets: list[dict[str, Any]],
        all_chunks: list[dict[str, Any]],
    ) -> str:
        """
        Generates a single cheatsheet for one single document
        from relevant chunksets (relevant for a certain query)
        and from all chunks of that document (providing the textual content).

        Args:
            relevant_chunksets (list[dict]): A list of chunksets, each containing a "chunks" key with a list of chunk IDs.
            all_chunks (list[dict]): A list of all chunk dictionaries of the same document, each representing a chunk of content.
        Returns:
            str: The textual content of the generated cheatsheet.

        Raises:
            InvalidInputError: If ``relevant_chunksets`` or ``all_chunks`` is not provided (None),
                not a list, or empty.
            InvalidResponseError: If the generated result has an unexpected shape (e.g. no content).
            ValueError: If the underlying generator raises (e.g. invalid chunkset structure,
                duplicate chunk_index, or file_id mismatch).
        """
        if relevant_chunksets is None or all_chunks is None:
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be provided (cannot be None)."
            )
        if not isinstance(relevant_chunksets, list) or not isinstance(all_chunks, list):
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be lists."
            )
        if not relevant_chunksets or not all_chunks:
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be non-empty."
            )
        cheatsheets = generate_cheatsheets(relevant_chunksets, all_chunks)
        if (
            not cheatsheets
            or not isinstance(cheatsheets, list)
            or len(cheatsheets) == 0
            or "content" not in cheatsheets[0]
        ):
            raise InvalidResponseError(
                "Cheatsheet could not be created from input chunks (unexpected result shape)."
            )
        return cheatsheets[0]["content"]

    def create_cheatsheets(
        self,
        relevant_chunksets: list[dict[str, Any]],
        all_chunks: list[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        """
        Generates cheatsheets from relevant chunksets (relevant for a certain query)
        and from all the chunks of all affected documents (providing the textual content).
        One cheatsheet is created for each document found in the chunks (tagged with file_id).

        Args:
            relevant_chunksets (list[dict]): A list of chunksets, each containing a "chunks" key with a list of chunk IDs.
            all_chunks (list[dict]): A list of all available chunk dictionaries of affected documents, each representing a chunk of content.
        Returns:
            list[dict]: A list of dictionaries representing the generated cheatsheets, each containing:
                - 'file_id': The tag associated with the respective document.
                - 'content': The textual content of the generated cheatsheet.

        Raises:
            InvalidInputError: If ``relevant_chunksets`` or ``all_chunks`` is not provided (None) or not a list.
            ValueError: If the underlying generator raises (e.g. invalid chunkset structure,
                duplicate chunk_index, or file_id mismatch).
        """
        if relevant_chunksets is None or all_chunks is None:
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be provided (cannot be None)."
            )
        if not isinstance(relevant_chunksets, list) or not isinstance(all_chunks, list):
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be lists."
            )
        return generate_cheatsheets(relevant_chunksets, all_chunks)

    def download_file(
        self,
        download_url: str,
        filename: str | None = None,
        *,
        save_directory: str | os.PathLike[str] | None = None,
    ) -> str:
        """
        Download a file from the given download URL.
        Args:
            download_url (str):
                The URL to download the file from.
            filename (str, optional):
                The filename to save the file as. If not provided, will be extracted from URL.
            save_directory (str | os.PathLike[str], optional):
                Directory to save the file in. If not provided, saves to current directory.
        Returns:
            str: The path to the downloaded file.
        """
        if not download_url:
            raise ValueError("download_url cannot be empty")

        # Determine filename
        if not filename:
            filename = Path(download_url).name or "downloaded_file"

        # Determine save directory (default current directory so path has a parent)
        if save_directory:
            save_path = Path(save_directory) / filename
        else:
            save_path = Path(".") / filename

        # Create the directory if it doesn't exist (skip when file is in cwd)
        parent = os.path.dirname(save_path)
        if parent:
            os.makedirs(parent, exist_ok=True)

        # Download the file data
        content = self.download_bytes(download_url)

        # Save the file
        with open(save_path, "wb") as f:
            f.write(content)

        return str(save_path)

    def download_bytes(
        self,
        download_url: str,
    ) -> bytes:
        """
        Download a file from the given download URL and return the bytes content.

        On HTTP error, raises :exc:`ApiError` with :attr:`ApiError.status_code` and
        :attr:`ApiError.response_body` set from the API. For job downloads
        (``GET /jobs/{job_id}/download``), typical codes: 400 (invalid job_id/UUID),
        401 (authorization), 403 (account not found), 404 (job download not found),
        500 (failed to download job data).

        Args:
            download_url (str):
                The URL to download the file from.
        Returns:
            bytes: The content of the downloaded file as bytes.

        Raises:
            ApiError: On API HTTP error (see :attr:`ApiError.status_code` and
                :attr:`ApiError.response_body`).
        """
        if not download_url:
            raise ValueError("download_url cannot be empty")

        # Construct the full URL if it's a relative path
        if download_url.startswith("/"):
            full_url = f"{self.base_api_url}{download_url}"
        else:
            full_url = download_url

        try:
            response = self._client.get(full_url)
            response.raise_for_status()
            return response.content
        except httpx.HTTPStatusError as error:
            status = error.response.status_code
            body = (error.response.text or "").strip()
            msg = f"Failed to download '{download_url}'"
            if body:
                msg = f"{msg}: {body}"
            raise ApiError(
                msg,
                status_code=status,
                response_body=body,
                url=download_url,
            ) from error
        except PomaSDKError:
            raise
        except Exception as error:
            raise RuntimeError(f"File download failed: {error}") from error

    def close(self):
        self._client.close()

    def __enter__(self) -> "Poma":
        return self

    def __exit__(self, *exc):
        self.close()


class PrimeCut:
    """
    Client for interacting with the POMA API.

    Provides a simple interface to ingest documents and get structured
    chunks and chunksets back as typed objects.
    """

    def __init__(
        self,
        api_key: str | None = None,
        *,
        timeout: float = 600.0,
        client: httpx.Client | None = None,
    ):
        """
        Initialize the PrimeCut client.

        Args:
            api_key: API key for authenticating with POMA. Falls back to
                the ``POMA_API_KEY`` environment variable.
            timeout: Timeout (in seconds) for all HTTP requests.
            client: Optional custom ``httpx.Client``.
        """
        api_base_url = os.environ.get("API_BASE_URL") or API_BASE_URL
        if not api_base_url:
            raise ValueError("API base URL cannot be empty.")

        self.base_api_url = api_base_url.rstrip("/")
        self._client = client or httpx.Client(
            timeout=timeout, headers={"user-agent": USER_AGENT}
        )
        if not (api_key := api_key or os.environ.get("POMA_API_KEY", "")):
            raise Exception("POMA_API_KEY environment variable not set.")
        self._client.headers.update({"Authorization": f"Bearer {api_key}"})

    def submit(
        self,
        file_path: str | os.PathLike[str],
        *,
        base_url: str | None = None,
        show_progress: bool = False,
    ) -> str:
        """
        Submit a file for ingestion and return the job ID.

        Uses ``POST /ingest`` with a raw ``application/octet-stream``
        body and automatically retries on ``POST /ingestForm`` when the
        primary v2 upload path is rejected. Use :meth:`poll` to wait for
        the result, or :meth:`ingest` to do both in one call.

        Args:
            file_path: Path to the input file.
            base_url: Optional base URL to resolve relative links in the file.
            show_progress: Retained for compatibility. Upload requests do not log verbose debug output.

        Returns:
            The job ID string.

        Raises:
            ApiError: On API HTTP error.
            InvalidResponseError: If the response body is not valid JSON.
        """
        data = _submit_ingest_job(
            self._client,
            self.base_api_url,
            file_path,
            base_url=base_url,
            raw_endpoint="ingest",
            form_endpoint="ingestForm",
            show_progress=show_progress,
        )
        job_id = data.get("job_id", "")
        if not job_id:
            raise InvalidResponseError("Server did not return a job_id.")
        return job_id

    def submit_eco(
        self,
        file_path: str | os.PathLike[str],
        *,
        base_url: str | None = None,
        show_progress: bool = False,
    ) -> str:
        """
        Submit a file for eco ingestion and return the job ID.

        Uses ``POST /ingestEco`` with a raw ``application/octet-stream``
        body and automatically retries on ``POST /ingestEcoForm`` when
        the primary v2 upload path is rejected.
        """
        data = _submit_ingest_job(
            self._client,
            self.base_api_url,
            file_path,
            base_url=base_url,
            raw_endpoint="ingestEco",
            form_endpoint="ingestEcoForm",
            show_progress=show_progress,
        )
        job_id = data.get("job_id", "")
        if not job_id:
            raise InvalidResponseError("Server did not return a job_id.")
        return job_id

    def poll(
        self,
        job_id: str,
        *,
        initial_delay: float = 5.0,
        poll_interval: float = 3.0,
        max_interval: float = 15.0,
        show_progress: bool = False,
        download_dir: str | os.PathLike[str] | None = None,
        filename: str | None = None,
    ) -> "PomaArchive":
        """
        Poll for the result of a previously submitted ingestion job.

        Args:
            job_id: The job ID returned by :meth:`submit`.
            initial_delay: Seconds to wait before the first poll.
            poll_interval: Starting interval between polls.
            max_interval: Maximum poll interval.
            show_progress: Log progress messages while polling.
            download_dir: Directory to save the ``.poma`` archive to.
            filename: Name for the saved ``.poma`` file.

        Returns:
            A :class:`PomaArchive` wrapping the raw archive data.

        Raises:
            ApiError: On HTTP or job-level errors.
            InvalidResponseError: On unexpected job status.
        """
        time.sleep(initial_delay)
        current_interval = poll_interval
        last_status = None

        while True:
            time.sleep(current_interval)
            try:
                status_url = f"{self.base_api_url}/jobs/{job_id}/status"
                resp = self._client.get(status_url)
                resp.raise_for_status()
                status_data = resp.json()
                status = status_data.get("status", "")

                if status == "done":
                    download = status_data.get("download", {})
                    download_url = download.get("download_url", "")
                    if not download_url:
                        raise RuntimeError(
                            "Failed to receive download URL from server."
                        )

                    if download_dir is None and filename is None:
                        file_bytes = self._download_bytes(download_url)
                        return PomaArchive(data=file_bytes)
                    else:
                        save_filename = (
                            filename
                            or download.get("filename")
                            or f"{job_id}.poma"
                        )
                        if not save_filename.endswith(".poma"):
                            save_filename = f"{save_filename}.poma"
                        save_dir = (
                            download_dir
                            if download_dir not in (None, "")
                            else "."
                        )
                        saved_path = self._download_file(
                            download_url, save_filename, save_directory=save_dir
                        )
                        return PomaArchive(path=saved_path)

                elif status == "failed":
                    error_code = status_data.get("code", "unknown")
                    error_details = _extract_status_error_details(status_data)
                    raise ApiError(
                        f"Job {job_id} failed with code {error_code}: "
                        f"{error_details}",
                        job_id=job_id,
                        response_body=_format_debug_dump(status_data),
                    )
                elif status == "processing":
                    if show_progress:
                        print(f"Job {job_id} is still processing...")
                    if last_status == "pending":
                        current_interval = poll_interval
                    current_interval = min(current_interval * 1.5, max_interval)
                elif status == "pending":
                    if show_progress:
                        print(
                            f"Job {job_id} is pending (queued due to rate "
                            "limiting, sequential processing - common on "
                            "demo accounts)..."
                        )
                    current_interval = min(current_interval * 1.5, max_interval)
                else:
                    raise InvalidResponseError(
                        f"Unexpected job status for job {job_id}: {status}"
                    )
                last_status = status

            except httpx.HTTPStatusError as error:
                s = error.response.status_code
                body = (error.response.text or "").strip()
                msg = f"Failed to get job status for '{job_id}'"
                if body:
                    msg = f"{msg}: {body}"
                raise ApiError(
                    msg, status_code=s, response_body=body, job_id=job_id
                ) from error
            except PomaSDKError:
                raise
            except Exception as error:
                raise RuntimeError(
                    f"POMA-AI job polling failed: {error}"
                ) from error

    def ingest(
        self,
        file_path: str | os.PathLike[str],
        *,
        base_url: str | None = None,
        initial_delay: float = 5.0,
        poll_interval: float = 3.0,
        max_interval: float = 15.0,
        show_progress: bool = False,
        download_dir: str | os.PathLike[str] | None = None,
        filename: str | None = None,
    ) -> PomaResult:
        """
        Submit a file for ingestion and wait for the result.

        Convenience method that combines :meth:`submit`, :meth:`poll`,
        and unpacking into a single call.

        Args:
            file_path: Path to the input file.
            base_url: Optional base URL to resolve relative links in the file.
            initial_delay: Seconds to wait before the first poll.
            poll_interval: Starting interval between polls.
            max_interval: Maximum poll interval.
            show_progress: Log progress messages while polling.
            download_dir: Directory to save the ``.poma`` archive to.
            filename: Name for the saved ``.poma`` file.

        Returns:
            A :class:`PomaResult` with typed chunks and chunksets.
        """
        job_id = self.submit(
            file_path,
            base_url=base_url,
            show_progress=show_progress,
        )
        archive = self.poll(
            job_id,
            initial_delay=initial_delay,
            poll_interval=poll_interval,
            max_interval=max_interval,
            show_progress=show_progress,
            download_dir=download_dir,
            filename=filename,
        )
        return archive.unpack()

    def ingest_eco(
        self,
        file_path: str | os.PathLike[str],
        *,
        base_url: str | None = None,
        initial_delay: float = 5.0,
        poll_interval: float = 3.0,
        max_interval: float = 15.0,
        show_progress: bool = False,
        download_dir: str | os.PathLike[str] | None = None,
        filename: str | None = None,
    ) -> PomaResult:
        """
        Submit a file for eco ingestion and wait for the result.

        Convenience method that combines :meth:`submit_eco`,
        :meth:`poll`, and unpacking into a single call.
        """
        job_id = self.submit_eco(
            file_path,
            base_url=base_url,
            show_progress=show_progress,
        )
        archive = self.poll(
            job_id,
            initial_delay=initial_delay,
            poll_interval=poll_interval,
            max_interval=max_interval,
            show_progress=show_progress,
            download_dir=download_dir,
            filename=filename,
        )
        return archive.unpack()

    def create_cheatsheet(
        self,
        relevant_chunksets: list[dict[str, Any]],
        all_chunks: list[dict[str, Any]],
    ) -> str:
        """
        Generate a single cheatsheet string from relevant chunksets and all
        chunks of one document.
        """
        if relevant_chunksets is None or all_chunks is None:
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be provided."
            )
        if not isinstance(relevant_chunksets, list) or not isinstance(all_chunks, list):
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be lists."
            )
        if not relevant_chunksets or not all_chunks:
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be non-empty."
            )
        cheatsheets = generate_cheatsheets(relevant_chunksets, all_chunks)
        if (
            not cheatsheets
            or not isinstance(cheatsheets, list)
            or len(cheatsheets) == 0
            or "content" not in cheatsheets[0]
        ):
            raise InvalidResponseError(
                "Cheatsheet could not be created from input chunks."
            )
        return cheatsheets[0]["content"]

    def create_cheatsheets(
        self,
        relevant_chunksets: list[dict[str, Any]],
        all_chunks: list[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        """
        Generate cheatsheets for multiple documents from relevant chunksets
        and all chunks.
        """
        if relevant_chunksets is None or all_chunks is None:
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be provided."
            )
        if not isinstance(relevant_chunksets, list) or not isinstance(all_chunks, list):
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be lists."
            )
        return generate_cheatsheets(relevant_chunksets, all_chunks)

    # --- private helpers ---

    def _download_bytes(self, download_url: str) -> bytes:
        if download_url.startswith("/"):
            full_url = f"{self.base_api_url}{download_url}"
        else:
            full_url = download_url
        try:
            response = self._client.get(full_url)
            response.raise_for_status()
            return response.content
        except httpx.HTTPStatusError as error:
            s = error.response.status_code
            body = (error.response.text or "").strip()
            msg = f"Failed to download '{download_url}'"
            if body:
                msg = f"{msg}: {body}"
            raise ApiError(
                msg, status_code=s, response_body=body, url=download_url
            ) from error

    def _download_file(
        self,
        download_url: str,
        filename: str,
        *,
        save_directory: str | os.PathLike[str] | None = None,
    ) -> str:
        if save_directory:
            save_path = Path(save_directory) / filename
        else:
            save_path = Path(".") / filename
        parent = os.path.dirname(save_path)
        if parent:
            os.makedirs(parent, exist_ok=True)
        content = self._download_bytes(download_url)
        with open(save_path, "wb") as f:
            f.write(content)
        return str(save_path)

    def close(self):
        self._client.close()

    def __enter__(self) -> "PrimeCut":
        return self

    def __exit__(self, *exc):
        self.close()


class PomaArchive:
    """
    Wrapper around a raw ``.poma`` archive (zip containing chunks.json
    and chunksets.json).

    Returned by :meth:`PrimeCut.ingest`. Call :meth:`unpack` to get a
    typed :class:`PomaResult`.
    """

    def __init__(
        self,
        data: bytes | None = None,
        path: str | os.PathLike[str] | None = None,
    ):
        self._data = data
        self._path = str(path) if path else None

    def unpack(self) -> PomaResult:
        """Extract and return a :class:`PomaResult` with typed chunks, chunksets, and images.

        Images from ``image_sources.json`` are resolved to base64 data URIs.
        If an entry is not already base64 (e.g. ``None`` or a URL), the
        corresponding file from the ``assets/`` folder is used as a fallback.
        """
        raw = unpack_poma_archive(
            poma_archive_data=self._data,
            poma_archive_path=self._path,
        )
        return PomaResult.from_dict(raw)
