#include <bitset>
#include <cstdint>
#include <iostream>
#include <string>
#include <vector>

#include <cpr/cpr.h>
#include <nlohmann/json.hpp>
#include <pgvector/pqxx.hpp>
#include <pqxx/pqxx>

using json = nlohmann::json;

// https://docs.cohere.com/reference/embed
std::vector<std::string> embed(
    const std::vector<std::string>& texts,
    const std::string& input_type,
    char* api_key
) {
    std::string url{"https://api.cohere.com/v2/embed"};
    json data{
        {"texts", texts},
        {"model", "embed-v4.0"},
        {"input_type", input_type},
        {"embedding_types", {"ubinary"}}
    };

    cpr::Response r = cpr::Post(
        cpr::Url{url},
        cpr::Body{data.dump()},
        cpr::Bearer{api_key},
        cpr::Header{{"Content-Type", "application/json"}}
    );
    if (r.status_code != 200) {
        throw std::runtime_error{"Bad status: " + std::to_string(r.status_code)};
    }
    json response = json::parse(r.text);

    std::vector<std::string> embeddings;
    for (const auto& v : response["embeddings"]["ubinary"]) {
        std::stringstream buf;
        for (uint8_t c : v) {
            std::bitset<8> b{c};
            buf << b.to_string();
        }
        embeddings.emplace_back(buf.str());
    }
    return embeddings;
}

int main() {
    char* api_key = std::getenv("CO_API_KEY");
    if (!api_key) {
        std::cout << "Set CO_API_KEY" << std::endl;
        return 1;
    }

    pqxx::connection conn{"dbname=pgvector_example"};

    pqxx::nontransaction tx{conn};
    tx.exec("CREATE EXTENSION IF NOT EXISTS vector");
    tx.exec("DROP TABLE IF EXISTS documents");
    tx.exec("CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1536))");

    std::vector<std::string> input{
        "The dog is barking", "The cat is purring", "The bear is growling"
    };
    std::vector<std::string> embeddings = embed(input, "search_document", api_key);
    for (size_t i = 0; i < input.size(); i++) {
        tx.exec(
            "INSERT INTO documents (content, embedding) VALUES ($1, $2)",
            pqxx::params{input[i], embeddings[i]}
        );
    }

    std::string query{"forest"};
    std::string query_embedding = embed({query}, "search_query", api_key)[0];
    pqxx::result result = tx.exec(
        "SELECT content FROM documents ORDER BY embedding <~> $1 LIMIT 5",
        pqxx::params{query_embedding}
    );
    for (const auto& row : result) {
        std::cout << row[0].as<std::string>() << std::endl;
    }

    return 0;
}