"""Query/chat helper functions for the ARBI SDK.

Provides ``query()`` (simple, returns answer string) and ``query_streaming()``
(returns a full :class:`QueryResult` with typed events and callbacks).
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

from arbi_client.models.agent_step_event import AgentStepEvent
from arbi_client.sdk._types import QueryResult, StreamingEvent
from arbi_client.sdk.streaming import run_streaming_query
from arbi_client.sdk.tools import retrieval_chunk as _build_retrieval_chunk


def query(
    base_url: str,
    access_token: str,
    content: str,
    *,
    doc_ext_ids: list[str] | None = None,
    search_mode: str = "hybrid",
    tools: dict | None = None,
    config_ext_id: str | None = None,
    timeout: float = 600.0,
    verify: bool = True,
) -> str:
    """Run a query and return the answer string.

    If *tools* is not provided, a default ``retrieval_chunk`` tool is built
    from *doc_ext_ids* and *search_mode*.
    """
    if tools is None:
        tools = _build_retrieval_chunk(doc_ext_ids=doc_ext_ids, search_mode=search_mode)

    result = run_streaming_query(
        base_url=base_url,
        access_token=access_token,
        query=content,
        tools=tools,
        config_ext_id=config_ext_id,
        timeout=timeout,
        verify=verify,
    )
    return result.content


def query_streaming(
    base_url: str,
    access_token: str,
    content: str,
    *,
    doc_ext_ids: list[str] | None = None,
    search_mode: str = "hybrid",
    tools: dict | None = None,
    config_ext_id: str | None = None,
    on_token: Callable[[str], Any] | None = None,
    on_agent_step: Callable[[AgentStepEvent], Any] | None = None,
    timeout: float = 600.0,
    verify: bool = True,
) -> QueryResult:
    """Run a streaming query with optional per-event callbacks.

    Args:
        base_url: API base URL.
        access_token: Bearer token (workspace-scoped JWT).
        content: The user query.
        doc_ext_ids: Optional list of document IDs to scope search.
        search_mode: Search mode.
        tools: Explicit tools dict.
        config_ext_id: Optional agent config.
        on_token: Called with each content token string as it arrives.
        on_agent_step: Called with a typed :class:`AgentStepEvent` for each step.
        timeout: HTTP timeout in seconds.
        verify: Whether to verify SSL.

    Returns:
        A :class:`QueryResult` with typed events.
    """
    if tools is None:
        tools = _build_retrieval_chunk(doc_ext_ids=doc_ext_ids, search_mode=search_mode)

    def _on_event(event: StreamingEvent) -> None:
        if event.event_type == "response.output_text.delta" and on_token is not None:
            delta = event.data.get("delta", "") if isinstance(event.data, dict) else ""
            if delta:
                on_token(delta)
        elif isinstance(event.data, AgentStepEvent) and on_agent_step is not None:
            on_agent_step(event.data)

    return run_streaming_query(
        base_url=base_url,
        access_token=access_token,
        query=content,
        tools=tools,
        config_ext_id=config_ext_id,
        on_event=_on_event,
        timeout=timeout,
        verify=verify,
    )
