package ankurm.qa;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
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.converter.BeanOutputConverter;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;

/**
 * Small, runnable facts behind the answers in "Top 25 Spring AI and GenAI interview questions".
 * Spring AI 2.0.1. No model is called: embeddings are a deterministic toy and nothing leaves the machine.
 */
public class InterviewDemos {

	/** Toy embedding: counts of four keywords. Good enough to show how similarity search behaves. */
	static final class KeywordEmbedding implements EmbeddingModel {
		final AtomicInteger calls = new AtomicInteger();
		static final String[] WORDS = { "refund", "invoice", "password", "shipping" };

		float[] vec(String t) {
			calls.incrementAndGet();
			float[] v = new float[WORDS.length];
			String s = t.toLowerCase();
			for (int i = 0; i < WORDS.length; i++) {
				v[i] = s.contains(WORDS[i]) ? 1f : 0f;
			}
			return v;
		}

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

		public float[] embed(Document d) {
			return vec(d.getText());
		}

		public int dimensions() {
			return WORDS.length;
		}
	}

	static class WeatherTools {
		@Tool(description = "Current temperature in a city")
		String temperature(@ToolParam(description = "City name") String city) {
			return "31C in " + city;
		}
	}

	record Ticket(String category, int urgency) {
	}

	static double cosine(float[] a, float[] b) {
		double dot = 0, na = 0, nb = 0;
		for (int i = 0; i < a.length; i++) {
			dot += a[i] * b[i];
			na += a[i] * a[i];
			nb += b[i] * b[i];
		}
		return dot / (Math.sqrt(na) * Math.sqrt(nb));
	}

	public static void main(String[] args) {
		System.out.println("== Q: what does similarity search return? (toy embeddings, SimpleVectorStore)");
		KeywordEmbedding model = new KeywordEmbedding();
		SimpleVectorStore store = SimpleVectorStore.builder(model).build();
		store.add(List.of(
				new Document("How to request a refund for an invoice", Map.of("team", "billing")),
				new Document("Reset your password from the login page", Map.of("team", "support")),
				new Document("Shipping times for international orders", Map.of("team", "logistics"))));
		System.out.println("embedding calls to index 3 documents = " + model.calls.get());
		List<Document> hits = store.similaritySearch(SearchRequest.builder().query("refund please").topK(2).build());
		hits.forEach(d -> System.out.println("hit: " + d.getText() + " score=" + d.getScore()));
		System.out.println("embedding calls after one query      = " + model.calls.get());
		System.out.println("cosine(refund, invoice-refund doc)   = " + cosine(model.vec("refund"), model.vec("refund invoice")));
		System.out.println("cosine(refund, password doc)         = " + cosine(model.vec("refund"), model.vec("password")));

		System.out.println("== Q: how does chunk size change the number of chunks? (TokenTextSplitter)");
		String text = "Spring AI makes it easy to build retrieval pipelines. ".repeat(200);
		for (int size : new int[] { 100, 400, 800 }) {
			TokenTextSplitter splitter = TokenTextSplitter.builder().withChunkSize(size).build();
			System.out.println("chunkSize " + size + " -> chunks = " + splitter.split(new Document(text)).size());
		}

		System.out.println("== Q: what does the model actually receive for a tool? (@Tool)");
		ToolCallback tool = ToolCallbacks.from(new WeatherTools())[0];
		System.out.println("name        = " + tool.getToolDefinition().name());
		System.out.println("description= " + tool.getToolDefinition().description());
		System.out.println("inputSchema = " + tool.getToolDefinition().inputSchema().replaceAll("\\s+", " "));
		System.out.println("call result = " + tool.call("{\"city\":\"Mumbai\"}"));

		System.out.println("== Q: what does structured output add to the prompt? (BeanOutputConverter)");
		BeanOutputConverter<Ticket> converter = new BeanOutputConverter<>(Ticket.class);
		System.out.println(converter.getFormat().replaceAll("\\s+", " "));
		Ticket t = converter.convert("{\"category\":\"billing\",\"urgency\":2}");
		System.out.println("parsed = " + t);
		try {
			converter.convert("Sure! Here is the JSON: {\"category\":\"billing\"");
		}
		catch (RuntimeException e) {
			System.out.println("malformed reply -> " + e.getClass().getSimpleName());
		}
	}
}