import com.fasterxml.jackson.databind.ObjectMapper;
import com.openai.client.OpenAIClient;
import com.openai.client.okhttp.OpenAIOkHttpClient;
import com.openai.models.ChatModel;
import com.openai.models.chat.completions.ChatCompletion;
import com.openai.models.chat.completions.ChatCompletionCreateParams;
import com.openai.models.conversations.Message;
import com.openai.models.embeddings.Embedding;
import com.openai.models.embeddings.EmbeddingCreateParams;
import com.openai.models.embeddings.EmbeddingModel;
import com.openai.models.responses.*;

import java.sql.*;
import java.util.List;

public class Main {
    static OpenAIClient client = null;
    static Connection conn = null;
    public static void main(String[] args) throws SQLException , Exception {
        System.out.println("Start");

        //Create connection to MariaDB. Before this, you need to create the "db12" database using mariadb console
        conn = DriverManager.getConnection("jdbc:mariadb://localhost/db12", "root", null);

        //create table where we store documents with embeddings, if it does not already exist
        Statement stmt = conn.createStatement();
        stmt.executeUpdate("CREATE TABLE IF NOT Exists documents ("
                +"id BIGINT AUTO_INCREMENT PRIMARY KEY,"
                +"content TEXT NOT NULL,"
                +"embedding VECTOR(1536) NOT NULL,"
                +"VECTOR INDEX (embedding)"
        +")");

        System.out.println("Table created");


        //create a connection to OpenAI API
        String api_key = "place the key here";
        client = OpenAIOkHttpClient.builder()
                .apiKey(api_key)
                .build();

        /*
        //Demo how to get AI chat response
        ResponseCreateParams params = ResponseCreateParams.builder()
                .model(ChatModel.GPT_5_1_CHAT_LATEST)
                .input("You are a helpful teacher. What is cosine similarity?").build();

        System.out.println("AI working");
        Response response = client.responses().create(params);

        for (ResponseOutputItem item : response.output()) {
            if (item.message().isEmpty()) continue;
            ResponseOutputMessage msg = item.message().get();

            for (com.openai.models.responses.ResponseOutputMessage.Content cnt : msg.content()) {
                if (cnt.outputText().isEmpty()) continue;
                System.out.println( cnt.asOutputText().text());

            }
        }
        */


        //Ask OpenAI to create embedding for text in the variable "text"
        String text = "Vector databases store embeddings for semantic search";
        EmbeddingCreateParams eparams = EmbeddingCreateParams.builder()
                .model(EmbeddingModel.TEXT_EMBEDDING_3_SMALL)
                .input(text)
                .build();
        List<Embedding> embeddings = client.embeddings().create(eparams).data();
        List<Float> vector = embeddings.getFirst().embedding();

        //Encode the returned vector as json so that we can store it into the MariaDb
        ObjectMapper mapper = new ObjectMapper();
        String embeddingJson = mapper.writeValueAsString(vector);

        //Store the vector with text to db
        String sql = "INSERT INTO documents (content, embedding) VALUES (?, VEC_FromText(?))";
        PreparedStatement ps = conn.prepareStatement(sql);
        ps.setString(1, text);
        ps.setString(2, embeddingJson);
        ps.executeUpdate();





    }

    static void vectorSearch(String searchQuery) throws Exception {
        //First, we create embedding for a query
        EmbeddingCreateParams eparams = EmbeddingCreateParams.builder()
                .model(EmbeddingModel.TEXT_EMBEDDING_3_SMALL)
                .input(searchQuery)
                .build();
        List<Embedding> embeddings = client.embeddings().create(eparams).data();
        List<Float> vector = embeddings.getFirst().embedding();

        //Encode the returned vector as json so that we can store it into the MariaDb
        ObjectMapper mapper = new ObjectMapper();
        String embeddingJson = mapper.writeValueAsString(vector);

        //Then we search for nearest embeddings
        String sql = "SELECT d.content FROM documents AS d ORDER BY VEC_DISTANCE_EUCLIDEAN(d.embedding,VEC_FromText(?)) LIMIT 10";
        PreparedStatement ps = conn.prepareStatement(sql);
        ps.setString(1, embeddingJson);
        ResultSet rs = ps.executeQuery();
        while (rs.next()) {
            System.out.println(rs.getString(1));
        }
    }
}