#include <string>
#include <thread>
#include <utility>
#include <vector>

#include <agent_smith/sync_wait.hpp>
#include <agent_smith/tool.hpp>
#include <gtest/gtest.h>

namespace agent_smith
{
namespace
{

ToolSchema schema(std::string name)
{
    return ToolSchema{
        .name = std::move(name),
        .description = "test tool",
        .parameters =
            {
                {"type", "object"},
                {"properties", nlohmann::json::object()},
            },
    };
}

TEST(CancellationTest, TokensObserveSourceCancellation)
{
    CancellationSource source;
    const CancellationToken token = source.token();

    EXPECT_FALSE(token.cancelled());
    source.cancel();
    EXPECT_TRUE(token.cancelled());
}

TEST(ToolTest, SyncAdapterUsesRuntimeWorkerPool)
{
    Runtime runtime;
    const std::thread::id caller_thread = std::this_thread::get_id();
    std::thread::id invoked_thread;
    auto tool = makeTool(schema("thread_id"),
                         [&invoked_thread](nlohmann::json, ToolContext)
                         {
                             invoked_thread = std::this_thread::get_id();
                             return "completed";
                         });

    auto result = syncWait(tool->invoke(nlohmann::json::object(),
                                        ToolContext{.runtime = &runtime}));

    ASSERT_TRUE(result.has_value());
    EXPECT_EQ(result->content, "completed");
    EXPECT_NE(invoked_thread, caller_thread);
}

mw::Task<mw::E<ToolResult>> asyncThreadTool(std::thread::id& invoked_thread,
                                            nlohmann::json, ToolContext)
{
    invoked_thread = std::this_thread::get_id();
    co_return ToolResult{
        .content = "completed",
    };
}

TEST(ToolTest, AsyncAdapterAwaitsCallableDirectly)
{
    Runtime runtime;
    const std::thread::id caller_thread = std::this_thread::get_id();
    std::thread::id invoked_thread;
    auto tool = makeTool(
        schema("async_thread_id"),
        [&invoked_thread](nlohmann::json arguments, ToolContext context)
        {
            return asyncThreadTool(invoked_thread, std::move(arguments),
                                   std::move(context));
        });

    auto result = syncWait(tool->invoke(nlohmann::json::object(),
                                        ToolContext{.runtime = &runtime}));

    ASSERT_TRUE(result.has_value());
    EXPECT_EQ(result->content, "completed");
    EXPECT_EQ(invoked_thread, caller_thread);
}

TEST(ToolTest, NormalizesStructuredAndFallibleResults)
{
    auto structured = normalizeToolResult(nlohmann::json{{"answer", 42}});
    ASSERT_TRUE(structured.has_value());
    EXPECT_EQ(structured->content, "{\"answer\":42}");
    EXPECT_EQ(structured->structured, nlohmann::json({{"answer", 42}}));

    const mw::E<ToolResult> failed =
        std::unexpected(mw::runtimeError("tool failed"));
    auto forwarded = normalizeToolResult(failed);
    ASSERT_FALSE(forwarded.has_value());
    EXPECT_EQ(forwarded.error().msg(), "tool failed");
}

TEST(ToolRegistryTest, PreservesInsertionOrderAndReplacesByName)
{
    ToolRegistry registry;
    EXPECT_TRUE(registry.add(makeTool(
        schema("first"), [](nlohmann::json, ToolContext) { return "first"; })));
    EXPECT_TRUE(
        registry.add(makeTool(schema("second"), [](nlohmann::json, ToolContext)
                              { return "second"; }),
                     ToolOptions{.timeout = std::chrono::seconds(3)}));
    EXPECT_FALSE(
        registry.add(makeTool(schema("first"), [](nlohmann::json, ToolContext)
                              { return "collision"; })));

    registry.replace(makeTool(schema("first"), [](nlohmann::json, ToolContext)
                              { return "replaced"; }));

    EXPECT_EQ(registry.size(), 2U);
    ASSERT_NE(registry.find("first"), nullptr);
    const auto schemas = registry.schemas();
    ASSERT_EQ(schemas.size(), 2U);
    EXPECT_EQ(schemas[0].name, "first");
    EXPECT_EQ(schemas[1].name, "second");
    ASSERT_TRUE(registry.options("second").timeout.has_value());
    EXPECT_EQ(*registry.options("second").timeout, std::chrono::seconds(3));
}

} // namespace
} // namespace agent_smith