#include <gtest/gtest.h>
#include "embeddings.hpp"
#include "embedding_cache.hpp"
#include "errors.hpp"
#include <httplib.h>
#include <atomic>
#include <chrono>
#include <thread>
#include <vector>
using namespace svapi;

TEST(Embeddings, ParsesWellFormedResponse) {
    auto v = parseEmbeddingResponse(nlohmann::json::parse(R"({"data":[{"embedding":[0.5,1.0,1.5]}]})"));
    ASSERT_EQ(v.size(), 3u); EXPECT_FLOAT_EQ(v[0], 0.5f); EXPECT_FLOAT_EQ(v[2], 1.5f);
}
TEST(Embeddings, MalformedThrows) {
    EXPECT_THROW(parseEmbeddingResponse(nlohmann::json::parse(R"({"data":[]})")), ApiError);
    EXPECT_THROW(parseEmbeddingResponse(nlohmann::json::object()), ApiError);
}
TEST(Embeddings, SplitBaseHandlesPathPrefix) {
    EXPECT_EQ(splitEmbeddingBase("https://api.openai.com").origin, "https://api.openai.com");
    EXPECT_EQ(splitEmbeddingBase("https://api.openai.com").path, "");
    EXPECT_EQ(splitEmbeddingBase("https://api.openai.com/").path, "");
    auto orr = splitEmbeddingBase("https://openrouter.ai/api");
    EXPECT_EQ(orr.origin, "https://openrouter.ai");
    EXPECT_EQ(orr.path, "/api");
    EXPECT_EQ(splitEmbeddingBase("https://openrouter.ai/api/").path, "/api");
    EXPECT_EQ(splitEmbeddingBase("http://h:8080/a/b").origin, "http://h:8080");
    EXPECT_EQ(splitEmbeddingBase("http://h:8080/a/b").path, "/a/b");
}
TEST(Embeddings, ClientHitsMockServer) {
    httplib::Server mock;
    mock.Post("/v1/embeddings", [](const httplib::Request& req, httplib::Response& res) {
        auto j = nlohmann::json::parse(req.body);
        EXPECT_EQ(j["model"], "m"); EXPECT_EQ(j["input"], "hello");
        res.set_content(R"({"data":[{"embedding":[1,2,3,4]}]})", "application/json");
    });
    int port = mock.bind_to_any_port("127.0.0.1");
    std::thread t([&]{ mock.listen_after_bind(); });
    EmbeddingClient cli("http://127.0.0.1:" + std::to_string(port), "k");
    auto v = cli.embed("m", "hello");
    EXPECT_EQ(v.size(), 4u); EXPECT_FLOAT_EQ(v[3], 4.0f);
    mock.stop(); t.join();
}
TEST(Embeddings, ClientHonorsBasePathPrefix) {
    // Provider whose embeddings live under a path prefix (e.g. OpenRouter /api/v1/embeddings)
    httplib::Server mock;
    mock.Post("/api/v1/embeddings", [](const httplib::Request&, httplib::Response& res) {
        res.set_content(R"({"data":[{"embedding":[9,9]}]})", "application/json");
    });
    int port = mock.bind_to_any_port("127.0.0.1");
    std::thread t([&]{ mock.listen_after_bind(); });
    EmbeddingClient cli("http://127.0.0.1:" + std::to_string(port) + "/api", "k");
    auto v = cli.embed("qwen/qwen3-embedding-8b", "hello");
    EXPECT_EQ(v.size(), 2u); EXPECT_FLOAT_EQ(v[0], 9.0f);
    mock.stop(); t.join();
}

TEST(Embeddings, PoolRunsRequestsConcurrently) {
    // A serialized impl (single Client + mutex) would pin maxInFlight at 1.
    // A real pool of N clients lets up to N requests overlap on independent sockets.
    std::atomic<int> inFlight{0};
    std::atomic<int> maxInFlight{0};
    httplib::Server mock;
    mock.Post("/v1/embeddings", [&](const httplib::Request&, httplib::Response& res) {
        int now = ++inFlight;
        int prev = maxInFlight.load();
        while (now > prev && !maxInFlight.compare_exchange_weak(prev, now)) {}
        std::this_thread::sleep_for(std::chrono::milliseconds(180));
        --inFlight;
        res.set_content(R"({"data":[{"embedding":[1,2,3,4]}]})", "application/json");
    });
    int port = mock.bind_to_any_port("127.0.0.1");
    mock.set_keep_alive_max_count(1000);
    std::thread srv([&]{ mock.listen_after_bind(); });
    while (!mock.is_running()) std::this_thread::sleep_for(std::chrono::milliseconds(5));

    const std::string origin = "http://127.0.0.1:" + std::to_string(port);
    constexpr int K = 4;
    std::vector<std::thread> workers;
    std::atomic<int> ok{0};
    for (int i = 0; i < K; ++i) {
        workers.emplace_back([&]{
            EmbeddingClient cli(origin, "", 10, 60, /*poolSize=*/4);
            auto v = cli.embed("m", "t");
            if (v.size() == 4u && v[3] == 4.0f) ++ok;
        });
    }
    for (auto& w : workers) w.join();
    mock.stop(); srv.join();

    EXPECT_EQ(ok.load(), K);                 // all requests succeeded
    EXPECT_GE(maxInFlight.load(), 2);        // proves overlap — pool did NOT serialize
}

TEST(Embeddings, PoolSizeOneSerializes) {
    // poolSize=1 means a single keep-alive client → requests cannot overlap.
    std::atomic<int> inFlight{0};
    std::atomic<int> maxInFlight{0};
    httplib::Server mock;
    mock.Post("/v1/embeddings", [&](const httplib::Request&, httplib::Response& res) {
        int now = ++inFlight;
        int prev = maxInFlight.load();
        while (now > prev && !maxInFlight.compare_exchange_weak(prev, now)) {}
        std::this_thread::sleep_for(std::chrono::milliseconds(120));
        --inFlight;
        res.set_content(R"({"data":[{"embedding":[1,2,3,4]}]})", "application/json");
    });
    int port = mock.bind_to_any_port("127.0.0.1");
    mock.set_keep_alive_max_count(1000);
    std::thread srv([&]{ mock.listen_after_bind(); });
    while (!mock.is_running()) std::this_thread::sleep_for(std::chrono::milliseconds(5));

    const std::string origin = "http://127.0.0.1:" + std::to_string(port);
    constexpr int K = 4;
    std::vector<std::thread> workers;
    std::atomic<int> ok{0};
    for (int i = 0; i < K; ++i) {
        workers.emplace_back([&]{
            EmbeddingClient cli(origin, "", 10, 60, /*poolSize=*/1);
            auto v = cli.embed("m", "t");
            if (v.size() == 4u) ++ok;
        });
    }
    for (auto& w : workers) w.join();
    mock.stop(); srv.join();

    EXPECT_EQ(ok.load(), K);
    EXPECT_EQ(maxInFlight.load(), 1);        // serialized by design (single client)
}

// ---- EmbeddingCache tests -----------------------------------------------

TEST(EmbeddingCache, ReturnsCachedVectorOnExactKey) {
    EmbeddingCache cache(16, 1024*1024, 0);
    EmbeddingCache::Key k{"https://api.openai.com", "text-embedding-3-small", "hello world"};
    std::vector<float> vec = {0.1f, 0.2f, 0.3f};
    cache.put(k, vec);
    auto result = cache.get(k);
    ASSERT_TRUE(result.has_value());
    ASSERT_EQ(result->size(), 3u);
    EXPECT_FLOAT_EQ((*result)[0], 0.1f);
    EXPECT_FLOAT_EQ((*result)[1], 0.2f);
    EXPECT_FLOAT_EQ((*result)[2], 0.3f);
}

TEST(EmbeddingCache, MissOnAbsentKey) {
    EmbeddingCache cache(16, 1024*1024, 0);
    EmbeddingCache::Key k{"https://api.openai.com", "text-embedding-3-small", "not stored"};
    auto result = cache.get(k);
    EXPECT_FALSE(result.has_value());
    auto s = cache.stats();
    EXPECT_EQ(s.misses, 1u);
    EXPECT_EQ(s.hits, 0u);
}

TEST(EmbeddingCache, EvictsLeastRecentlyUsedWhenFull) {
    // capacity=3: insert A,B,C; then get(A) to refresh it; insert D → B should be evicted
    EmbeddingCache cache(3, 1024*1024, 0);
    EmbeddingCache::Key kA{"base", "m", "A"};
    EmbeddingCache::Key kB{"base", "m", "B"};
    EmbeddingCache::Key kC{"base", "m", "C"};
    EmbeddingCache::Key kD{"base", "m", "D"};
    cache.put(kA, {1.0f});
    cache.put(kB, {2.0f});
    cache.put(kC, {3.0f});
    // Touch A so it's recently used; B becomes the LRU
    cache.get(kA);
    // Insert D — evicts B (least recently used)
    cache.put(kD, {4.0f});
    EXPECT_TRUE(cache.get(kA).has_value());
    EXPECT_FALSE(cache.get(kB).has_value()); // evicted
    EXPECT_TRUE(cache.get(kC).has_value());
    EXPECT_TRUE(cache.get(kD).has_value());
}

TEST(EmbeddingCache, DistinguishesModels) {
    EmbeddingCache cache(16, 1024*1024, 0);
    EmbeddingCache::Key k1{"https://api.openai.com", "model-A", "hello"};
    EmbeddingCache::Key k2{"https://api.openai.com", "model-B", "hello"};
    cache.put(k1, {1.0f});
    cache.put(k2, {2.0f});
    auto r1 = cache.get(k1);
    auto r2 = cache.get(k2);
    ASSERT_TRUE(r1.has_value());
    ASSERT_TRUE(r2.has_value());
    EXPECT_FLOAT_EQ((*r1)[0], 1.0f);
    EXPECT_FLOAT_EQ((*r2)[0], 2.0f);
    EXPECT_EQ(cache.stats().size, 2u);
}

TEST(EmbeddingCache, DistinguishesApiBase) {
    EmbeddingCache cache(16, 1024*1024, 0);
    EmbeddingCache::Key k1{"https://api.openai.com", "text-embedding-3-small", "hello"};
    EmbeddingCache::Key k2{"https://openrouter.ai/api", "text-embedding-3-small", "hello"};
    cache.put(k1, {1.0f});
    cache.put(k2, {2.0f});
    auto r1 = cache.get(k1);
    auto r2 = cache.get(k2);
    ASSERT_TRUE(r1.has_value());
    ASSERT_TRUE(r2.has_value());
    EXPECT_FLOAT_EQ((*r1)[0], 1.0f);
    EXPECT_FLOAT_EQ((*r2)[0], 2.0f);
    EXPECT_EQ(cache.stats().size, 2u);
}

TEST(EmbeddingCache, ExpiresAfterTtl) {
    EmbeddingCache cache(16, 1024*1024, /*ttlSec=*/60);
    int64_t fakeNow = 1000;
    cache.setClockForTesting([&fakeNow]{ return fakeNow; });
    EmbeddingCache::Key k{"base", "m", "text"};
    cache.put(k, {5.0f});
    // Before expiry — hit
    EXPECT_TRUE(cache.get(k).has_value());
    // Advance clock past TTL
    fakeNow = 1061;
    auto result = cache.get(k);
    EXPECT_FALSE(result.has_value()); // expired
    auto s = cache.stats();
    EXPECT_GE(s.misses, 1u); // at least the expired-get counted as miss
}

TEST(EmbeddingCache, EvictsToStayUnderByteBudget) {
    // Each vector of 4 floats = 16 bytes. Budget = 40 bytes → max 2 full entries.
    EmbeddingCache cache(/*capacity=*/100, /*maxBytes=*/40, /*ttlSec=*/0);
    EmbeddingCache::Key kA{"b", "m", "A"};
    EmbeddingCache::Key kB{"b", "m", "B"};
    EmbeddingCache::Key kC{"b", "m", "C"};
    cache.put(kA, {1.0f, 2.0f, 3.0f, 4.0f}); // 16 bytes
    cache.put(kB, {5.0f, 6.0f, 7.0f, 8.0f}); // 16 bytes; total=32
    // inserting C (16 bytes) would push total to 48 > 40 → evict LRU (A)
    cache.put(kC, {9.0f, 10.0f, 11.0f, 12.0f});
    EXPECT_LE(cache.stats().bytes, 40u);
    EXPECT_FALSE(cache.get(kA).has_value()); // evicted
    EXPECT_TRUE(cache.get(kB).has_value());
    EXPECT_TRUE(cache.get(kC).has_value());
}

TEST(EmbeddingCache, NormalizationCollapsesWhitespaceAndCaseWhenEnabled) {
    // With normalize=true: "Max  DC" and "max dc" should map to ONE entry
    EmbeddingCache cacheOn(16, 1024*1024, 0, /*normalize=*/true);
    EmbeddingCache::Key k1{"base", "m", "Max  DC"};
    EmbeddingCache::Key k2{"base", "m", "max dc"};
    cacheOn.put(k1, {1.0f});
    auto r = cacheOn.get(k2);
    ASSERT_TRUE(r.has_value());
    EXPECT_FLOAT_EQ((*r)[0], 1.0f);
    EXPECT_EQ(cacheOn.stats().size, 1u);

    // With normalize=false (default): they stay distinct
    EmbeddingCache cacheOff(16, 1024*1024, 0, /*normalize=*/false);
    EmbeddingCache::Key k3{"base", "m", "Max  DC"};
    EmbeddingCache::Key k4{"base", "m", "max dc"};
    cacheOff.put(k3, {1.0f});
    auto r2 = cacheOff.get(k4);
    EXPECT_FALSE(r2.has_value());
    EXPECT_EQ(cacheOff.stats().size, 1u);
}

TEST(EmbeddingCache, StatsCountsHitsAndMisses) {
    EmbeddingCache cache(16, 1024*1024, 0);
    EmbeddingCache::Key k{"base", "m", "text"};
    // 2 misses
    cache.get(k);
    cache.get(k);
    // put then 3 hits
    cache.put(k, {1.0f});
    cache.get(k);
    cache.get(k);
    cache.get(k);
    auto s = cache.stats();
    EXPECT_EQ(s.misses, 2u);
    EXPECT_EQ(s.hits, 3u);
    EXPECT_EQ(s.size, 1u);
    EXPECT_EQ(s.capacity, 16u);
    EXPECT_EQ(s.bytes, sizeof(float) * 1u);
}

TEST(EmbeddingCache, ReconfigureEvictsToTighterBounds) {
    EmbeddingCache cache(16, 1024*1024, 0);
    EmbeddingCache::Key kA{"b", "m", "A"};
    EmbeddingCache::Key kB{"b", "m", "B"};
    EmbeddingCache::Key kC{"b", "m", "C"};
    cache.put(kA, {1.0f});   // A is least-recently used
    cache.put(kB, {2.0f});
    cache.put(kC, {3.0f});
    ASSERT_EQ(cache.stats().size, 3u);
    // Shrink capacity to 1 — reconfigure must evict immediately, keeping the MRU entry.
    cache.reconfigure(1, 1024*1024, 0, false);
    auto s = cache.stats();
    EXPECT_EQ(s.size, 1u);
    EXPECT_EQ(s.capacity, 1u);
    EXPECT_FALSE(cache.get(kA).has_value());  // evicted
    EXPECT_TRUE(cache.get(kC).has_value());   // most-recently used survives
}