"""xAI tools: function calling (cheap) and server-side tools (web_search, x_search, code_interpreter, mcp, file_search — RUN_EXPENSIVE_TESTS)."""
from __future__ import annotations

import json

import pytest

TOOLS_CHAT = [{"type": "function", "function": {"name": "get_weather", "description": "Get weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]
TOOLS_RESP = [{"type": "function", "name": "get_weather", "description": "Get weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}]


def _text(body):
    return next((c["text"] for o in body["output"] if o["type"] == "message" for c in o["content"] if c["type"] == "output_text"), None)


def test_chat_function_call_roundtrip(xai, models):
    msgs = [{"role": "user", "content": "Weather in Paris?"}]
    st, b, _ = xai("POST", "/v1/chat/completions", {"model": models["xai"], "messages": msgs, "max_completion_tokens": 64, "tools": TOOLS_CHAT, "tool_choice": {"type": "function", "function": {"name": "get_weather"}}}, est_cost_usd=0.0005, note="test_tools chat forced")
    assert st == 200 and b["choices"][0]["finish_reason"] == "tool_calls"
    tc = b["choices"][0]["message"]["tool_calls"][0]
    assert tc["type"] == "function" and tc["function"]["name"] == "get_weather" and "city" in json.loads(tc["function"]["arguments"])
    msgs += [b["choices"][0]["message"], {"role": "tool", "tool_call_id": tc["id"], "content": "{\"temp_c\": 18}"}]
    st, b2, _ = xai("POST", "/v1/chat/completions", {"model": models["xai"], "messages": msgs, "max_completion_tokens": 32, "tools": TOOLS_CHAT}, est_cost_usd=0.0005, note="test_tools chat round trip")
    assert st == 200 and "18" in b2["choices"][0]["message"]["content"]


def test_responses_function_call_roundtrip(xai, models):
    st, r, _ = xai("POST", "/v1/responses", {"model": models["xai"], "input": "Weather in Paris?", "tools": TOOLS_RESP, "tool_choice": {"type": "function", "name": "get_weather"}, "max_output_tokens": 2000}, est_cost_usd=0.0005, note="test_tools responses forced")
    assert st == 200
    fc = next(o for o in r["output"] if o["type"] == "function_call")
    assert fc["name"] == "get_weather" and fc["call_id"].startswith("call-")
    st, r2, _ = xai("POST", "/v1/responses", {"model": models["xai"], "previous_response_id": r["id"], "input": [{"type": "function_call_output", "call_id": fc["call_id"], "output": "{\"temp_c\": 18}"}], "tools": TOOLS_RESP, "max_output_tokens": 2000}, est_cost_usd=0.0005, note="test_tools responses function_call_output")
    assert st == 200 and "18" in _text(r2)
    for rid in (r["id"], r2["id"]):
        xai("DELETE", f"/v1/responses/{rid}", note="test_tools cleanup")


def test_chat_rejects_server_tools(xai, models):
    st, b, _ = xai("POST", "/v1/chat/completions", {"model": models["xai"], "messages": [{"role": "user", "content": "Reply with OK."}], "max_completion_tokens": 32, "tools": [{"type": "web_search"}]}, est_cost_usd=0, note="test_tools chat web_search 422")
    assert st == 422 and "unknown variant" in (b.decode() if isinstance(b, bytes) else json.dumps(b))


def test_unknown_tool_type_lists_accepted_types(xai, models):
    st, b, _ = xai("POST", "/v1/responses", {"model": models["xai"], "input": "Reply with OK.", "store": False, "tools": [{"type": "web_search_preview"}]}, est_cost_usd=0, note="test_tools web_search_preview 422")
    txt = b.decode() if isinstance(b, bytes) else json.dumps(b)
    assert st == 422
    for t in ("function", "web_search", "x_search", "code_interpreter", "file_search", "mcp"):
        assert f"`{t}`" in txt


def _run_tool(xai, models, tools, prompt, note, include=None):
    body = {"model": models["xai"], "input": prompt, "tools": tools, "store": False, "max_output_tokens": 4000, "max_turns": 3}
    if include:
        body["include"] = include
    st, r, _ = xai("POST", "/v1/responses", body, est_cost_usd=0.02, note=note, timeout=300)
    assert st == 200, r
    return r


@pytest.mark.run_expensive_tests
def test_web_search_tool(xai, models):
    r = _run_tool(xai, models, [{"type": "web_search"}], "Use the web_search tool to open https://x.ai/news and give the newest article title. Max 8 words.", "test_tools web_search", ["web_search_call.action.sources"])
    calls = [o for o in r["output"] if o["type"] == "web_search_call"]
    assert calls and calls[0]["action"]["type"] in ("search", "open_page", "find_in_page")
    assert r["usage"]["server_side_tool_usage_details"]["web_search_calls"] >= 1


@pytest.mark.run_expensive_tests
def test_x_search_tool(xai, models):
    r = _run_tool(xai, models, [{"type": "x_search", "allowed_x_handles": ["xai"]}], "What is the latest post by @xai about? Max 8 words.", "test_tools x_search")
    types = {o["type"] for o in r["output"]}
    assert types & {"custom_tool_call", "x_search_call"}  # live: custom_tool_call (x_keyword_search…)
    assert r["usage"]["server_side_tool_usage_details"]["x_search_calls"] >= 1


@pytest.mark.run_expensive_tests
def test_code_interpreter_tool(xai, models):
    r = _run_tool(xai, models, [{"type": "code_interpreter"}], "Run print(2+2) in Python and reply with the output only.", "test_tools code_interpreter", ["code_interpreter_call.outputs"])
    ci = next(o for o in r["output"] if o["type"] == "code_interpreter_call")
    assert "print" in ci["code"] and ci["outputs"] and ci["outputs"][0]["type"] == "logs"
    assert "4" in _text(r)


@pytest.mark.run_expensive_tests
def test_mcp_tool(xai, models):
    r = _run_tool(xai, models, [{"type": "mcp", "server_url": "https://mcp.deepwiki.com/mcp", "server_label": "deepwiki"}], "Using the deepwiki tool, what does the xai-org/xai-sdk-python repo do? Max 10 words.", "test_tools mcp")
    mc = next(o for o in r["output"] if o["type"] == "mcp_call")
    assert mc["server_label"] == "deepwiki" and mc["output"]
    assert r["usage"]["server_side_tool_usage_details"]["mcp_calls"] >= 1


@pytest.mark.run_expensive_tests
def test_file_search_tool_shape(xai, models):
    st, col, _ = xai("POST", "/v1/collections", {"collection_name": "atlas-test", "field_definitions": []}, note="test_tools collection create")
    assert st == 200
    cid = col["collection_id"]
    try:
        r = _run_tool(xai, models, [{"type": "file_search", "vector_store_ids": [cid]}], "What is in my documents? Three words.", "test_tools file_search", ["file_search_call.results"])
        fs = next(o for o in r["output"] if o["type"] == "file_search_call")
        assert fs["status"] in ("completed", "failed") and isinstance(fs["queries"], list)
    finally:
        xai("DELETE", f"/v1/collections/{cid}", note="test_tools collection delete")