"""Ask-docs tool for querying GPU documentation via wafer-api.

Calls the wafer-api docs endpoint which runs a docs agent in a persistent
Modal sandbox with the corpus mounted. The user never sees corpus files.
"""
from __future__ import annotations

import json
import os

import httpx

from wafer_core.rollouts.dtypes import (
    Tool,
    ToolCall,
    ToolFunction,
    ToolFunctionParameter,
    ToolResult,
)

VALID_CORPORA = [
    "cuda", "cutlass", "hip", "amd", "cdna3",
    "hopper", "rdna35", "llvm-amdgpu", "gcnasm",
]

DOCS_QUERY_TIMEOUT = 120  # seconds

ASK_DOCS_TOOL = Tool(
    type="function",
    function=ToolFunction(
        name="ask_docs",
        description=(
            "Query GPU documentation. Searches wafer's documentation corpus "
            "and returns an answer grounded in official docs."
        ),
        parameters=ToolFunctionParameter(
            type="object",
            properties={
                "query": {
                    "type": "string",
                    "description": "The question to answer",
                },
                "corpus": {
                    "type": "string",
                    "enum": VALID_CORPORA,
                    "description": "Documentation corpus to search",
                },
            },
        ),
        required=["query", "corpus"],
    ),
)


async def exec_ask_docs(tool_call: ToolCall) -> ToolResult:
    """Execute ask_docs by calling wafer-api /v1/docs/query.

    Reads WAFER_API_URL and WAFER_AUTH_TOKEN from env.
    Streams SSE response from the API and accumulates the full answer.
    """
    args = tool_call.args
    query = args.get("query", "")
    corpus = args.get("corpus", "")

    if not query:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="'query' is required",
        )
    if corpus not in VALID_CORPORA:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error=f"Unknown corpus: {corpus!r}. Valid: {', '.join(VALID_CORPORA)}",
        )

    api_url = os.environ.get("WAFER_API_URL")
    auth_token = os.environ.get("WAFER_AUTH_TOKEN")

    if not api_url:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="WAFER_API_URL not set. Run 'wafer settings login' first.",
        )
    if not auth_token:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="WAFER_AUTH_TOKEN not set. Run 'wafer settings login' first.",
        )

    url = f"{api_url.rstrip('/')}/v1/docs/query"
    headers = {
        "Authorization": f"Bearer {auth_token}",
        "Content-Type": "application/json",
        "Accept": "text/event-stream",
    }
    payload = {"query": query, "corpus": corpus}

    accumulated = []
    try:
        async with httpx.AsyncClient(timeout=DOCS_QUERY_TIMEOUT) as client:
            async with client.stream(
                "POST", url, json=payload, headers=headers
            ) as response:
                if response.status_code == 401:
                    return ToolResult(
                        tool_call_id=tool_call.id,
                        is_error=True,
                        content="",
                        error="Authentication failed. Run 'wafer settings login'.",
                    )
                if response.status_code == 402:
                    return ToolResult(
                        tool_call_id=tool_call.id,
                        is_error=True,
                        content="",
                        error="Insufficient credits. Check 'wafer settings billing'.",
                    )
                if response.status_code != 200:
                    body = await response.aread()
                    return ToolResult(
                        tool_call_id=tool_call.id,
                        is_error=True,
                        content="",
                        error=f"API error ({response.status_code}): {body.decode()[:500]}",
                    )

                async for line in response.aiter_lines():
                    if not line.startswith("data: "):
                        continue
                    data_str = line[len("data: "):]
                    if data_str == "[DONE]":
                        break
                    event = json.loads(data_str)
                    if event.get("type") == "text":
                        accumulated.append(event.get("content", ""))
                    elif event.get("type") == "error":
                        return ToolResult(
                            tool_call_id=tool_call.id,
                            is_error=True,
                            content="",
                            error=event.get("content", "Unknown error from docs agent"),
                        )

    except httpx.TimeoutException:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error=f"Docs query timed out after {DOCS_QUERY_TIMEOUT}s",
        )
    except httpx.ConnectError:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error=f"Cannot reach wafer API at {api_url}",
        )

    answer = "".join(accumulated)
    if not answer:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="Docs agent returned empty response",
        )

    return ToolResult(
        tool_call_id=tool_call.id,
        is_error=False,
        content=answer,
    )
