import hashlib
import threading
import uuid
from datetime import UTC, datetime
from typing import Any, Dict, List, Optional

from langchain_community.embeddings import (
    HuggingFaceEmbeddings,
)
from langchain_core.documents import Document
from loguru import logger

from ...database.models.library import DocumentChunk
from ...database.session_context import get_user_db_session
from ...utilities.url_utils import normalize_url


class LocalEmbeddingManager:
    """Handles embedding generation and storage for local document search"""

    def __init__(
        self,
        embedding_model: str = "all-MiniLM-L6-v2",
        embedding_device: str = "cpu",
        embedding_model_type: str = "sentence_transformers",  # or 'ollama'
        ollama_base_url: Optional[str] = None,
        settings_snapshot: Optional[Dict[str, Any]] = None,
    ):
        """
        Initialize the embedding manager for local document search.

        Args:
            embedding_model: Name of the embedding model to use
            embedding_device: Device to run embeddings on ('cpu' or 'cuda')
            embedding_model_type: Type of embedding model ('sentence_transformers' or 'ollama')
            ollama_base_url: Base URL for Ollama API if using ollama embeddings
            settings_snapshot: Optional settings snapshot for background threads
        """

        self.embedding_model = embedding_model
        self.embedding_device = embedding_device
        self.embedding_model_type = embedding_model_type
        self.ollama_base_url = ollama_base_url
        self.settings_snapshot = settings_snapshot or {}

        # Username for database access (extracted from settings if available)
        self.username = (
            settings_snapshot.get("_username") if settings_snapshot else None
        )
        # Password for encrypted database access (can be set later)
        self.db_password = None

        # Initialize the embedding model (with lock for thread-safe lazy init)
        self._embeddings = None
        self._embedding_lock = threading.Lock()

        # Vector store cache
        self.vector_stores: dict[str, Any] = {}

        # Track if this manager has been closed
        self._closed = False

    def close(self):
        """Release embedding model resources."""
        if self._closed:
            return
        self._closed = True
        # Clear embedding model reference to allow garbage collection
        self._embeddings = None
        # Clear vector store cache
        self.vector_stores.clear()
        logger.debug("LocalEmbeddingManager closed")

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - ensures resources are released."""
        self.close()
        return False

    @property
    def embeddings(self):
        """
        Lazily initialize embeddings when first accessed.
        This allows the LocalEmbeddingManager to be created without
        immediately loading models, which is helpful when no local search is performed.

        Uses double-checked locking to ensure thread-safe initialization.
        Concurrent SentenceTransformer model loading causes meta tensor errors
        in PyTorch when multiple threads call model.to(device) simultaneously.
        """
        if self._embeddings is None:
            with self._embedding_lock:
                if self._embeddings is None:
                    logger.info("Initializing embeddings on first use")
                    self._embeddings = self._initialize_embeddings()
        return self._embeddings

    def _initialize_embeddings(self):
        """Initialize the embedding model based on configuration"""
        try:
            # Use the new unified embedding system
            from ...embeddings import get_embeddings

            # Prepare kwargs for provider-specific parameters
            kwargs = {}

            # Add device for sentence transformers
            if self.embedding_model_type == "sentence_transformers":
                kwargs["device"] = self.embedding_device

            # Add base_url for ollama if specified
            if self.embedding_model_type == "ollama" and self.ollama_base_url:
                kwargs["base_url"] = normalize_url(self.ollama_base_url)

            logger.info(
                f"Initializing embeddings with provider={self.embedding_model_type}, model={self.embedding_model}"
            )

            return get_embeddings(
                provider=self.embedding_model_type,
                model=self.embedding_model,
                settings_snapshot=self.settings_snapshot,
                **kwargs,
            )
        except Exception:
            logger.exception("Error initializing embeddings")
            logger.warning(
                "Falling back to HuggingFaceEmbeddings with all-MiniLM-L6-v2"
            )
            return HuggingFaceEmbeddings(
                model_name="sentence-transformers/all-MiniLM-L6-v2"
            )

    def _store_chunks_to_db(
        self,
        chunks: List[Document],
        collection_name: str,
        source_path: Optional[str] = None,
        source_id: Optional[int] = None,
        source_type: str = "local_file",
    ) -> List[str]:
        """
        Store document chunks in the database.

        Args:
            chunks: List of LangChain Document chunks
            collection_name: Name of the collection (e.g., 'personal_notes', 'library')
            source_path: Path to source file (for local files)
            source_id: ID of source document (for library documents)
            source_type: Type of source ('local_file' or 'library')

        Returns:
            List of chunk embedding IDs (UUIDs) for FAISS mapping
        """
        if not self.username:
            logger.warning(
                "No username available, cannot store chunks in database"
            )
            return []

        chunk_ids = []

        try:
            with get_user_db_session(
                self.username, self.db_password
            ) as session:
                for idx, chunk in enumerate(chunks):
                    # Generate unique hash for chunk
                    chunk_text = chunk.page_content
                    chunk_hash = hashlib.sha256(chunk_text.encode()).hexdigest()

                    # Generate unique embedding ID
                    embedding_id = uuid.uuid4().hex

                    # Extract metadata
                    metadata = chunk.metadata or {}
                    document_title = metadata.get(
                        "filename", metadata.get("title", "Unknown")
                    )

                    # Calculate word count
                    word_count = len(chunk_text.split())

                    # Get character positions from metadata if available
                    start_char = metadata.get("start_char", 0)
                    end_char = metadata.get("end_char", len(chunk_text))

                    # Check if chunk already exists
                    existing_chunk = (
                        session.query(DocumentChunk)
                        .filter_by(chunk_hash=chunk_hash)
                        .first()
                    )

                    if existing_chunk:
                        # Update existing chunk
                        existing_chunk.last_accessed = datetime.now(UTC)
                        chunk_ids.append(existing_chunk.embedding_id)
                        logger.debug(
                            f"Chunk already exists, reusing: {existing_chunk.embedding_id}"
                        )
                    else:
                        # Create new chunk
                        db_chunk = DocumentChunk(
                            chunk_hash=chunk_hash,
                            source_type=source_type,
                            source_id=source_id,
                            source_path=str(source_path)
                            if source_path
                            else None,
                            collection_name=collection_name,
                            chunk_text=chunk_text,
                            chunk_index=idx,
                            start_char=start_char,
                            end_char=end_char,
                            word_count=word_count,
                            embedding_id=embedding_id,
                            embedding_model=self.embedding_model,
                            embedding_model_type=self.embedding_model_type,
                            document_title=document_title,
                            document_metadata=metadata,
                        )
                        session.add(db_chunk)
                        chunk_ids.append(embedding_id)

                session.commit()
                logger.info(
                    f"Stored {len(chunk_ids)} chunks to database for collection '{collection_name}'"
                )

        except Exception:
            logger.exception(
                f"Error storing chunks to database for collection '{collection_name}'"
            )
            return []

        return chunk_ids

    def _delete_chunks_from_db(
        self,
        collection_name: str,
        source_path: Optional[str] = None,
        source_id: Optional[int] = None,
    ) -> int:
        """
        Delete chunks from database.

        Args:
            collection_name: Name of the collection
            source_path: Path to source file (for local files)
            source_id: ID of source document (for library documents)

        Returns:
            Number of chunks deleted
        """
        if not self.username:
            logger.warning(
                "No username available, cannot delete chunks from database"
            )
            return 0

        try:
            with get_user_db_session(
                self.username, self.db_password
            ) as session:
                query = session.query(DocumentChunk).filter_by(
                    collection_name=collection_name
                )

                if source_path:
                    query = query.filter_by(source_path=str(source_path))
                if source_id:
                    query = query.filter_by(source_id=source_id)

                count = int(query.delete())
                session.commit()

                logger.info(
                    f"Deleted {count} chunks from database for collection '{collection_name}'"
                )
                return count

        except Exception:
            logger.exception(
                f"Error deleting chunks from database for collection '{collection_name}'"
            )
            return 0
