#pragma once

#include <chrono>
#include <cstddef>
#include <atomic>
#include <functional>
#include <memory>
#include <optional>

#include "llm_client.hpp"
#include "memory.hpp"
#include "mcp_client.hpp"

namespace agent_smith
{

/// Shared cancellation and iteration accounting for an agent tree.
struct AgentTreeState
{
    /// Number of LLM iterations started across the tree.
    std::atomic<int> iterations{0};
    /// Cooperative cancellation shared by the tree.
    CancellationSource cancel;
};

/// Agent execution limits and default request behavior.
struct AgentConfig
{
    /// Maximum LLM requests per send.
    int max_iterations = 16;
    /// Wall-clock limit per send, checked between awaited operations.
    std::optional<std::chrono::seconds> deadline = std::nullopt;
    /// Cap on reported total token usage per send; missing usage is uncounted.
    std::optional<std::uint64_t> max_total_tokens = std::nullopt;
    /// Default per-call budget when a tool has no override.
    std::chrono::milliseconds default_tool_timeout{30'000};
    /// Maximum retained content bytes, excluding the truncation marker.
    std::size_t tool_result_max_bytes = 64 * 1024;
    /// Whether send uses streaming unless the request overrides it.
    bool stream_default = true;
    /// Maximum iterations across the shared tree state.
    int tree_max_iterations = 50;
    /// Shared tree state, or null to create a state owned by this agent.
    std::shared_ptr<AgentTreeState> tree_state;
};

/// Synchronous observers. Exceptions are logged and ignored; do not block.
struct AgentHooks
{
    /// Observe each assistant message appended to memory.
    std::function<void(const Message&)> on_message;
    /// Observe each requested tool call before dispatch.
    std::function<void(const ToolCall&)> on_tool_call;
    /// Observe each tool result after it is appended to memory.
    std::function<void(const ToolMessage&)> on_tool_result;
    /// Observe streamed events as they arrive.
    std::function<void(const StreamEvent&)> on_stream_event;
    /// Observe a framework error before send returns it.
    std::function<void(const mw::Error&)> on_error;
    /// Observe MCP notifications and connection transitions.
    mcp::EventHook on_mcp_event;
};

/// One user turn with optional sampling and streaming overrides.
struct AgentRequest
{
    /// User message to append before running the loop.
    UserMessage message;
    /// Per-turn LLM sampling overrides.
    SamplingParams sampling;
    /// Streaming override; absent uses AgentConfig::stream_default.
    std::optional<bool> stream;
};

/// Drives conversation, LLM requests, and native tool execution.
/// Only one coroutine may use an agent at a time. Dependencies must outlive it.
class Agent
{
public:
    /// Construct an agent; null memory creates a FullMemory instance.
    Agent(Runtime& runtime, LlmClient& client, std::shared_ptr<Memory> memory,
          AgentConfig config = {});
    /// Replace the base system prompt used for subsequent requests.
    void setSystemPrompt(std::string prompt);
    /// Register a native tool; duplicate or empty names return an error.
    mw::E<void> addTool(std::unique_ptr<Tool> tool, ToolOptions options = {});
    /// Bind prefixed MCP tools and resource access; retain the client.
    mw::E<void> addMcpClient(std::shared_ptr<mcp::Client> client);
    /// Connect and bind an MCP server, observing initialization events.
    mw::Task<mw::E<void>> connectMcp(mcp::ServerConfig config);
    /// Access synchronous observers for this agent.
    AgentHooks& hooks();
    /// Cancel this agent and its shared tree; tools observe linked tokens and
    /// streaming HTTP stops on the next chunk. Safe to call from any thread.
    void cancel();
    /// Append a user message and run until completion or a stop condition.
    mw::Task<mw::E<AssistantMessage>> send(std::string user);
    /// Run one turn with sampling and streaming overrides.
    mw::Task<mw::E<AssistantMessage>> send(AgentRequest request);
private:
    mw::Task<mw::E<AssistantMessage>> run(AgentRequest request);
    Runtime* runtime_;
    LlmClient* client_;
    std::shared_ptr<Memory> memory_;
    AgentConfig config_;
    std::shared_ptr<AgentHooks> hooks_ = std::make_shared<AgentHooks>();
    std::vector<std::shared_ptr<mcp::Client>> mcp_clients_;
    ToolRegistry tools_;
    SystemMessage system_;
    CancellationSource cancel_;
};

} // namespace agent_smith