"""Database setup and session management with singleton registry pattern."""

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from contextlib import contextmanager
from typing import Generator

from .config import settings
from .database_registry import registry, get_base

# Get the singleton Base from registry
Base = get_base()

# Create database engine
engine = create_engine(
    settings.DATABASE_URL,
    connect_args={"check_same_thread": False} if settings.DATABASE_URL.startswith("sqlite") else {},
    echo=settings.DEBUG,
)

# Create session factory
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def get_db() -> Generator[Session, None, None]:
    """
    Dependency for getting database session.
    
    Yields:
        Database session
    """
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


@contextmanager
def get_db_context() -> Generator[Session, None, None]:
    """
    Context manager for database session.
    
    Yields:
        Database session
    """
    db = SessionLocal()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


def init_db() -> None:
    """Initialize database with all tables."""
    # Import all models to register them with Base
    from backend.models import (
        User, RefreshToken, APIKey,
        EmailVerificationToken, PasswordResetToken,
        Summary, ExportHistory
    )
    
    # Use registry to create tables safely
    registry.create_all_tables(engine)


def drop_db() -> None:
    """Drop all database tables. Use with caution!"""
    registry.drop_all_tables(engine)


def reset_db() -> None:
    """Reset database by dropping and recreating all tables."""
    drop_db()
    init_db()


def get_test_db(db_url: str = "sqlite:///./test.db") -> tuple:
    """
    Create a test database configuration.
    
    Returns:
        Tuple of (engine, SessionLocal, Base)
    """
    test_engine = create_engine(
        db_url,
        connect_args={"check_same_thread": False} if db_url.startswith("sqlite") else {},
        echo=False,
    )
    
    TestSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=test_engine)
    
    return test_engine, TestSessionLocal, Base