#pragma once

#include <chrono>
#include <functional>
#include <optional>
#include <string>
#include <vector>

#include <mw/http_client_async.hpp>
#include <mw/url.hpp>
#include <nlohmann/json.hpp>

#include "message.hpp"
#include "runtime.hpp"
#include "stream_event.hpp"
#include "tool.hpp"

namespace agent_smith
{

/// Sampling controls sent to an OpenAI-compatible chat-completions endpoint.
struct SamplingParams
{
    /// Sampling temperature, when the provider should receive one.
    std::optional<double> temperature;

    /// Nucleus-sampling threshold, when the provider should receive one.
    std::optional<double> top_p;

    /// Maximum number of tokens the provider may generate.
    std::optional<int> max_tokens;

    /// Sequences that cause generation to stop.
    std::vector<std::string> stop;

    /// Opaque structured-output or JSON-mode configuration.
    nlohmann::json response_format = nullptr;

    /// Provider-specific request fields outside the standard API surface.
    nlohmann::json extra = nullptr;
};

/// Connection and request defaults for an LLM client.
struct LlmClientConfig
{
    /// Base URL of the OpenAI-compatible endpoint.
    mw::URL endpoint;

    /// Provider model identifier.
    std::string model;

    /// Optional bearer token for the endpoint.
    std::optional<std::string> api_key;

    /// Requested maximum duration of one HTTP transfer.
    std::chrono::seconds request_timeout{60};

    /// Number of retry attempts after the initial request.
    int max_retries = 3;

    /// Sampling values used when a request does not override them.
    SamplingParams default_sampling;
};

/// One non-streaming LLM completion request.
struct LlmRequest
{
    /// Ordered conversation history sent to the model.
    std::vector<Message> messages;

    /// Function schemas made available to the model for this request.
    std::vector<ToolSchema> tools;

    /// Sampling values that override the client's defaults.
    SamplingParams sampling;
    /// Cooperative cancellation checked before requests and streamed chunks.
    CancellationToken cancel;
};

/// The final assistant message returned by an LLM completion request.
struct LlmResponse
{
    /// Parsed assistant response, including usage when supplied by the server.
    AssistantMessage message;
};

/// Calls an OpenAI-compatible chat-completions endpoint.
class LlmClient
{
public:
    /// Construct a client that uses the runtime's shared HTTP session.
    LlmClient(Runtime& runtime, LlmClientConfig config);

    /// Construct a client with an HTTP session override, primarily for tests.
    LlmClient(Runtime& runtime, mw::HTTPSessionAsyncInterface& http,
              LlmClientConfig config);

    /// Send a non-streaming request to `/v1/chat/completions`.
    mw::Task<mw::E<LlmResponse>> complete(LlmRequest request);

    /// Send a streaming request and synchronously deliver typed stream events.
    mw::Task<mw::E<LlmResponse>>
    completeStream(LlmRequest request,
                   std::function<void(const StreamEvent&)> on_event);

private:
    Runtime* runtime_;
    mw::HTTPSessionAsyncInterface* http_;
    LlmClientConfig config_;
};

} // namespace agent_smith