"""
Emulated file_search for providers that don't support the tool natively.

Flow:
  1. Convert file_search tools to a single function tool definition.
  2. Call the provider with the function tool.
  3. If the provider issues a file_search function_call, execute vector search
     via litellm.vector_stores.main.asearch().
  4. Feed results back and get the final answer.
  5. Wrap everything in OpenAI Responses-API format:
       [file_search_call output item] + [message output item with file_citation annotations]
"""

import json
import time
import uuid
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, cast

from litellm._internal_context import is_internal_call
from litellm._logging import verbose_logger
from litellm.types.llms.openai import ResponseOutputItem, ResponsesAPIResponse
from litellm.types.vector_stores import VectorStoreSearchResult

# Keep ToolParam broad so we stay compatible with both dict and Pydantic forms
ToolParam = Any

FILE_SEARCH_FUNCTION_NAME = "litellm_file_search"


# ---------------------------------------------------------------------------
# Detection
# ---------------------------------------------------------------------------


def should_use_emulated_file_search(
    tools: Optional[Iterable[ToolParam]],
    provider_config: Any,  # BaseResponsesAPIConfig
) -> bool:
    """Return True when there is a file_search tool and the provider can't handle it natively."""
    if not tools:
        return False
    has_fs = any(isinstance(t, dict) and t.get("type") == "file_search" for t in tools)
    if not has_fs:
        return False
    return provider_config is None or not provider_config.supports_native_file_search()


# ---------------------------------------------------------------------------
# Tool conversion
# ---------------------------------------------------------------------------


def _build_function_tool(vector_store_ids: List[str]) -> Dict[str, Any]:
    """
    Create a Responses API function-tool definition that describes file search.
    The function accepts one or more natural-language queries (like OpenAI's native
    file_search); LiteLLM runs the actual vector search against the configured
    vector stores.

    Note: Uses Responses API format (name/description/parameters at top level),
    NOT Chat Completion format (nested under "function"), so that the
    LiteLLMCompletionResponsesConfig transformation picks up name and description.
    """
    return {
        "type": "function",
        "name": FILE_SEARCH_FUNCTION_NAME,
        "description": (
            "Search the knowledge base for information relevant to the query. "
            "Use this whenever you need to look up specific facts, documents, "
            "or content from the vector store. You can provide multiple queries "
            "to search for different aspects of the information."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "queries": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": (
                        "One or more search queries to look up in the vector store. "
                        "Multiple queries help find comprehensive information from "
                        "different angles."
                    ),
                },
                "vector_store_id": {
                    "type": "string",
                    "description": "ID of the vector store to search.",
                    "enum": vector_store_ids,
                },
            },
            "required": ["queries"],
        },
    }


def _replace_file_search_tools(
    tools: Optional[Iterable[ToolParam]],
) -> Tuple[List[Dict[str, Any]], List[str]]:
    """
    Replace all file_search tools with a single function tool.

    Returns:
        (new_tools_list, all_vector_store_ids)
    """
    non_file_search: List[Dict[str, Any]] = []
    vector_store_ids: List[str] = []

    for tool in tools or []:
        if isinstance(tool, dict) and tool.get("type") == "file_search":
            ids = tool.get("vector_store_ids") or []
            vector_store_ids.extend(ids)
        else:
            non_file_search.append(tool)

    # Deduplicate while preserving order
    unique_ids: List[str] = list(dict.fromkeys(vector_store_ids))
    if unique_ids:
        non_file_search.append(_build_function_tool(unique_ids))

    return non_file_search, unique_ids


# ---------------------------------------------------------------------------
# Search execution
# ---------------------------------------------------------------------------


async def _run_vector_searches(
    queries: List[str],
    vector_store_ids: List[str],
) -> Tuple[List[str], List[VectorStoreSearchResult]]:
    """
    Run `asearch` against all vector stores for all queries and collect results.

    Args:
        queries: List of search queries to execute (like OpenAI's multi-query approach)
        vector_store_ids: Vector store IDs to search

    Returns:
        (queries_list, combined_results)
    """
    import litellm.vector_stores.main as vs_main

    all_results: List[VectorStoreSearchResult] = []
    ids_to_search = vector_store_ids

    # Execute each query against all vector stores
    for query in queries:
        for vs_id in ids_to_search:
            try:
                response = await vs_main.asearch(
                    vector_store_id=vs_id,
                    query=query,
                )
                results_data = (
                    response.get("data")
                    if isinstance(response, dict)
                    else getattr(response, "data", None)
                )
                if results_data:
                    all_results.extend(results_data)
            except Exception as exc:
                verbose_logger.warning(
                    "file_search emulated: search failed for query='%s', vector_store_id='%s': %s",
                    query,
                    vs_id,
                    exc,
                )

    return queries, all_results


# ---------------------------------------------------------------------------
# Result formatting
# ---------------------------------------------------------------------------


def _get_field(result: Any, key: str, default: Any = None) -> Any:
    """Read a field from either a dict/TypedDict or an attribute-based object."""
    if isinstance(result, dict):
        return result.get(key, default)
    return getattr(result, key, default)


def _format_search_results_as_tool_output(
    results: List[VectorStoreSearchResult],
) -> str:
    """Serialize search results into a string to pass back as the tool's output."""
    if not results:
        return "No results found in the vector store."

    parts: List[str] = []
    for i, result in enumerate(results, 1):
        score = _get_field(result, "score")
        file_id = _get_field(result, "file_id")
        filename = _get_field(result, "filename")
        content_items = _get_field(result, "content") or []
        text_chunks = [
            c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "")
            for c in content_items
        ]
        text = " ".join(t for t in text_chunks if t)

        header = f"[Result {i}"
        if filename:
            header += f" | {filename}"
        if file_id:
            header += f" | file_id={file_id}"
        if score is not None:
            header += f" | score={score:.3f}"
        header += "]"

        parts.append(f"{header}\n{text}")

    return "\n\n".join(parts)


def _build_search_results_for_include(
    results: List[VectorStoreSearchResult],
) -> List[Dict[str, Any]]:
    """
    Convert VectorStoreSearchResult objects to the format expected in
    file_search_call.search_results (mirrors OpenAI's include= format).

    All chunks are returned — no deduplication by file_id — matching the
    behaviour of OpenAI's native file_search which surfaces every relevant
    chunk even when multiple chunks originate from the same document.
    """
    formatted: List[Dict[str, Any]] = []
    for result in results:
        file_id = _get_field(result, "file_id") or ""
        content_items = _get_field(result, "content") or []
        text_chunks = [
            c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "")
            for c in content_items
        ]
        text = " ".join(t for t in text_chunks if t)
        formatted.append(
            {
                "file_id": file_id,
                "filename": _get_field(result, "filename") or "",
                "score": _get_field(result, "score"),
                "text": text,
                "attributes": _get_field(result, "attributes") or {},
            }
        )
    return formatted


def _build_file_search_call_output(
    call_id: str,
    queries: List[str],
    results: Optional[List[VectorStoreSearchResult]] = None,
    include_search_results: bool = False,
) -> Dict[str, Any]:
    """Build the file_search_call output item (mirrors OpenAI's format).

    Args:
        call_id: Unique ID for this file_search call.
        queries: List of search queries used.
        results: The raw search results (used when include_search_results=True).
        include_search_results: Populate search_results when the caller passed
            ``include=["file_search_call.results"]``.
    """
    search_results = None
    if include_search_results and results:
        search_results = _build_search_results_for_include(results)
    return {
        "type": "file_search_call",
        "id": call_id,
        "status": "completed",
        "queries": queries,
        "search_results": search_results,
    }


def _build_file_citation_annotations(
    results: List[VectorStoreSearchResult],
    text: str,
) -> List[Dict[str, Any]]:
    """
    Build file_citation annotations for the text.
    Each result with a file_id gets a citation at the end of the text.
    """
    annotations: List[Dict[str, Any]] = []
    index = len(text)  # cite at end of text block
    seen_file_ids: set = set()

    for result in results:
        file_id = _get_field(result, "file_id")
        filename = _get_field(result, "filename")
        if not file_id or file_id in seen_file_ids:
            continue
        seen_file_ids.add(file_id)
        annotations.append(
            {
                "type": "file_citation",
                "index": index,
                "file_id": file_id,
                "filename": filename or "",
            }
        )

    return annotations


def _build_message_output(
    response_text: str,
    results: List[VectorStoreSearchResult],
) -> Dict[str, Any]:
    """Build the message output item with optional file_citation annotations."""
    annotations = _build_file_citation_annotations(results, response_text)
    return {
        "type": "message",
        "role": "assistant",
        "content": [
            {
                "type": "output_text",
                "text": response_text,
                "annotations": annotations,
            }
        ],
    }


def _extract_text_from_responses_output(response: ResponsesAPIResponse) -> str:
    """Pull the assistant's text from the provider's response."""
    for item in response.output:
        item_type = (
            item.get("type") if isinstance(item, dict) else getattr(item, "type", None)
        )
        if item_type == "message":
            content = (
                item.get("content")
                if isinstance(item, dict)
                else getattr(item, "content", [])
            )
            for block in content or []:
                block_type = (
                    block.get("type")
                    if isinstance(block, dict)
                    else getattr(block, "type", None)
                )
                if block_type == "output_text":
                    raw = (
                        block.get("text")
                        if isinstance(block, dict)
                        else getattr(block, "text", "")
                    )
                    return str(raw) if raw is not None else ""
    return ""


def _synthesize_responses_api_response(
    original_response: ResponsesAPIResponse,
    file_search_call_output: Dict[str, Any],
    message_output: Dict[str, Any],
    first_response: Optional[ResponsesAPIResponse] = None,
) -> ResponsesAPIResponse:
    """
    Return a new ResponsesAPIResponse with:
      output[0] = file_search_call item
      output[1] = message item (with citations)

    When first_response is provided, its response_cost is accumulated into the
    synthesized _hidden_params so that billing callbacks see the total cost of
    both provider calls that the emulated flow makes.
    """
    synthesized_output: List[Dict[str, Any]] = [file_search_call_output, message_output]
    synthesized = ResponsesAPIResponse(
        id=getattr(original_response, "id", f"resp_{uuid.uuid4().hex}"),
        object="response",
        created_at=getattr(original_response, "created_at", int(time.time())),
        status="completed",
        model=getattr(original_response, "model", ""),
        output=cast(
            List[Union[ResponseOutputItem, Dict[str, Any]]], synthesized_output
        ),
        usage=getattr(original_response, "usage", None),
        error=None,
    )
    if hasattr(original_response, "_hidden_params"):
        hidden = dict(getattr(original_response, "_hidden_params") or {})
        if first_response is not None and hasattr(first_response, "_hidden_params"):
            first_hidden = getattr(first_response, "_hidden_params") or {}
            first_cost = (
                first_hidden.get("response_cost")
                if isinstance(first_hidden, dict)
                else getattr(first_hidden, "response_cost", None)
            )
            if first_cost is not None:
                current_cost = (
                    hidden.get("response_cost") if isinstance(hidden, dict) else 0
                )
                hidden["response_cost"] = (current_cost or 0) + first_cost
        synthesized._hidden_params = hidden
    return synthesized


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------


async def _call_aresponses(
    input, model, tools, **kwargs
):  # pragma: no cover – thin wrapper for patching in tests
    from litellm.responses.main import aresponses

    return await aresponses(input=input, model=model, tools=tools, **kwargs)


def _prepare_emulated_file_search_call(
    kwargs: Dict[str, Any],
) -> Tuple[bool, Dict[str, Any]]:
    include_items: List[str] = list(kwargs.get("include") or [])
    include_search_results = "file_search_call.results" in include_items

    original_stream = kwargs.get("stream")
    updated_kwargs = kwargs
    if original_stream:
        verbose_logger.debug(
            "Streaming is not yet supported for emulated file_search. "
            "Disabling stream for this request."
        )
        updated_kwargs = {**kwargs, "stream": False}

    return include_search_results, updated_kwargs


def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> Tuple[str, str]:
    """Extract (call_id, raw_arguments_string) from a dict or Pydantic tool_call item."""
    if isinstance(tool_call, dict):
        call_id = str(
            tool_call.get("call_id") or tool_call.get("id") or fallback_call_id
        )
        raw_args = tool_call.get("arguments") or "{}"
    else:
        raw_call_id = (
            getattr(tool_call, "call_id", None)
            or getattr(tool_call, "id", None)
            or fallback_call_id
        )
        call_id = str(raw_call_id)
        raw_args = getattr(tool_call, "arguments", "{}") or "{}"
    return call_id, raw_args


def _resolve_queries_from_args(args: Dict[str, Any], input: Any) -> List[str]:
    """Pull the queries list out of parsed tool-call arguments, with backward-compat fallbacks."""
    queries_from_call = args.get("queries")
    if not queries_from_call:
        # Fallback: check for single "query" field (backward compat)
        single_query = args.get("query")
        return [single_query] if single_query else [str(input)]
    if not isinstance(queries_from_call, list):
        return [str(queries_from_call)]
    return queries_from_call


async def _execute_file_search_tool_calls(
    file_search_calls: List[Any],
    all_vs_ids: List[str],
    input: Any,
    file_search_call_id: str,
) -> Tuple[List[Dict[str, Any]], List[str], List[VectorStoreSearchResult]]:
    """Run the vector search for each file_search tool_call and collect results."""
    tool_results: List[Dict[str, Any]] = []
    all_queries: List[str] = []
    all_results: List[VectorStoreSearchResult] = []

    for tool_call in file_search_calls:
        call_id, raw_args = _extract_tool_call_fields(
            tool_call, fallback_call_id=file_search_call_id
        )

        try:
            args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
        except json.JSONDecodeError:
            args = {}

        queries_from_call = _resolve_queries_from_args(args, input)

        vs_id_arg = args.get("vector_store_id")
        vs_ids_for_call = [vs_id_arg] if vs_id_arg else all_vs_ids

        queries, results = await _run_vector_searches(
            queries=queries_from_call,
            vector_store_ids=vs_ids_for_call,
        )
        all_queries.extend(queries)
        all_results.extend(results)

        tool_results.append(
            {
                "type": "function_call_output",
                "call_id": call_id,
                "output": _format_search_results_as_tool_output(results),
            }
        )

    return tool_results, all_queries, all_results


def _build_follow_up_input(
    input: Any,
    first_response: ResponsesAPIResponse,
    tool_results: List[Dict[str, Any]],
) -> List[Any]:
    """Assemble the follow-up call input: original messages + first-response output + tool results.

    Including all output items (text blocks, reasoning, non-file-search calls) ensures providers
    like Anthropic that emit text before the tool call have complete conversation context.
    Serializes Pydantic model instances to plain dicts so the transformation layer can call .get().
    """
    original_input_items = (
        list(input)
        if isinstance(input, (list, tuple))
        else [{"role": "user", "content": str(input)}]
    )
    first_response_output_items: List[Any] = []
    for _item in first_response.output:
        if isinstance(_item, dict):
            first_response_output_items.append(_item)
        elif hasattr(_item, "model_dump"):
            first_response_output_items.append(_item.model_dump(exclude_none=True))  # type: ignore[union-attr]
        else:
            first_response_output_items.append(_item)

    return original_input_items + first_response_output_items + tool_results


async def aresponses_with_emulated_file_search(
    input: Any,
    model: str,
    tools: Optional[Iterable[ToolParam]] = None,
    # Pass-through params — forwarded as-is to the underlying aresponses call
    **kwargs: Any,
) -> ResponsesAPIResponse:
    """
    Emulated file_search for providers that don't support it natively.

    Replaces file_search tools with a function tool, intercepts the tool call,
    runs vector search, and synthesizes an OpenAI-format response.
    """
    # Determine whether caller wants search_results populated in the output.
    _include_search_results, kwargs = _prepare_emulated_file_search_call(kwargs=kwargs)

    # 1. Replace file_search tools with function tool
    transformed_tools, all_vs_ids = _replace_file_search_tools(tools)

    # 2. First provider call — provider will call the file_search function.
    # Mark as an internal sub-call so wrapper_async skips billing callbacks;
    # the parent litellm_logging_obj (propagated via kwargs) fires once at the end.
    _prev_internal = is_internal_call.get()
    is_internal_call.set(True)
    try:
        first_response: ResponsesAPIResponse = cast(
            ResponsesAPIResponse,
            await _call_aresponses(
                input=input,
                model=model,
                tools=transformed_tools or None,
                **kwargs,
            ),
        )
    finally:
        is_internal_call.set(_prev_internal)

    # 3. Look for a file_search function_call in the output
    file_search_calls = [
        item
        for item in first_response.output
        if (
            isinstance(item, dict)
            and item.get("type") == "function_call"
            and item.get("name") == FILE_SEARCH_FUNCTION_NAME
        )
        or (
            hasattr(item, "type")
            and getattr(item, "type") == "function_call"
            and getattr(item, "name", None) == FILE_SEARCH_FUNCTION_NAME
        )
    ]

    if not file_search_calls:
        # Provider answered without calling the tool (e.g. it had enough context).
        # Return as-is wrapped in OpenAI format.
        call_id = f"fs_{uuid.uuid4().hex[:24]}"
        response_text = _extract_text_from_responses_output(first_response)
        return _synthesize_responses_api_response(
            original_response=first_response,
            file_search_call_output=_build_file_search_call_output(
                call_id=call_id,
                queries=[str(input)],
                results=None,
                include_search_results=False,
            ),
            message_output=_build_message_output(response_text, []),
        )

    # 4. Execute each file_search tool call
    file_search_call_id = f"fs_{uuid.uuid4().hex[:24]}"
    tool_results, all_queries, all_results = await _execute_file_search_tool_calls(
        file_search_calls=file_search_calls,
        all_vs_ids=all_vs_ids,
        input=input,
        file_search_call_id=file_search_call_id,
    )

    # 5. Build follow-up input: original messages + ALL first-response output items + tool results
    follow_up_input = _build_follow_up_input(
        input=input,
        first_response=first_response,
        tool_results=tool_results,
    )

    # 6. Follow-up call — provider writes the final answer given search results.
    # Also an internal sub-call; billing is suppressed so the outer call fires once.
    is_internal_call.set(True)
    try:
        final_response: ResponsesAPIResponse = cast(
            ResponsesAPIResponse,
            await _call_aresponses(
                input=follow_up_input,
                model=model,
                tools=None,  # no tools needed for the answer step
                **kwargs,
            ),
        )
    finally:
        is_internal_call.set(_prev_internal)

    # 7. Synthesize OpenAI-format output
    response_text = _extract_text_from_responses_output(final_response)

    return _synthesize_responses_api_response(
        original_response=final_response,
        file_search_call_output=_build_file_search_call_output(
            call_id=file_search_call_id,
            queries=all_queries or [str(input)],
            results=all_results,
            include_search_results=_include_search_results,
        ),
        message_output=_build_message_output(response_text, all_results),
        first_response=first_response,
    )
