# Handles conversion of raw data into vector embeddings
import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List, Dict, Any

class EmbeddingService:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)
        self.dimension = self.model.get_sentence_embedding_dimension()

    def embed_text(self, text: str) -> np.ndarray:
        return self.model.encode(text, normalize_embeddings=True)

    def embed_batch(self, texts: List[str]) -> np.ndarray:
        return self.model.encode(
            texts,
            normalize_embeddings=True,
            batch_size=32,
            show_progress_bar=False
        )

    def get_dimension(self) -> int:
        return self.dimension

# Implements common distance functions for vector similarity
import numpy as np
from enum import Enum
from typing import Callable

class DistanceMetric(Enum):
    COSINE = "cosine"
    EUCLIDEAN = "euclidean"
    DOT_PRODUCT = "dot_product"

def cosine_distance(a: np.ndarray, b: np.ndarray) -> float:
    dot = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    return 1 - (dot / (norm_a * norm_b))

def euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
    return np.linalg.norm(a - b)

def dot_product_distance(a: np.ndarray, b: np.ndarray) -> float:
    return -np.dot(a, b)

def get_distance_function(metric: DistanceMetric) -> Callable:
    functions = {
        DistanceMetric.COSINE: cosine_distance,
        DistanceMetric.EUCLIDEAN: euclidean_distance,
        DistanceMetric.DOT_PRODUCT: dot_product_distance
    }
    return functions[metric]

# Minimal vector index (flat search with metric support)
import numpy as np
from dataclasses import dataclass
from typing import List, Tuple, Optional, Dict, Any

@dataclass
class VectorRecord:
    id: str
    vector: np.ndarray
    metadata: Dict[str, Any]

class VectorIndex:
    def __init__(self, dimension: int, metric: str = "cosine", index_type: str = "flat"):
        self.dimension = dimension
        self.metric = metric
        self.index_type = index_type
        self.vectors: List[VectorRecord] = []
        self.id_to_index: Dict[str, int] = {}

    def add(self, id: str, vector: np.ndarray, metadata: Dict[str, Any] = None):
        if vector.shape[0] != self.dimension:
            raise ValueError(f"Expected dimension {self.dimension}, got {vector.shape[0]}")
        if self.metric == "cosine":
            vector = vector / np.linalg.norm(vector)
        record = VectorRecord(id=id, vector=vector, metadata=metadata or {})
        self.id_to_index[id] = len(self.vectors)
        self.vectors.append(record)

    def search(self, query_vector: np.ndarray, k: int = 10, filter_fn: Optional[callable] = None):
        if self.metric == "cosine":
            query_vector = query_vector / np.linalg.norm(query_vector)
        results = []
        for record in self.vectors:
            if filter_fn and not filter_fn(record.metadata):
                continue
            distance = self._calculate_distance(query_vector, record.vector)
            results.append((record.id, distance, record.metadata))
        results.sort(key=lambda x: x[1])
        return results[:k]

    def _calculate_distance(self, a: np.ndarray, b: np.ndarray) -> float:
        if self.metric == "cosine":
            return 1 - np.dot(a, b)
        elif self.metric == "euclidean":
            return np.linalg.norm(a - b)
        else:
            return -np.dot(a, b)

# Creates composable filters for metadata-based filtering
from typing import Dict, Any, Callable, List
from dataclasses import dataclass

@dataclass
class FilterCondition:
    field: str
    operator: str
    value: Any

class FilterBuilder:
    def __init__(self):
        self.conditions: List[FilterCondition] = []

    def equals(self, field: str, value: Any) -> "FilterBuilder":
        self.conditions.append(FilterCondition(field, "eq", value))
        return self

    def greater_than(self, field: str, value: Any) -> "FilterBuilder":
        self.conditions.append(FilterCondition(field, "gt", value))
        return self

    def less_than(self, field: str, value: Any) -> "FilterBuilder":
        self.conditions.append(FilterCondition(field, "lt", value))
        return self

    def in_list(self, field: str, values: List[Any]) -> "FilterBuilder":
        self.conditions.append(FilterCondition(field, "in", values))
        return self

    def build(self) -> Callable[[Dict[str, Any]], bool]:
        conditions = self.conditions.copy()
        def filter_fn(metadata: Dict[str, Any]) -> bool:
            for cond in conditions:
                value = metadata.get(cond.field)
                if value is None:
                    return False
                if cond.operator == "eq" and value != cond.value:
                    return False
                elif cond.operator == "gt" and value <= cond.value:
                    return False
                elif cond.operator == "lt" and value >= cond.value:
                    return False
                elif cond.operator == "in" and value not in cond.value:
                    return False
            return True
        return filter_fn

from embedding_service import EmbeddingService
from vector_index import VectorIndex
from filter_builder import FilterBuilder
from typing import List, Dict, Any, Optional

class VectorDatabase:
    def __init__(self, collection_name: str, dimension: int = 384):
        self.collection_name = collection_name
        self.embedding_service = EmbeddingService()
        self.index = VectorIndex(dimension=dimension, metric="cosine", index_type="flat")

    def upsert(self, id: str, text: str, metadata: Dict[str, Any] = None):
        vector = self.embedding_service.embed_text(text)
        full_metadata = {"text": text, **(metadata or {})}
        self.index.add(id, vector, full_metadata)

    def query(self, query_text: str, k: int = 10, filter_builder: Optional[FilterBuilder] = None):
        query_vector = self.embedding_service.embed_text(query_text)
        filter_fn = filter_builder.build() if filter_builder else None
        results = self.index.search(query_vector, k=k, filter_fn=filter_fn)
        return [
            {"id": id, "score": 1 - distance, "metadata": metadata}
            for id, distance, metadata in results
        ]

if __name__ == "__main__":
    db = VectorDatabase("documents")
    db.upsert("doc1", "Python is great for machine learning", {"category": "tech"})
    db.upsert("doc2", "TensorFlow and PyTorch are popular ML frameworks", {"category": "tech"})
    db.upsert("doc3", "The weather is sunny today", {"category": "general"})

    filt = FilterBuilder().equals("category", "tech")
    results = db.query("deep learning frameworks", k=2, filter_builder=filt)
    for r in results:
        print(f"Score: {r['score']:.3f} - {r['metadata']['text']}")