// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors

// --8<-- [start:openai_embeddings]
use std::{iter::once, sync::Arc};

use arrow_array::{record_batch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use futures::StreamExt;
use lancedb::{
    connect,
    embeddings::{openai::OpenAIEmbeddingFunction, EmbeddingDefinition, EmbeddingFunction},
    query::{ExecutableQuery, QueryBase},
    Result,
};

#[tokio::main]
async fn main() -> Result<()> {
    let db = connect("./mydb").execute().await?;
    let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not set");
    let embedding = Arc::new(OpenAIEmbeddingFunction::new_with_model(
        api_key,
        "text-embedding-3-large",
    )?);

    db.embedding_registry().register("openai", embedding.clone())?;

    let schema = Arc::new(Schema::new(vec![Field::new("text", DataType::Utf8, false)]));
    let table = db
        .create_empty_table("mytable", schema)
        .add_embedding(EmbeddingDefinition::new("text", "openai", Some("vector")))?
        .execute()
        .await?;

    table
        .add(record_batch!(("text", Utf8, ["This is a test.", "Another example."]))?)
        .execute()
        .await?;

    let query = Arc::new(StringArray::from_iter_values(once("test example")));
    let query_vector = embedding.compute_query_embeddings(query)?;
    let mut results = table.vector_search(query_vector)?.limit(5).execute().await?;

    while let Some(batch) = results.next().await {
        println!("{:?}", batch?);
    }

    Ok(())
}
// --8<-- [end:openai_embeddings]

// --8<-- [start:create_embedding_function]
use std::sync::Arc;

use lancedb::embeddings::openai::OpenAIEmbeddingFunction;

let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not set");
let embedding = Arc::new(
    OpenAIEmbeddingFunction::new_with_model(api_key, "text-embedding-3-small")
        .expect("failed to create OpenAI embedding function"),
);
// --8<-- [end:create_embedding_function]

// --8<-- [start:manual_query_embeddings]
use std::{iter::once, sync::Arc};

use arrow_array::{record_batch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use futures::StreamExt;
use lancedb::{
    connect,
    embeddings::{openai::OpenAIEmbeddingFunction, EmbeddingDefinition, EmbeddingFunction},
    query::{ExecutableQuery, QueryBase},
    Result,
};

#[tokio::main]
async fn main() -> Result<()> {
    let db = connect("./mydb").execute().await?;
    let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not set");
    let embedding = Arc::new(OpenAIEmbeddingFunction::new_with_model(
        api_key,
        "text-embedding-3-large",
    )?);
    db.embedding_registry().register("openai", embedding.clone())?;

    let schema = Arc::new(Schema::new(vec![Field::new("text", DataType::Utf8, false)]));
    let table = db
        .create_empty_table("mytable", schema)
        .add_embedding(EmbeddingDefinition::new("text", "openai", Some("vector")))?
        .execute()
        .await?;

    table
        .add(record_batch!(("text", Utf8, ["This is a test.", "Another example."]))?)
        .execute()
        .await?;

    // Manually generate embeddings for the query (Enterprise path)
    let query = Arc::new(StringArray::from_iter_values(once("test example")));
    let query_vector = embedding.compute_query_embeddings(query)?;
    // --8<-- [start:manual_query_search]
    // query_vector is assumed to already be generated by your embedding function
    let mut results = table.vector_search(query_vector)?.limit(5).execute().await?;

    while let Some(batch) = results.next().await {
        println!("{:?}", batch?);
    }
    // --8<-- [end:manual_query_search]

    Ok(())
}
// --8<-- [end:manual_query_embeddings]

// --8<-- [start:embedding_function]
use std::{borrow::Cow, sync::Arc};

use arrow_array::{Array, FixedSizeListArray, Float32Array};
use arrow_schema::{DataType, Field, Schema};
use lancedb::{
    connect,
    embeddings::{EmbeddingDefinition, EmbeddingFunction},
    Result,
};

#[derive(Debug, Clone)]
struct MyTextEmbedder {
    dim: usize,
}

impl EmbeddingFunction for MyTextEmbedder {
    fn name(&self) -> &str {
        "my-embedder"
    }

    fn source_type(&self) -> Result<Cow<'_, DataType>> {
        Ok(Cow::Owned(DataType::Utf8))
    }

    fn dest_type(&self) -> Result<Cow<'_, DataType>> {
        Ok(Cow::Owned(DataType::new_fixed_size_list(
            DataType::Float32,
            self.dim as i32,
            true,
        )))
    }

    fn compute_source_embeddings(&self, source: Arc<dyn Array>) -> Result<Arc<dyn Array>> {
        let values = Arc::new(Float32Array::from(vec![1.0f32; source.len() * self.dim]));
        let field = Arc::new(Field::new("item", DataType::Float32, true));
        Ok(Arc::new(FixedSizeListArray::new(
            field,
            self.dim as i32,
            values,
            None,
        )))
    }

    fn compute_query_embeddings(&self, _input: Arc<dyn Array>) -> Result<Arc<dyn Array>> {
        unimplemented!()
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let db = connect("./mydb").execute().await?;
    db.embedding_registry()
        .register("my-embedder", Arc::new(MyTextEmbedder { dim: 3 }))?;

    let schema = Arc::new(Schema::new(vec![Field::new("text", DataType::Utf8, false)]));
    db.create_empty_table("mytable", schema)
        .add_embedding(EmbeddingDefinition::new(
            "text",
            "my-embedder",
            Some("vector"),
        ))?
        .execute()
        .await?;

    Ok(())
}
// --8<-- [end:embedding_function]