import os
import logging
import faiss
import numpy as np
from typing import List, Optional, Dict, Any
from sqlalchemy.orm import Session
from sqlalchemy import select
from .base import VectorStore
from app.db.models import VectorMetadata, Document  # Import your SQLAlchemy models


class FaissVectorStore(VectorStore):
    """
    A FAISS vector store that uses a relational database (via SQLAlchemy)
    for persistence and filtering of document metadata.
    """

    def __init__(self, index_file_path: str, dimension: int, embedder):
        self.index_file_path = index_file_path
        self.dimension = dimension
        self.embedder = embedder
        self.index = None
        self.doc_id_map = []
        import threading
        self._save_lock = threading.Lock()

    def initialize_index(self, db_session: Session):
        """Initializes the FAISS index and syncs it with the database."""
        if os.path.exists(self.index_file_path):
            logging.info(f"Loading FAISS index from {self.index_file_path}")
            self.index = faiss.read_index(self.index_file_path)
        else:
            logging.info("Creating a new FAISS index.")
            quantizer = faiss.IndexFlatL2(self.dimension)
            self.index = faiss.IndexIDMap(quantizer)

        self.sync_with_db(db_session)

    def sync_with_db(self, db_session: Session):
        """Syncs the in-memory FAISS ID map with the database metadata."""
        logging.info("Synchronizing FAISS index with database metadata.")
        faiss_ids_from_db = db_session.execute(select(VectorMetadata.id)).scalars().all()
        self.doc_id_map = faiss_ids_from_db

    def add_document(self, text: str, document_id: int,
                     session_id: Optional[int] = None,
                     embedding_model: str = "default_model",
                     db_session: Session = None) -> int:
        """Embeds and adds a single document to FAISS and DB."""
        if db_session is None:
            raise ValueError("db_session must be provided")

        logging.debug("Embedding document text for FAISS index...")
        vector = self.embedder.embed_text(text).reshape(1, -1).astype('float32')

        metadata_entry = VectorMetadata(
            document_id=document_id,
            session_id=session_id,
            embedding_model=embedding_model,
        )
        db_session.add(metadata_entry)
        db_session.flush()

        faiss_index_id = metadata_entry.id
        if self.index is None:
            self.initialize_index(db_session)
        self.index.add_with_ids(vector, np.array([faiss_index_id], dtype='int64'))
        self.doc_id_map.append(faiss_index_id)

        db_session.commit()
        self.save_index()

        logging.info(f"Document added to FAISS index with ID: {faiss_index_id}")
        return faiss_index_id

    def add_multiple_documents(self, texts: List[str], document_ids: List[int],
                               session_id: Optional[int] = None,
                               embedding_model: str = "default_model",
                               db_session: Session = None) -> List[int]:
        """Adds multiple documents to FAISS and DB."""
        if db_session is None:
            raise ValueError("db_session must be provided")

        if len(texts) != len(document_ids):
            raise ValueError("The number of texts must match the number of document IDs.")

        logging.debug("Embedding multiple document texts for FAISS index...")
        vectors = self.embedder.embed_multiple_texts(texts).astype('float32')

        new_metadata_entries = []
        for doc_id in document_ids:
            new_metadata_entries.append(VectorMetadata(
                document_id=doc_id,
                session_id=session_id,
                embedding_model=embedding_model,
            ))

        db_session.add_all(new_metadata_entries)
        db_session.flush()

        new_faiss_ids = [entry.id for entry in new_metadata_entries]
        if self.index is None:
            self.initialize_index(db_session)
        self.index.add_with_ids(vectors, np.array(new_faiss_ids, dtype='int64'))
        self.doc_id_map.extend(new_faiss_ids)

        db_session.commit()
        self.save_index()

        logging.info(f"Added {len(new_faiss_ids)} documents to FAISS index.")
        return new_faiss_ids

    def remove_vectors(self, faiss_ids: List[int]) -> int:
        """Removes vectors with specified IDs from the FAISS index."""
        if self.index is None:
            logging.warning("FAISS index is not initialized.")
            return 0
        
        logging.info(f"Removing {len(faiss_ids)} vectors from FAISS index...")
        selector = faiss.IDSelectorBatch(faiss_ids)
        num_removed = self.index.remove_ids(selector)
        
        self.doc_id_map = [i for i in self.doc_id_map if i not in faiss_ids]
        
        self.save_index()
        logging.info(f"Successfully removed {num_removed} vectors from FAISS index.")
        return num_removed

    def search_similar_documents(self, query_text: str, k: int = 5,
                                 prefilter_tags: Optional[Dict[str, Any]] = None,
                                 db_session: Session = None) -> List[int]:
        """Searches FAISS for similar documents, with optional DB filter."""
        if db_session is None:
            raise ValueError("db_session must be provided")

        logging.debug("Searching FAISS index with database-backed filtering.")
        if self.index.ntotal == 0:
            logging.warning("FAISS index is empty, no documents to search.")
            return []

        query_vector = self.embedder.embed_text(query_text).reshape(1, -1).astype('float32')

        if prefilter_tags:
            db_query = select(VectorMetadata.id)
            for key, value in prefilter_tags.items():
                if hasattr(VectorMetadata, key):
                    db_query = db_query.where(getattr(VectorMetadata, key) == value)
                else:
                    logging.warning(f"Metadata key '{key}' not found in VectorMetadata model.")
                    return []

            valid_faiss_ids = db_session.execute(db_query).scalars().all()

            if not valid_faiss_ids:
                logging.warning("No documents match the filter criteria.")
                return []

            id_selector = faiss.IDSelectorBatch(valid_faiss_ids)
            D, I = self.index.search(query_vector, min(k, len(valid_faiss_ids)),
                                     params=faiss.SearchParameters(sel=id_selector))
            result_faiss_ids = [int(i) for i in I.flatten() if i >= 0]
        else:
            D, I = self.index.search(query_vector, k)
            result_faiss_ids = [int(i) for i in I.flatten() if i >= 0]

        result_document_ids = db_session.execute(
            select(VectorMetadata.document_id).where(VectorMetadata.id.in_(result_faiss_ids))
        ).scalars().all()

        logging.info(f"Search complete, found {len(result_document_ids)} similar documents.")
        return result_document_ids

    def save_index(self):
        """Saves the FAISS index to the file system with concurrency protection."""
        if self.index:
            with self._save_lock:
                logging.info(f"Saving FAISS index to {self.index_file_path}")
                faiss.write_index(self.index, self.index_file_path)
