#include <agent_smith/agent_smith.hpp>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mw/http_client_async_mock.hpp>
#include <span>
#include <stdexcept>

namespace agent_smith
{
namespace
{
using ::testing::_;
using ::testing::Invoke;

mw::HTTPResponse reply(const nlohmann::json& message)
{
    return {200, nlohmann::json({{"choices", {{{"message", message}}}}}).dump()};
}

nlohmann::json callMessage(std::string name = "echo")
{
    return {{"role", "assistant"}, {"content", ""},
            {"tool_calls", {{{"id", "a"}, {"type", "function"},
                {"function", {{"name", name}, {"arguments", "{}"}}}}}}};
}

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

class AgentTest : public ::testing::Test
{
protected:
    Runtime runtime;
    mw::HTTPSessionAsyncMock http;
    LlmClient client{runtime, http, {
        .endpoint = *mw::URL::fromStr("https://example.test"),
        .model = "test", .max_retries = 0}};
    std::shared_ptr<FullMemory> memory = std::make_shared<FullMemory>();
};

TEST_F(AgentTest, RunsToolLoopAndReusesHistory)
{
    Agent agent(runtime, client, memory, {.stream_default = false});
    agent.setSystemPrompt("system");
    ASSERT_TRUE(agent.addTool(makeTool(schema(),
        [](nlohmann::json, ToolContext) { return "tool output"; })));
    std::vector<std::string> events;
    agent.hooks().on_message = [&](const Message&) { events.push_back("message"); };
    agent.hooks().on_tool_call = [&](const ToolCall&) { events.push_back("call"); };
    agent.hooks().on_tool_result = [&](const ToolMessage&) { events.push_back("result"); };
    int requests = 0;
    EXPECT_CALL(http, post(_)).Times(3).WillRepeatedly(Invoke(
        [&](mw::HTTPRequest request)
        {
            auto body = nlohmann::json::parse(request.request_data);
            EXPECT_EQ(body["messages"][0]["content"], "system");
            EXPECT_EQ(body["tools"][0]["function"]["name"], "echo");
            ++requests;
            if(requests == 1)
            {
                EXPECT_EQ(body["temperature"], 0.25);
                return mw::HTTPSessionAsyncMock::complete(reply(callMessage()));
            }
            EXPECT_EQ(body["messages"][3]["role"], "tool");
            EXPECT_EQ(body["messages"][3]["tool_call_id"], "a");
            EXPECT_EQ(body["messages"][3]["content"], "tool output");
            return mw::HTTPSessionAsyncMock::complete(reply(
                {{"role", "assistant"}, {"content", "done"}}));
        }));
    auto result = syncWait(agent.send(AgentRequest{
        .message = {"hello"}, .sampling = {.temperature = 0.25}}));
    ASSERT_TRUE(result);
    EXPECT_EQ(result->content, "done");
    EXPECT_EQ(events, (std::vector<std::string>{"message", "call", "result", "message"}));
    EXPECT_EQ(memory->size(), 4U);
    EXPECT_TRUE(syncWait(agent.send("again")));
    EXPECT_EQ(memory->size(), 6U);
}

mw::Task<mw::E<mw::HTTPResponse>> streamReply(mw::ChunkCallback callback)
{
    const std::string data =
        "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},"
        "\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n";
    callback(std::as_bytes(std::span(data.data(), data.size())));
    co_return mw::HTTPResponse(200, "");
}

TEST_F(AgentTest, StreamsAndIgnoresThrowingObservers)
{
    Agent agent(runtime, client, memory, {.stream_default = false});
    int events = 0;
    agent.hooks().on_stream_event = [&](const StreamEvent&)
    { ++events; throw std::runtime_error("observer"); };
    agent.hooks().on_message = [](const Message&) { throw 1; };
    EXPECT_CALL(http, postStream(_, _)).WillOnce(Invoke(
        [](mw::HTTPRequest request, mw::ChunkCallback callback)
        {
            EXPECT_TRUE(nlohmann::json::parse(request.request_data)["stream"]);
            return streamReply(std::move(callback));
        }));
    auto result = syncWait(agent.send(AgentRequest{
        .message = {"hi"}, .stream = true}));
    ASSERT_TRUE(result);
    EXPECT_EQ(result->content, "hello");
    EXPECT_GT(events, 0);
}

TEST_F(AgentTest, ToolDoneStopsWithoutAnotherLlmRequest)
{
    Agent agent(runtime, client, memory, {.stream_default = false});
    ASSERT_TRUE(agent.addTool(makeTool(schema(),
        [](nlohmann::json, ToolContext)
        { return nlohmann::json{{"_done", true}, {"final_message", "finished"}}; })));
    EXPECT_CALL(http, post(_)).WillOnce(Invoke([](mw::HTTPRequest)
    { return mw::HTTPSessionAsyncMock::complete(reply(callMessage())); }));
    auto result = syncWait(agent.send("go"));
    ASSERT_TRUE(result);
    EXPECT_EQ(result->content, "finished");
    EXPECT_EQ(memory->size(), 4U);
}

TEST_F(AgentTest, UnknownToolIsConversationErrorAndIterationCapStopsLoop)
{
    Agent agent(runtime, client, memory,
                {.max_iterations = 2, .stream_default = false});
    int errors = 0;
    agent.hooks().on_error = [&](const mw::Error& error)
    { ++errors; EXPECT_NE(error.as<CancelledError>(), nullptr); throw 1; };
    int results = 0;
    agent.hooks().on_tool_result = [&](const ToolMessage& result)
    { ++results; EXPECT_TRUE(result.is_error); };
    EXPECT_CALL(http, post(_)).Times(2).WillRepeatedly(Invoke([](mw::HTTPRequest)
    { return mw::HTTPSessionAsyncMock::complete(reply(callMessage("missing"))); }));
    auto result = syncWait(agent.send("go"));
    EXPECT_FALSE(result);
    EXPECT_EQ(errors, 1);
    EXPECT_EQ(results, 2);
}

TEST_F(AgentTest, EnforcesDeadlineCancellationAndTreeCapBeforeRequest)
{
    EXPECT_CALL(http, post(_)).Times(0);
    EXPECT_CALL(http, postStream(_, _)).Times(0);
    Agent deadline(runtime, client, memory, {.deadline = std::chrono::seconds(0)});
    EXPECT_FALSE(syncWait(deadline.send("go")));
    Agent cancelled(runtime, client, memory);
    cancelled.cancel();
    EXPECT_FALSE(syncWait(cancelled.send("go")));
    Agent tree(runtime, client, memory, {.tree_max_iterations = 0});
    EXPECT_FALSE(syncWait(tree.send("go")));
}

TEST_F(AgentTest, TokenCapAndHttpFailureReturnErrors)
{
    Agent agent(runtime, client, memory,
                {.max_total_tokens = 1, .stream_default = false});
    EXPECT_CALL(http, post(_)).WillOnce(Invoke([](mw::HTTPRequest)
    {
        return mw::HTTPSessionAsyncMock::complete(mw::HTTPResponse(200,
            R"({"choices":[{"message":{"role":"assistant","content":"done"}}],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}})"));
    }));
    auto result = syncWait(agent.send("go"));
    ASSERT_FALSE(result);
    EXPECT_NE(result.error().as<CancelledError>(), nullptr);
    EXPECT_CALL(http, post(_)).WillOnce(Invoke([](mw::HTTPRequest)
    { return mw::HTTPSessionAsyncMock::complete(mw::HTTPResponse(400, "bad")); }));
    result = syncWait(agent.send("again"));
    ASSERT_FALSE(result);
    EXPECT_NE(result.error().as<LlmApiError>(), nullptr);
}

class BrokenMemory final : public Memory
{
public:
    void append(Message) override {}
    std::size_t size() const override { return 0; }
    void clear() override {}
    mw::Task<mw::E<std::vector<Message>>>
    snapshot(const SystemMessage&) override
    {
        co_return std::unexpected(error::parse("snapshot failed"));
    }
};

TEST_F(AgentTest, MemoryFailureReachesErrorHookWithoutHttp)
{
    Agent agent(runtime, client, std::make_shared<BrokenMemory>());
    int errors = 0;
    agent.hooks().on_error = [&](const mw::Error& error)
    { ++errors; EXPECT_NE(error.as<ParseError>(), nullptr); };
    EXPECT_CALL(http, postStream(_, _)).Times(0);
    auto result = syncWait(agent.send("go"));
    ASSERT_FALSE(result);
    EXPECT_EQ(result.error().msg(), "snapshot failed");
    EXPECT_EQ(errors, 1);
}

mw::Task<mw::E<mw::HTTPResponse>> streamTool(mw::ChunkCallback callback)
{
    const std::string data =
        R"(data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"a","type":"function","function":{"name":"echo","arguments":"{}"}}]},"finish_reason":"tool_calls"}]})"
        "\n\ndata: [DONE]\n\n";
    callback(std::as_bytes(std::span(data.data(), data.size())));
    co_return mw::HTTPResponse(200, "");
}

TEST_F(AgentTest, StreamingToolRoundTripWorksWithoutStreamHook)
{
    Agent agent(runtime, client, memory);
    ASSERT_TRUE(agent.addTool(makeTool(schema(),
        [](nlohmann::json, ToolContext) { return "output"; })));
    EXPECT_FALSE(agent.addTool(makeTool(schema(),
        [](nlohmann::json, ToolContext) { return "duplicate"; })));
    EXPECT_FALSE(agent.addTool(nullptr));
    ::testing::InSequence sequence;
    EXPECT_CALL(http, postStream(_, _)).WillOnce(Invoke(
        [](mw::HTTPRequest, mw::ChunkCallback callback)
        { return streamTool(std::move(callback)); }));
    EXPECT_CALL(http, postStream(_, _)).WillOnce(Invoke(
        [](mw::HTTPRequest request, mw::ChunkCallback callback)
        {
            const auto body = nlohmann::json::parse(request.request_data);
            EXPECT_EQ(body["messages"].back()["content"], "output");
            return streamReply(std::move(callback));
        }));
    auto result = syncWait(agent.send("go"));
    ASSERT_TRUE(result);
    EXPECT_EQ(result->content, "hello");
    EXPECT_EQ(memory->size(), 4U);
}

TEST_F(AgentTest, CancellationFromToolHookPreventsInvocation)
{
    Agent agent(runtime, client, memory, {.stream_default = false});
    bool invoked = false;
    ASSERT_TRUE(agent.addTool(makeTool(schema(),
        [&invoked](nlohmann::json, ToolContext)
        { invoked = true; return "output"; })));
    agent.hooks().on_tool_call = [&](const ToolCall&) { agent.cancel(); };
    EXPECT_CALL(http, post(_)).WillOnce(Invoke([](mw::HTTPRequest)
    { return mw::HTTPSessionAsyncMock::complete(reply(callMessage())); }));
    EXPECT_FALSE(syncWait(agent.send("go")));
    EXPECT_FALSE(invoked);
}

mw::Task<mw::E<mw::HTTPResponse>> cancellableStream(
    mw::ChunkCallback callback, Runtime& runtime)
{
    const std::string first =
        "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n";
    callback(std::as_bytes(std::span(first.data(), first.size())));
    co_await runtime.sleepFor(std::chrono::milliseconds(10));
    const bool accepted = callback(
        std::as_bytes(std::span(first.data(), first.size())));
    EXPECT_FALSE(accepted);
    co_return std::unexpected(mw::runtimeError("aborted"));
}

TEST_F(AgentTest, CancellationAbortsStreamOnNextChunk)
{
    Agent agent(runtime, client, memory);
    agent.hooks().on_stream_event = [&](const StreamEvent&) { agent.cancel(); };
    EXPECT_CALL(http, postStream(_, _)).WillOnce(Invoke(
        [&](mw::HTTPRequest, mw::ChunkCallback callback)
        { return cancellableStream(std::move(callback), runtime); }));
    auto result = syncWait(agent.send("go"));
    ASSERT_FALSE(result);
    EXPECT_NE(result.error().as<CancelledError>(), nullptr);
}

TEST_F(AgentTest, ToolObservesParentCancellation)
{
    auto tree = std::make_shared<AgentTreeState>();
    Agent agent(runtime, client, memory,
                {.stream_default = false, .tree_state = tree});
    bool observed = false;
    ASSERT_TRUE(agent.addTool(makeTool(schema(),
        [&](nlohmann::json, ToolContext context)
        {
            tree->cancel.cancel();
            observed = context.cancel.cancelled();
            return "cancelled";
        })));
    EXPECT_CALL(http, post(_)).WillOnce(Invoke([](mw::HTTPRequest)
    { return mw::HTTPSessionAsyncMock::complete(reply(callMessage())); }));
    auto result = syncWait(agent.send("go"));
    ASSERT_FALSE(result);
    EXPECT_TRUE(observed);
    EXPECT_NE(result.error().as<CancelledError>(), nullptr);
}

} // namespace
} // namespace agent_smith