import os
import shutil
import tempfile
import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
import numpy as np
from typing import List, Dict

from app.core.vector_store.faiss_store import FaissVectorStore
from app.db.models import Base, VectorMetadata, Document, Session

# -----------------------------
# Fixtures
# -----------------------------

@pytest.fixture(scope="function")
def engine():
    return create_engine("sqlite:///:memory:")

@pytest.fixture(scope="function")
def setup_database(engine):
    Base.metadata.create_all(engine)
    yield
    Base.metadata.drop_all(engine)

@pytest.fixture
def db_session(setup_database, engine):
    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    session = SessionLocal()
    session.begin()
    try:
        yield session
    finally:
        session.rollback()
        session.close()

@pytest.fixture
def mock_embedder():
    class MockEmbedder:
        def embed_text(self, text):
            np.random.seed(abs(hash(text)) % 2**32)
            return np.random.rand(768).astype('float32')

        def embed_multiple_texts(self, texts):
            return np.vstack([self.embed_text(text) for text in texts])
    return MockEmbedder()

@pytest.fixture
def temp_faiss_file():
    tmp_dir = tempfile.mkdtemp()
    index_file = os.path.join(tmp_dir, "test.index")
    yield index_file
    shutil.rmtree(tmp_dir)

@pytest.fixture
def faiss_store(temp_faiss_file, mock_embedder, db_session):
    store = FaissVectorStore(
        index_file_path=temp_faiss_file,
        dimension=768,
        embedder=mock_embedder
    )
    store.initialize_index(db_session)
    return store

# -----------------------------
# Helper Function for Test Setup
# -----------------------------
def create_test_data(db_session, faiss_store, texts_and_metadata: List[Dict]):
    dummy_session = Session(user_id="test_user", title="Test Session")
    db_session.add(dummy_session)
    db_session.commit()

    document_ids = []
    texts = []

    for item in texts_and_metadata:
        dummy_doc = Document(title=item.get("title", "Test Doc"), text=item["text"])
        db_session.add(dummy_doc)
        db_session.flush()
        document_ids.append(dummy_doc.id)
        texts.append(dummy_doc.text)

    db_session.commit()

    faiss_store.add_multiple_documents(
        texts=texts,
        document_ids=document_ids,
        session_id=dummy_session.id,
        embedding_model="test_model",
        db_session=db_session
    )

    return document_ids, dummy_session.id

# -----------------------------
# Tests
# -----------------------------

def test_add_document(faiss_store, db_session):
    test_text = "This is a test document."
    test_doc = Document(title="Test", text=test_text)
    db_session.add(test_doc)
    db_session.commit()
    document_id = test_doc.id

    assert faiss_store.index.ntotal == 0
    assert db_session.query(VectorMetadata).count() == 0

    faiss_index_id = faiss_store.add_document(test_text, document_id=document_id, db_session=db_session)

    assert faiss_store.index.ntotal == 1
    assert os.path.exists(faiss_store.index_file_path)

    metadata_entry = db_session.query(VectorMetadata).filter_by(id=faiss_index_id).first()
    assert metadata_entry is not None
    assert metadata_entry.document_id == document_id
    assert metadata_entry.id == faiss_index_id

def test_add_multiple_documents(faiss_store, db_session):
    docs = [
        {"text": "Doc 1"},
        {"text": "Doc 2"},
        {"text": "Doc 3"}
    ]

    doc_ids, _ = create_test_data(db_session, faiss_store, docs)

    assert faiss_store.index.ntotal == 3
    assert db_session.query(VectorMetadata).count() == 3

    db_faiss_ids = db_session.query(VectorMetadata.id).order_by(VectorMetadata.id).all()
    assert [row[0] for row in db_faiss_ids] == faiss_store.doc_id_map

def test_load_existing_index_and_sync_with_db(temp_faiss_file, mock_embedder, db_session):
    store1 = FaissVectorStore(temp_faiss_file, 768, mock_embedder)
    store1.initialize_index(db_session)
    create_test_data(db_session, store1, [{"text": "Persistence test 1"}, {"text": "Persistence test 2"}])

    store2 = FaissVectorStore(temp_faiss_file, 768, mock_embedder)
    store2.initialize_index(db_session)

    assert store2.index.ntotal == 2
    assert db_session.query(VectorMetadata).count() == 2

    expected_ids = db_session.query(VectorMetadata.id).order_by(VectorMetadata.id).all()
    expected_ids = [row[0] for row in expected_ids]
    assert sorted(store2.doc_id_map) == sorted(expected_ids)

def test_search_similar_documents_without_filter(faiss_store, db_session):
    docs = [
        {"text": "The sun is a star.", "title": "astronomy"},
        {"text": "Mars is a planet.", "title": "astronomy"},
        {"text": "The moon orbits the Earth.", "title": "astronomy"}
    ]
    create_test_data(db_session, faiss_store, docs)

    results = faiss_store.search_similar_documents("What is a star?", k=2, db_session=db_session)

    assert len(results) == 2
    assert isinstance(results[0], int)

def test_search_similar_documents_with_filter(faiss_store, db_session):
    dummy_session1 = Session(user_id="user1", title="Session 1")
    dummy_session2 = Session(user_id="user2", title="Session 2")
    db_session.add_all([dummy_session1, dummy_session2])
    db_session.commit()

    doc1 = Document(title="Python", text="Python is a programming language.")
    doc2 = Document(title="Dog", text="A dog is a loyal pet.")
    doc3 = Document(title="Java", text="Java is another programming language.")
    db_session.add_all([doc1, doc2, doc3])
    db_session.commit()

    faiss_store.add_document(doc1.text, doc1.id, session_id=dummy_session1.id, db_session=db_session)
    faiss_store.add_document(doc2.text, doc2.id, session_id=dummy_session2.id, db_session=db_session)
    faiss_store.add_document(doc3.text, doc3.id, session_id=dummy_session1.id, db_session=db_session)

    results = faiss_store.search_similar_documents(
        "Which is a programming language?", k=2,
        prefilter_tags={"session_id": dummy_session1.id},
        db_session=db_session
    )

    assert len(results) == 2
    for doc_id in results:
        metadata = db_session.query(VectorMetadata).filter_by(document_id=doc_id).first()
        assert metadata.session_id == dummy_session1.id

def test_search_with_no_matching_filter(faiss_store, db_session):
    dummy_session = Session(user_id="user", title="Some Session")
    dummy_doc = Document(title="Cat", text="A document about cats.")
    db_session.add_all([dummy_session, dummy_doc])
    db_session.commit()
    faiss_store.add_document(dummy_doc.text, dummy_doc.id, session_id=dummy_session.id, db_session=db_session)

    results = faiss_store.search_similar_documents(
        "What is a dog?", k=5, prefilter_tags={"session_id": 9999}, db_session=db_session
    )

    assert len(results) == 0

def test_search_with_invalid_filter_key(faiss_store, db_session):
    results = faiss_store.search_similar_documents(
        "A query.", k=1, prefilter_tags={"invalid_key": "some_value"}, db_session=db_session
    )
    assert len(results) == 0
