package ankurm.ai21;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;

import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.vectorstore.EmbeddedDocument;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;

/**
 * VectorStore.upsert(List<EmbeddedDocument>) in Spring AI 2.1.0-M1: what the contract says, which stores
 * refuse it, and why vectors you already hold never need to be embedded twice.
 */
public class UpsertDemo {

	/** Deterministic 4-d "embedding" that counts how often it is asked to embed something. */
	static final class CountingModel implements EmbeddingModel {
		final AtomicInteger texts = new AtomicInteger();

		float[] vector(String text) {
			texts.incrementAndGet();
			int h = text.hashCode();
			return new float[] { (h & 0xff) / 255f, ((h >> 8) & 0xff) / 255f, ((h >> 16) & 0xff) / 255f, ((h >> 24) & 0xff) / 255f };
		}

		@Override
		public EmbeddingResponse call(EmbeddingRequest request) {
			List<Embedding> out = new ArrayList<>();
			int i = 0;
			for (String t : request.getInstructions()) {
				out.add(new Embedding(vector(t), i++));
			}
			return new EmbeddingResponse(out);
		}

		@Override
		public float[] embed(Document document) {
			return vector(document.getText());
		}

		@Override
		public int dimensions() {
			return 4;
		}
	}

	/** A teaching store that implements upsert by id. Real stores do this with ON CONFLICT, HSET, PUT or bulk index. */
	static final class ToyStore implements VectorStore {
		final Map<String, EmbeddedDocument> rows = new LinkedHashMap<>();
		final CountingModel model;

		ToyStore(CountingModel model) {
			this.model = model;
		}

		@Override
		public void add(List<Document> documents) {
			upsert(documents.stream().map(d -> new EmbeddedDocument(d, model.embed(d))).toList());
		}

		@Override
		public void upsert(List<EmbeddedDocument> documents) {
			documents.forEach(d -> rows.put(d.document().getId(), d));
		}

		@Override
		public void delete(List<String> idList) {
			idList.forEach(rows::remove);
		}

		@Override
		public void delete(Filter.Expression filterExpression) {
			throw new UnsupportedOperationException();
		}

		@Override
		public List<Document> similaritySearch(SearchRequest request) {
			return rows.values().stream().map(EmbeddedDocument::document).toList();
		}

		@Override
		public <T> Optional<T> getNativeClient() {
			return Optional.empty();
		}
	}

	static Document doc(String text) {
		// Idempotent id: the same text always maps to the same id, so a re-run replaces instead of duplicating.
		String id = UUID.nameUUIDFromBytes(text.getBytes()).toString();
		return Document.builder().id(id).text(text).metadata(Map.of("source", "demo")).build();
	}

	public static void main(String[] args) {
		List<Document> docs = List.of(doc("Spring AI message parts"), doc("Responses API hosted tools"), doc("Share groups in Kafka"));

		System.out.println("== 1. SimpleVectorStore has no doUpsert");
		CountingModel simpleModel = new CountingModel();
		SimpleVectorStore simple = SimpleVectorStore.builder(simpleModel).build();
		try {
			simple.upsert(List.of(new EmbeddedDocument(docs.get(0), new float[] { 1, 0, 0, 0 })));
			System.out.println("upsert accepted");
		}
		catch (UnsupportedOperationException e) {
			System.out.println("UnsupportedOperationException: " + e.getMessage());
		}
		System.out.println("embedding calls = " + simpleModel.texts.get());

		System.out.println("== 2. add() embeds, upsert() with held vectors does not");
		CountingModel sourceModel = new CountingModel();
		ToyStore source = new ToyStore(sourceModel);
		source.add(docs);
		System.out.println("source after add          : rows=" + source.rows.size() + " embeddingCalls=" + sourceModel.texts.get());

		CountingModel targetModel = new CountingModel();
		ToyStore target = new ToyStore(targetModel);
		target.upsert(new ArrayList<>(source.rows.values()));
		System.out.println("target after upsert (copy): rows=" + target.rows.size() + " embeddingCalls=" + targetModel.texts.get());

		System.out.println("== 3. idempotent ids: running the copy again");
		target.upsert(new ArrayList<>(source.rows.values()));
		System.out.println("target after second run   : rows=" + target.rows.size() + " embeddingCalls=" + targetModel.texts.get());
	}
}