"""
Extended SplitRuns client with polling utilities.

Example:
    from extend_ai import Extend

    client = Extend(token="...")

    # Create and poll until completion
    result = client.split_runs.create_and_poll(
        file={"id": "file_xxx"},
        splitter={"id": "splitter_abc123"},
    )

    if result.status == "PROCESSED":
        print(result.output)
"""

from typing import Any, Dict, Optional

from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ...requests.split_config import SplitConfigParams
from ...split_runs.client import AsyncSplitRunsClient as GeneratedAsyncSplitRunsClient
from ...split_runs.client import SplitRunsClient as GeneratedSplitRunsClient
from ...split_runs.requests.split_runs_create_request_file import SplitRunsCreateRequestFileParams
from ...split_runs.requests.split_runs_create_request_splitter import SplitRunsCreateRequestSplitterParams
from ...types.run_metadata import RunMetadata
from ...types.run_priority import RunPriority
from ...types.split_run import SplitRun
from ..polling import PollingOptions, poll_until_done, poll_until_done_async

# Re-export for convenience
from ..polling import PollingTimeoutError

__all__ = ["SplitRunsClient", "AsyncSplitRunsClient", "PollingTimeoutError"]


def _is_terminal_status(status: str) -> bool:
    """
    Check if a ProcessorRunStatus is terminal (no longer processing).
    We check for non-terminal states rather than terminal states so that
    if new terminal states are added, polling will still complete.
    """
    return status not in ("PROCESSING", "PENDING", "CANCELLING")


class SplitRunsClient(GeneratedSplitRunsClient):
    """
    Extended SplitRuns client with create_and_poll method.

    Inherits all methods from SplitRunsClient and adds create_and_poll
    for convenient polling until completion.
    """

    def __init__(self, *, client_wrapper: SyncClientWrapper):
        super().__init__(client_wrapper=client_wrapper)

    def create_and_poll(
        self,
        *,
        file: SplitRunsCreateRequestFileParams,
        splitter: Optional[SplitRunsCreateRequestSplitterParams] = None,
        config: Optional[SplitConfigParams] = None,
        priority: Optional[RunPriority] = None,
        metadata: Optional[RunMetadata] = None,
        polling_options: Optional[PollingOptions] = None,
    ) -> SplitRun:
        """
        Creates a split run and polls until it reaches a terminal state.

        Terminal states: PROCESSED, FAILED, CANCELLED

        Args:
            file: The file to split.
            splitter: Reference to an existing splitter.
            config: Inline split configuration.
            priority: Priority of the run.
            metadata: Additional metadata for the run.
            polling_options: Options for polling behavior.

        Returns:
            The final split run when processing is complete.

        Raises:
            PollingTimeoutError: If the run doesn't complete within max_wait_ms.

        Example:
            result = client.split_runs.create_and_poll(
                file={"id": "file_xxx"},
                splitter={"id": "splitter_abc123"}
            )

            if result.status == "PROCESSED":
                print(result.output)
        """
        # Build kwargs, only including non-None values to avoid passing null
        kwargs: Dict[str, Any] = {"file": file}
        if splitter is not None:
            kwargs["splitter"] = splitter
        if config is not None:
            kwargs["config"] = config
        if priority is not None:
            kwargs["priority"] = priority
        if metadata is not None:
            kwargs["metadata"] = metadata

        # Create the split run
        create_response = self.create(**kwargs)
        run_id = create_response.id

        # Poll until terminal state
        return poll_until_done(
            retrieve=lambda: self.retrieve(run_id),
            is_terminal=lambda response: _is_terminal_status(response.status),
            options=polling_options,
        )


class AsyncSplitRunsClient(GeneratedAsyncSplitRunsClient):
    """
    Extended AsyncSplitRuns client with create_and_poll method.
    """

    def __init__(self, *, client_wrapper: AsyncClientWrapper):
        super().__init__(client_wrapper=client_wrapper)

    async def create_and_poll(
        self,
        *,
        file: SplitRunsCreateRequestFileParams,
        splitter: Optional[SplitRunsCreateRequestSplitterParams] = None,
        config: Optional[SplitConfigParams] = None,
        priority: Optional[RunPriority] = None,
        metadata: Optional[RunMetadata] = None,
        polling_options: Optional[PollingOptions] = None,
    ) -> SplitRun:
        """
        Creates a split run and polls until it reaches a terminal state (async version).
        """
        # Build kwargs, only including non-None values to avoid passing null
        kwargs: Dict[str, Any] = {"file": file}
        if splitter is not None:
            kwargs["splitter"] = splitter
        if config is not None:
            kwargs["config"] = config
        if priority is not None:
            kwargs["priority"] = priority
        if metadata is not None:
            kwargs["metadata"] = metadata

        # Create the split run
        create_response = await self.create(**kwargs)
        run_id = create_response.id

        # Poll until terminal state
        return await poll_until_done_async(
            retrieve=lambda: self.retrieve(run_id),
            is_terminal=lambda response: _is_terminal_status(response.status),
            options=polling_options,
        )
