"""
Database module for the application.
Handles connections to PostgreSQL (Zulip DB) and ChromaDB.
"""
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy.ext.declarative import declarative_base
import chromadb

# SQLAlchemy base class for models
Base = declarative_base()

# Global variables for SQLAlchemy
db_engine = None
db_session = None

# Global variable for ChromaDB
chroma_client = None
chroma_collection = None

def init_db(app):
    """
    Initialize database connections.
    
    Args:
        app: Flask application object
    """
    global db_engine, db_session, chroma_client, chroma_collection
    
    # Initialize SQLAlchemy engine and session
    db_engine = create_engine(app.config['SQLALCHEMY_DATABASE_URI'])
    db_session = scoped_session(sessionmaker(autocommit=False, autoflush=False, bind=db_engine))
    
    # Set query property for models
    Base.query = db_session.query_property()
    
    # Initialize ChromaDB
    try:
        # Set allow_reset to True to prevent "Add of existing embedding ID" warnings
        chroma_client = chromadb.PersistentClient(
            path=app.config['CHROMADB_PATH'],
            settings=chromadb.Settings(
                allow_reset=True,
                anonymized_telemetry=False,
                is_persistent=True
            )
        )
        
        # Import here to avoid circular imports
        from app.db.chroma_service import CustomEmbeddingFunction
        
        # Create embedding function with setting from config
        try:
            # Always use Ollama since it's more reliable
            embedding_function = CustomEmbeddingFunction(use_nomic=False)
            
            # Get or create ChromaDB collection for Zulip messages with custom embedding function
            chroma_collection = chroma_client.get_or_create_collection(
                name=app.config.get('CHROMADB_COLLECTION', 'zulip_messages'),
                metadata={
                    "hnsw:space": "cosine",
                    "hnsw:allow_replace_deleted": True  # Allow replacing deleted vectors
                },
                embedding_function=embedding_function
            )
        except Exception as e:
            print(f"Error with embedding function: {e}")
            print("Creating collection without embedding function")
            # Create collection without embedding function
            chroma_collection = chroma_client.get_or_create_collection(
                name=app.config.get('CHROMADB_COLLECTION', 'zulip_messages'),
                metadata={
                    "hnsw:space": "cosine",
                    "hnsw:allow_replace_deleted": True  # Allow replacing deleted vectors
                }
            )
    except Exception as e:
        print(f"Critical error initializing ChromaDB: {e}")
        print("ChromaDB functionality will not be available")
        chroma_client = None
        chroma_collection = None
    
    # Register teardown function to remove database sessions
    @app.teardown_appcontext
    def shutdown_session(exception=None):
        """Remove the database session at the end of the request."""
        db_session.remove()

def get_db_session():
    """
    Get the current database session.
    
    Returns:
        SQLAlchemy session object
    """
    return db_session

def get_chroma_collection():
    """
    Get the ChromaDB collection for Zulip messages.
    
    Returns:
        ChromaDB collection object
    """
    return chroma_collection