"""
Embeddings module using Ollama.

Provides a simple caching layer and a function to embed text using Ollama's
embedding endpoint.  No OpenAI services are used.
"""

import json
import os
from typing import List, Dict

import ollama
import numpy as np

# Cache to avoid repeated calls for the same text
_EMBED_CACHE: Dict[str, List[float]] = {}


def embed(text: str, model: str = "llama2") -> List[float]:
    """
    Generate an embedding vector for the given text using Ollama.

    Parameters
    ----------
    text : str
        The text to embed.
    model : str, optional
        The Ollama model to use for embeddings. Defaults to "llama2".

    Returns
    -------
    List[float]
        The embedding vector.
    """
    if text in _EMBED_CACHE:
        return _EMBED_CACHE[text]

    # Ollama expects a dict with "model" and "prompt"
    payload = {"model": model, "prompt": text}
    try:
        response = ollama.embeddings(payload)
    except Exception as exc:
        raise RuntimeError(f"Failed to get embeddings from Ollama: {exc}") from exc

    # Ollama returns a dict with "embedding" key
    embedding = response.get("embedding")
    if embedding is None:
        raise ValueError("Ollama response missing 'embedding' field")

    _EMBED_CACHE[text] = embedding
    return embedding


def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
    """
    Compute cosine similarity between two vectors.

    Parameters
    ----------
    vec1, vec2 : List[float]
        Input vectors.

    Returns
    -------
    float
        Cosine similarity score.
    """
    v1 = np.array(vec1)
    v2 = np.array(vec2)
    dot = np.dot(v1, v2)
    norm1 = np.linalg.norm(v1)
    norm2 = np.linalg.norm(v2)
    if norm1 == 0 or norm2 == 0:
        return 0.0
    return dot / (norm1 * norm2)