"""Tools a model may call while it answers.

One tool so far -- web search -- but the shape is the point: a registry of
named callables with a JSON schema each, offered to the endpoint and executed
here when it asks. Built-in tools, MCP servers and agentic execution all plug
in at the same place.

Two things gate whether a tool is offered at all:

* the administrator has configured and enabled it, and
* the chat's model is marked as supporting tools.

The second is not optional politeness. Sending a ``tools`` array to an endpoint
that does not implement tool calling fails the entire request, exactly the way
sending image parts to a model without vision does -- and for the same reason,
the capability flag on the model is what decides.
"""

from __future__ import annotations

import json
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any

from sqlalchemy.orm import Session as DBSession

from lembas.db.models import Chat, User
from lembas.services import search as search_service
from lembas.services import settings_store
from lembas.services.search.base import SearchError

log = logging.getLogger(__name__)

# How many times a model may call tools before it has to answer with words.
# Not a safety limit so much as a termination one: a small model that has
# decided searching is the answer will otherwise search until the context runs
# out, and each round costs a full request.
MAX_ROUNDS = 3

WEB_SEARCH = "web_search"

WEB_SEARCH_SCHEMA: dict[str, Any] = {
    "type": "function",
    "function": {
        "name": WEB_SEARCH,
        "description": (
            "Search the web for current information. Use this when the answer "
            "depends on recent events, on facts you are unsure of, or on "
            "anything that may have changed since your training data. Returns "
            "a numbered list of results with titles, URLs and short extracts."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search terms. Keep them short and specific.",
                },
                "max_results": {
                    "type": "integer",
                    "description": "How many results to return. Defaults to the site setting.",
                },
            },
            "required": ["query"],
        },
    },
}


@dataclass
class ToolOutcome:
    """What running a tool produced, for the model and for the reader.

    The two are deliberately different. `content` is the flat text the model
    reads back; `event` is what the transcript shows, and keeps the results
    structured so they can be rendered as links rather than as a wall of URLs.
    """

    content: str
    event: dict[str, Any] = field(default_factory=dict)


def enabled_tools(db: DBSession, chat: Chat, user: User | None) -> list[dict[str, Any]]:
    """The tool schemas to offer for this chat, which is usually none."""
    from lembas.security import permissions
    from lembas.services import chat as chat_service

    config = settings_store.search(db)
    if not config.get("enabled"):
        return []
    if not permissions.has(db, user, "tools.web_search"):
        return []
    if not chat_service.model_supports(db, chat, "tools"):
        return []
    if search_service.availability(str(config.get("provider") or "ddgs")):
        # Configured but unusable -- offering a tool that will fail on every
        # call is worse than not offering it.
        return []
    return [WEB_SEARCH_SCHEMA]


async def run_tool(config: dict[str, Any], name: str, arguments: str) -> ToolOutcome:
    """Execute one tool call.

    Never raises. A tool that fails hands the model an explanation and lets it
    carry on -- a failed search should produce "I could not look that up"
    rather than killing the whole reply.
    """
    if name != WEB_SEARCH:
        return ToolOutcome(
            content=f"There is no tool called {name!r}.",
            event={"name": name, "status": "error", "error": "Unknown tool."},
        )

    try:
        parsed = json.loads(arguments) if arguments.strip() else {}
    except json.JSONDecodeError:
        # Small models emit malformed argument JSON often enough that this is a
        # normal path, not an exceptional one. Treat the whole string as the
        # query rather than giving up.
        parsed = {"query": arguments.strip()}
    if not isinstance(parsed, dict):
        parsed = {"query": str(parsed)}

    query = str(parsed.get("query") or "").strip()
    if not query:
        return ToolOutcome(
            content="No search query was given.",
            event={"name": name, "status": "error", "error": "No query was given."},
        )

    limit = parsed.get("max_results")
    try:
        limit = int(limit) if limit is not None else None
    except (TypeError, ValueError):
        limit = None

    try:
        results = await search_service.run(config, query, limit=limit)
    except SearchError as exc:
        log.info("web search failed for %r: %s", query[:60], exc.message)
        return ToolOutcome(
            content=f"The search failed: {exc.message}",
            event={"name": name, "query": query, "status": "error", "error": exc.message},
        )

    event = {
        "name": name,
        "query": query,
        "status": "ok",
        "results": [
            {"title": r.title, "url": r.url, "snippet": r.snippet, "host": r.host}
            for r in results
        ],
    }

    if not results:
        return ToolOutcome(content=f"No results were found for {query!r}.", event=event)

    lines = [f"Search results for {query!r}:"]
    for index, result in enumerate(results, start=1):
        lines.append(f"\n[{index}] {result.title}\n{result.url}\n{result.snippet}")
    return ToolOutcome(content="\n".join(lines), event=event)


class ToolCallAccumulator:
    """Reassembles tool calls arriving as streamed fragments.

    An endpoint sends ``delta.tool_calls`` as a list of partial objects: the id
    and the function name arrive once, and ``arguments`` arrives as a string
    split across however many chunks the tokeniser produced. Entries are keyed
    by ``index`` because that is the only field guaranteed on every fragment --
    the id is absent from continuations, and matching on name breaks the moment
    a model calls the same tool twice in one turn.
    """

    def __init__(self) -> None:
        self._calls: dict[int, dict[str, Any]] = {}

    def feed(self, fragments: list[dict[str, Any]]) -> None:
        for fragment in fragments:
            if not isinstance(fragment, dict):
                continue
            index = fragment.get("index")
            if not isinstance(index, int):
                # Some servers omit index entirely when there is only one call.
                index = 0
            call = self._calls.setdefault(index, {"id": "", "name": "", "arguments": ""})

            if fragment.get("id"):
                call["id"] = str(fragment["id"])
            function = fragment.get("function") or {}
            if isinstance(function, dict):
                if function.get("name"):
                    call["name"] = str(function["name"])
                arguments = function.get("arguments")
                if isinstance(arguments, str):
                    call["arguments"] += arguments

    @property
    def calls(self) -> list[dict[str, Any]]:
        """Completed calls, in the order the endpoint indexed them."""
        return [
            {
                # An id is required when the results are sent back, and not
                # every server supplies one.
                "id": call["id"] or f"call_{index}",
                "name": call["name"],
                "arguments": call["arguments"],
            }
            for index, call in sorted(self._calls.items())
            if call["name"]
        ]

    def __bool__(self) -> bool:
        return bool(self.calls)


def assistant_turn(calls: list[dict[str, Any]], content: str) -> dict[str, Any]:
    """The assistant message to send back with the tool results.

    The endpoint needs its own tool_calls echoed before the tool replies, or it
    has nothing to match the tool_call_ids against.
    """
    return {
        "role": "assistant",
        "content": content or None,
        "tool_calls": [
            {
                "id": call["id"],
                "type": "function",
                "function": {"name": call["name"], "arguments": call["arguments"]},
            }
            for call in calls
        ],
    }


def tool_turn(call: dict[str, Any], content: str) -> dict[str, Any]:
    return {
        "role": "tool",
        "tool_call_id": call["id"],
        "name": call["name"],
        "content": content,
    }


ToolRunner = Callable[..., Awaitable[ToolOutcome]]

__all__ = [
    "MAX_ROUNDS",
    "WEB_SEARCH",
    "ToolCallAccumulator",
    "ToolOutcome",
    "assistant_turn",
    "enabled_tools",
    "run_tool",
    "tool_turn",
]