#include "agent_smith/detail/tool_dispatch.hpp"

#include "agent_smith/detail/gather_all.hpp"
#include "agent_smith/detail/timeout.hpp"

namespace agent_smith::internal
{

struct ToolAccess
{
    static std::shared_ptr<Tool> retain(ToolRegistry& registry,
                                      const std::string& name)
    {
        const auto found = registry.by_name_.find(name);
        return found == registry.by_name_.end() ? nullptr : found->second.tool;
    }
};

namespace
{

mw::Task<mw::E<ToolResult>> invokeOwned(std::shared_ptr<Tool> tool,
                                      nlohmann::json arguments,
                                      ToolContext context)
{
    co_return co_await tool->invoke(std::move(arguments), std::move(context));
}

mw::Task<mw::E<ToolResult>> unknownTool(std::string name)
{
    co_return std::unexpected(mw::runtimeError("unknown tool: " + name));
}

std::string errorContent(const std::string& message)
{
    return nlohmann::json({{"error", true}, {"message", message}}).dump();
}

} // namespace

mw::Task<mw::E<std::vector<ToolMessage>>> dispatchTools(
    Runtime& runtime, ToolRegistry& registry, std::vector<ToolCall> calls,
    AgentConfig config, std::optional<std::string>* final_message)
{
    std::vector<StartedTask<ToolResult>> started;
    started.reserve(calls.size());
    for(const auto& call : calls)
    {
        auto tool = ToolAccess::retain(registry, call.name);
        if(tool == nullptr)
        {
            started.push_back(StartedTask<ToolResult>::start(
                unknownTool(call.name)));
            continue;
        }
        const auto budget = registry.options(call.name).timeout.value_or(
            config.default_tool_timeout);
        CancellationSource cancel(config.tree_state
            ? config.tree_state->cancel.token() : CancellationToken{});
        ToolContext context{
            .cancel = cancel.token(),
            .budget = budget,
            .call_id = call.id,
            .logger = spdlog::default_logger_raw(),
            .runtime = &runtime,
        };
        started.push_back(StartedTask<ToolResult>::start(withTimeout(
            runtime, invokeOwned(std::move(tool), call.arguments,
                                 std::move(context)),
            budget, &cancel)));
    }

    auto results = co_await gatherAllStarted(std::move(started));
    std::vector<ToolMessage> messages;
    messages.reserve(calls.size());
    for(std::size_t index = 0; index < calls.size(); ++index)
    {
        auto& result = results[index];
        const bool is_error = !result || result->is_error;
        if(!is_error && final_message != nullptr && !*final_message &&
           result->structured.is_object())
        {
            const auto& data = result->structured;
            if(data.contains("_done") && data["_done"] == true &&
               data.contains("final_message") && data["final_message"].is_string())
            {
                *final_message = data["final_message"].get<std::string>();
            }
        }
        std::string content = !result ? errorContent(result.error().msg())
            : result->is_error ? errorContent(result->content)
                               : std::move(result->content);
        if(content.size() > config.tool_result_max_bytes)
        {
            const auto original_size = content.size();
            auto end = config.tool_result_max_bytes;
            // Do not split a UTF-8 code point in a serializable message.
            while(end > 0 &&
                  (static_cast<unsigned char>(content[end]) & 0xc0) == 0x80)
            {
                --end;
            }
            content.resize(end);
            content += "\n…[truncated, original size " +
                       std::to_string(original_size) + " bytes]";
        }
        messages.push_back(ToolMessage{
            .tool_call_id = calls[index].id,
            .name = calls[index].name,
            .content = std::move(content),
            .is_error = is_error,
        });
    }
    co_return messages;
}

} // namespace agent_smith::internal