import numpy as np
import pytest
from common.vector_index import create_index, save_index, load_index, add_vector, search

def test_add_vector_returns_sequential_ids():
    index = create_index(dim=4)
    v1 = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
    v2 = np.array([0.0, 1.0, 0.0, 0.0], dtype=np.float32)
    assert add_vector(index, v1) == 0
    assert add_vector(index, v2) == 1
    assert index.ntotal == 2

def test_search_returns_closest_vector_first():
    index = create_index(dim=4)
    add_vector(index, np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
    add_vector(index, np.array([0.0, 1.0, 0.0, 0.0], dtype=np.float32))
    query = np.array([0.9, 0.1, 0.0, 0.0], dtype=np.float32)
    results = search(index, query, k=2)
    assert results[0][0] == 0
    assert results[0][1] > results[1][1]

def test_save_and_load_roundtrip(tmp_path):
    index = create_index(dim=4)
    add_vector(index, np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
    path = str(tmp_path / "index.faiss")
    save_index(index, path)
    loaded = load_index(path)
    assert loaded.ntotal == 1
    assert loaded.d == 4

def test_load_missing_index_raises():
    with pytest.raises(FileNotFoundError):
        load_index("/nonexistent/path/index.faiss")