package com.ankurm.semanticcache;

import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
import org.springframework.core.io.DefaultResourceLoader;

/**
 * all-MiniLM-L6-v2 (384 dimensions) running inside the JVM. Spring AI's default download location for
 * the model is a GitHub media URL; here both files come from Hugging Face and are cached on disk, so the
 * first run downloads about 90 MB and later runs are offline.
 */
public final class Embedder {

    static final String HF = "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/";

    private Embedder() { }

    public static EmbeddingModel miniLm() {
        try {
            TransformersEmbeddingModel m = new TransformersEmbeddingModel();
            m.setTokenizerResource(HF + "tokenizer.json");
            m.setModelResource(HF + "onnx/model.onnx");
            m.setResourceCacheDirectory(System.getProperty("java.io.tmpdir") + "/minilm-cache");
            m.setTokenizerOptions(java.util.Map.of("padding", "true", "truncation", "true", "maxLength", "256"));
            m.afterPropertiesSet();
            return m;
        } catch (Exception e) {
            throw new IllegalStateException(e);
        }
    }

    public 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) + 1e-12);
    }
}