"""
Vector store module for Toxo - provides persistent vector database functionality.
"""

import asyncio
import uuid
import numpy as np
from typing import Dict, Any, List, Optional, Tuple, Union
from dataclasses import dataclass, asdict
from datetime import datetime
import json
from pathlib import Path

# Vector database integrations
try:
    import chromadb
    from chromadb.config import Settings
    CHROMADB_AVAILABLE = True
except (ImportError, AttributeError):
    # Handle both import errors and numpy compatibility issues
    chromadb = None
    CHROMADB_AVAILABLE = False

# SQL database for metadata
from sqlalchemy import create_engine, Column, String, Text, Float, Integer, DateTime, JSON, Boolean
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.dialects.postgresql import UUID
import sqlalchemy as sa

from ..utils.logger import get_logger
from ..utils.exceptions import ToxoError


Base = declarative_base()


class VectorDocument(Base):
    """SQLAlchemy model for vector document metadata."""
    
    __tablename__ = 'vector_documents'
    
    id = Column(String, primary_key=True)
    collection_name = Column(String, nullable=False, index=True)
    content = Column(Text, nullable=False)
    document_metadata = Column(JSON, default={})
    embedding_dimension = Column(Integer)
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
    is_active = Column(Boolean, default=True)


@dataclass
class VectorSearchResult:
    """Represents a vector search result."""
    id: str
    content: str
    metadata: Dict[str, Any]
    similarity_score: float
    distance: float


class VectorStore:
    """
    Vector store for persistent embeddings and similarity search.
    Supports both ChromaDB and in-memory fallback.
    """
    
    def __init__(
        self,
        collection_name: str = "toxo_vectors",
        persist_directory: Optional[str] = None,
        database_url: Optional[str] = None,
        embedding_dimension: int = 768
    ):
        self.collection_name = collection_name
        self.persist_directory = persist_directory or "./data/chroma"
        self.database_url = database_url or "sqlite:///./data/toxo_vectors.db"
        self.embedding_dimension = embedding_dimension
        self.logger = get_logger(__name__)
        
        # Initialize components
        self._init_chroma()
        self._init_sql_db()
        
        # Fallback in-memory storage
        self._memory_store: Dict[str, Dict[str, Any]] = {}
    
    def _init_chroma(self):
        """Initialize ChromaDB client."""
        if not CHROMADB_AVAILABLE:
            self.logger.warning("ChromaDB not available, using in-memory fallback")
            self.chroma_client = None
            self.chroma_collection = None
            return
        
        try:
            # Create persist directory
            Path(self.persist_directory).mkdir(parents=True, exist_ok=True)
            
            # Initialize ChromaDB client
            self.chroma_client = chromadb.PersistentClient(
                path=self.persist_directory,
                settings=Settings(
                    anonymized_telemetry=False,
                    allow_reset=True
                )
            )
            
            # Get or create collection
            self.chroma_collection = self.chroma_client.get_or_create_collection(
                name=self.collection_name,
                metadata={"hnsw:space": "cosine"}
            )
            
            self.logger.info(f"ChromaDB initialized at {self.persist_directory}")
            
        except Exception as e:
            self.logger.warning(f"ChromaDB initialization failed: {e}")
            self.chroma_client = None
            self.chroma_collection = None
    
    def _init_sql_db(self):
        """Initialize SQL database for metadata."""
        try:
            # Create database directory
            if self.database_url.startswith("sqlite"):
                db_path = Path(self.database_url.replace("sqlite:///", ""))
                db_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Create engine and session
            self.sql_engine = create_engine(self.database_url)
            Base.metadata.create_all(self.sql_engine)
            self.SessionLocal = sessionmaker(bind=self.sql_engine)
            
            self.logger.info(f"SQL database initialized: {self.database_url}")
            
        except Exception as e:
            self.logger.error(f"SQL database initialization failed: {e}")
            self.sql_engine = None
            self.SessionLocal = None
    
    async def add_document(
        self,
        content: str,
        embedding: np.ndarray,
        metadata: Optional[Dict[str, Any]] = None,
        document_id: Optional[str] = None
    ) -> str:
        """
        Add a document with its embedding to the vector store.
        
        Args:
            content: Document content
            embedding: Document embedding vector
            metadata: Additional metadata
            document_id: Optional document ID
            
        Returns:
            Document ID
        """
        document_id = document_id or str(uuid.uuid4())
        metadata = metadata or {}
        
        # Validate embedding
        if embedding.shape[0] != self.embedding_dimension:
            raise ToxoError(f"Embedding dimension mismatch: expected {self.embedding_dimension}, got {embedding.shape[0]}")
        
        try:
            # Store in ChromaDB if available
            if self.chroma_collection:
                self.chroma_collection.add(
                    ids=[document_id],
                    embeddings=[embedding.tolist()],
                    documents=[content],
                    metadatas=[metadata]
                )
            else:
                # Fallback to in-memory storage
                self._memory_store[document_id] = {
                    "content": content,
                    "embedding": embedding,
                    "metadata": metadata,
                    "created_at": datetime.now()
                }
            
            # Store metadata in SQL database
            if self.SessionLocal:
                with self.SessionLocal() as session:
                    doc = VectorDocument(
                        id=document_id,
                        collection_name=self.collection_name,
                        content=content,
                        document_metadata=metadata,
                        embedding_dimension=self.embedding_dimension
                    )
                    session.add(doc)
                    session.commit()
            
            self.logger.debug(f"Added document {document_id}")
            return document_id
            
        except Exception as e:
            self.logger.error(f"Failed to add document: {e}")
            raise ToxoError(f"Failed to add document: {e}")
    
    async def search_similar(
        self,
        query_embedding: np.ndarray,
        k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[VectorSearchResult]:
        """
        Search for similar documents.
        
        Args:
            query_embedding: Query embedding vector
            k: Number of results to return
            filter_metadata: Metadata filters
            
        Returns:
            List of search results
        """
        if self.chroma_collection:
            return await self._search_chroma(query_embedding, k, filter_metadata)
        else:
            return await self._search_memory(query_embedding, k, filter_metadata)
    
    async def _search_chroma(
        self,
        query_embedding: np.ndarray,
        k: int,
        filter_metadata: Optional[Dict[str, Any]]
    ) -> List[VectorSearchResult]:
        """Search using ChromaDB."""
        try:
            where_clause = None
            if filter_metadata:
                where_clause = {key: {"$eq": value} for key, value in filter_metadata.items()}
            
            results = self.chroma_collection.query(
                query_embeddings=[query_embedding.tolist()],
                n_results=k,
                where=where_clause
            )
            
            search_results = []
            for i in range(len(results['ids'][0])):
                result = VectorSearchResult(
                    id=results['ids'][0][i],
                    content=results['documents'][0][i],
                    metadata=results['metadatas'][0][i] or {},
                    similarity_score=1.0 - results['distances'][0][i],  # Convert distance to similarity
                    distance=results['distances'][0][i]
                )
                search_results.append(result)
            
            return search_results
            
        except Exception as e:
            self.logger.error(f"ChromaDB search failed: {e}")
            return []
    
    async def _search_memory(
        self,
        query_embedding: np.ndarray,
        k: int,
        filter_metadata: Optional[Dict[str, Any]]
    ) -> List[VectorSearchResult]:
        """Search using in-memory storage."""
        try:
            results = []
            
            for doc_id, doc_data in self._memory_store.items():
                # Apply metadata filter
                if filter_metadata:
                    doc_metadata = doc_data.get("metadata", {})
                    if not all(doc_metadata.get(key) == value for key, value in filter_metadata.items()):
                        continue
                
                # Calculate cosine similarity
                doc_embedding = doc_data["embedding"]
                similarity = np.dot(query_embedding, doc_embedding) / (
                    np.linalg.norm(query_embedding) * np.linalg.norm(doc_embedding)
                )
                
                result = VectorSearchResult(
                    id=doc_id,
                    content=doc_data["content"],
                    metadata=doc_data.get("metadata", {}),
                    similarity_score=similarity,
                    distance=1.0 - similarity
                )
                results.append(result)
            
            # Sort by similarity and return top k
            results.sort(key=lambda x: x.similarity_score, reverse=True)
            return results[:k]
            
        except Exception as e:
            self.logger.error(f"Memory search failed: {e}")
            return []
    
    async def get_document(self, document_id: str) -> Optional[Dict[str, Any]]:
        """
        Get a document by ID.
        
        Args:
            document_id: Document ID
            
        Returns:
            Document data if found
        """
        try:
            if self.chroma_collection:
                results = self.chroma_collection.get(ids=[document_id])
                if results['ids']:
                    return {
                        "id": results['ids'][0],
                        "content": results['documents'][0],
                        "metadata": results['metadatas'][0] or {}
                    }
            else:
                if document_id in self._memory_store:
                    doc_data = self._memory_store[document_id]
                    return {
                        "id": document_id,
                        "content": doc_data["content"],
                        "metadata": doc_data.get("metadata", {})
                    }
            
            return None
            
        except Exception as e:
            self.logger.error(f"Failed to get document: {e}")
            return None
    
    async def delete_document(self, document_id: str) -> bool:
        """
        Delete a document.
        
        Args:
            document_id: Document ID
            
        Returns:
            True if successful
        """
        try:
            # Delete from ChromaDB
            if self.chroma_collection:
                self.chroma_collection.delete(ids=[document_id])
            else:
                # Delete from memory
                self._memory_store.pop(document_id, None)
            
            # Mark as inactive in SQL database
            if self.SessionLocal:
                with self.SessionLocal() as session:
                    doc = session.query(VectorDocument).filter_by(id=document_id).first()
                    if doc:
                        doc.is_active = False
                        session.commit()
            
            self.logger.debug(f"Deleted document {document_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to delete document: {e}")
            return False
    
    async def update_document(
        self,
        document_id: str,
        content: Optional[str] = None,
        embedding: Optional[np.ndarray] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """
        Update a document.
        
        Args:
            document_id: Document ID
            content: New content
            embedding: New embedding
            metadata: New metadata
            
        Returns:
            True if successful
        """
        try:
            # Get existing document
            existing_doc = await self.get_document(document_id)
            if not existing_doc:
                return False
            
            # Prepare updated data
            updated_content = content or existing_doc["content"]
            updated_metadata = metadata or existing_doc["metadata"]
            
            # Delete old document
            await self.delete_document(document_id)
            
            # Add updated document
            if embedding is not None:
                await self.add_document(
                    content=updated_content,
                    embedding=embedding,
                    metadata=updated_metadata,
                    document_id=document_id
                )
            
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to update document: {e}")
            return False
    
    def get_collection_stats(self) -> Dict[str, Any]:
        """Get collection statistics."""
        try:
            if self.chroma_collection:
                count = self.chroma_collection.count()
            else:
                count = len(self._memory_store)
            
            return {
                "collection_name": self.collection_name,
                "document_count": count,
                "embedding_dimension": self.embedding_dimension,
                "storage_type": "chromadb" if self.chroma_collection else "memory",
                "persistent": self.chroma_collection is not None
            }
            
        except Exception as e:
            self.logger.error(f"Failed to get stats: {e}")
            return {"error": str(e)}
    
    async def create_index(self, index_type: str = "hnsw") -> bool:
        """
        Create an index for faster searching.
        
        Args:
            index_type: Type of index to create
            
        Returns:
            True if successful
        """
        # ChromaDB handles indexing automatically
        if self.chroma_collection:
            self.logger.info("ChromaDB handles indexing automatically")
            return True
        
        # For memory storage, no indexing needed
        self.logger.info("In-memory storage doesn't require explicit indexing")
        return True
    
    async def backup_collection(self, backup_path: str) -> bool:
        """
        Backup the collection.
        
        Args:
            backup_path: Path to save backup
            
        Returns:
            True if successful
        """
        try:
            backup_data = {
                "collection_name": self.collection_name,
                "embedding_dimension": self.embedding_dimension,
                "documents": []
            }
            
            # Export all documents
            if self.chroma_collection:
                # Get all documents from ChromaDB
                results = self.chroma_collection.get()
                for i, doc_id in enumerate(results['ids']):
                    backup_data["documents"].append({
                        "id": doc_id,
                        "content": results['documents'][i],
                        "metadata": results['metadatas'][i] or {},
                        "embedding": results['embeddings'][i] if results['embeddings'] else None
                    })
            else:
                # Export from memory
                for doc_id, doc_data in self._memory_store.items():
                    backup_data["documents"].append({
                        "id": doc_id,
                        "content": doc_data["content"],
                        "metadata": doc_data.get("metadata", {}),
                        "embedding": doc_data["embedding"].tolist()
                    })
            
            # Save backup
            Path(backup_path).parent.mkdir(parents=True, exist_ok=True)
            with open(backup_path, 'w') as f:
                json.dump(backup_data, f, indent=2, default=str)
            
            self.logger.info(f"Collection backed up to {backup_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"Backup failed: {e}")
            return False
    
    async def restore_collection(self, backup_path: str) -> bool:
        """
        Restore collection from backup.
        
        Args:
            backup_path: Path to backup file
            
        Returns:
            True if successful
        """
        try:
            if not Path(backup_path).exists():
                raise ToxoError(f"Backup file not found: {backup_path}")
            
            with open(backup_path, 'r') as f:
                backup_data = json.load(f)
            
            # Restore documents
            for doc_data in backup_data.get("documents", []):
                embedding = np.array(doc_data["embedding"]) if doc_data["embedding"] else None
                if embedding is not None:
                    await self.add_document(
                        content=doc_data["content"],
                        embedding=embedding,
                        metadata=doc_data.get("metadata", {}),
                        document_id=doc_data["id"]
                    )
            
            self.logger.info(f"Collection restored from {backup_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"Restore failed: {e}")
            return False
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state for serialization."""
        return {
            "collection_name": getattr(self, 'collection_name', 'unknown'),
            "chroma_available": getattr(self, 'chroma_client', None) is not None,
            "sql_available": getattr(self, 'sql_engine', None) is not None,
            "persist_directory": str(getattr(self, 'persist_directory', '')),
            "database_url": getattr(self, 'database_url', ''),
            "document_count": len(getattr(self, 'documents', {})),
            "system_type": "VectorStore"
        }


# Global vector store instances
_vector_stores: Dict[str, VectorStore] = {}

def get_vector_store(
    collection_name: str = "toxo_default",
    persist_directory: Optional[str] = None,
    database_url: Optional[str] = None
) -> VectorStore:
    """Get or create a vector store instance."""
    if collection_name not in _vector_stores:
        _vector_stores[collection_name] = VectorStore(
            collection_name=collection_name,
            persist_directory=persist_directory,
            database_url=database_url
        )
    
    return _vector_stores[collection_name] 