"""Document helper functions for the ARBI SDK.

Handles file uploads (with WebSocket-based wait), listing, and deletion.
All functions accept the low-level building blocks (client, workspace ID, key header)
so that ``ArbiSession`` can delegate to them without coupling.
"""

from __future__ import annotations

import asyncio
import mimetypes
from collections.abc import Callable
from io import BytesIO
from pathlib import Path
from typing import Any

from arbi_client.api.document import delete_documents as delete_documents_api
from arbi_client.api.document import upload_documents as upload_documents_api
from arbi_client.client import AuthenticatedClient
from arbi_client.models.delete_documents_request import DeleteDocumentsRequest
from arbi_client.models.task_update_message import TaskUpdateMessage
from arbi_client.models.upload_documents_body import UploadDocumentsBody
from arbi_client.sdk.websocket import ArbiWebSocket
from arbi_client.types import File


def _path_to_file(path: str | Path) -> File:
    """Read a file from disk into an :class:`File` suitable for upload."""
    p = Path(path)
    mime, _ = mimetypes.guess_type(p.name)
    if mime is None:
        mime = "application/octet-stream"
    return File(payload=BytesIO(p.read_bytes()), file_name=p.name, mime_type=mime)


def upload_documents(
    client: AuthenticatedClient,
    base_url: str,
    access_token: str,
    file_paths: list[str | Path],
    *,
    wait: bool = True,
    timeout: int = 300,
    shared: bool = True,
    config_ext_id: str | None = None,
    tag_ext_id: str | None = None,
    on_status: Callable[[TaskUpdateMessage], Any] | None = None,
) -> list[str]:
    """Upload files and optionally wait for processing via WebSocket.

    Args:
        client: Authenticated API client.
        base_url: API base URL (needed for WebSocket connection).
        access_token: JWT bearer token (workspace-scoped JWT).
        file_paths: Local file paths to upload.
        wait: If ``True`` (default), block until all documents reach a
            terminal status (``completed``, ``failed``, or ``skipped``)
            using a WebSocket connection.
        timeout: Maximum seconds to wait when *wait* is ``True``.
        shared: Whether documents are shared with workspace members.
        config_ext_id: Optional processing config.
        tag_ext_id: Optional tag to link documents to.
        on_status: Optional callback receiving a :class:`TaskUpdateMessage`
            for every status update via WebSocket while waiting.
            Only used when *wait* is ``True``.

    Returns:
        List of document external IDs in upload order.

    Raises:
        RuntimeError: If the upload API returns an error.
        TimeoutError: If *wait* is ``True`` and documents don't finish in time.
    """
    files = [_path_to_file(p) for p in file_paths]

    kwargs: dict[str, Any] = {
        "client": client,
        "body": UploadDocumentsBody(files=files),
        "shared": shared,
    }
    if config_ext_id is not None:
        kwargs["config_ext_id"] = config_ext_id
    if tag_ext_id is not None:
        kwargs["tag_ext_id"] = tag_ext_id

    result = upload_documents_api.sync(**kwargs)
    if result is None or not hasattr(result, "doc_ext_ids"):
        raise RuntimeError(f"Upload failed: {result}")

    doc_ext_ids: list[str] = result.doc_ext_ids

    if wait and doc_ext_ids:
        _wait_for_docs_sync(base_url, access_token, doc_ext_ids, timeout, on_status)

    return doc_ext_ids


def _wait_for_docs_sync(
    base_url: str,
    access_token: str,
    doc_ext_ids: list[str],
    timeout: int,
    on_status: Callable[[TaskUpdateMessage], Any] | None = None,
) -> dict[str, str]:
    """Block until documents reach terminal status using WebSocket."""

    async def _inner() -> dict[str, str]:
        async with ArbiWebSocket(base_url=base_url, access_token=access_token) as ws:
            return await ws.wait_for_docs(doc_ext_ids, timeout=timeout, on_status=on_status)

    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = None

    if loop is not None and loop.is_running():
        # Already inside an event loop — run in a new thread
        import concurrent.futures

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            return pool.submit(asyncio.run, _inner()).result(timeout=timeout + 5)
    else:
        return asyncio.run(_inner())


def list_documents(
    client: AuthenticatedClient,
) -> list:
    """List all documents in the active workspace.

    Returns the list of document objects from the API.
    The workspace is identified by the workspace-scoped JWT in the client's auth header.
    """
    from arbi_client.api.document import list_documents as list_documents_api

    result = list_documents_api.sync(client=client)
    if result is None:
        return []
    if isinstance(result, list):
        return result
    raise RuntimeError(f"List documents failed: {result}")


def delete_documents(
    client: AuthenticatedClient,
    doc_ext_ids: list[str],
) -> Any:
    """Delete documents by their external IDs.

    Returns the API response.
    """
    return delete_documents_api.sync(
        client=client,
        body=DeleteDocumentsRequest(external_ids=doc_ext_ids),
    )
