"""Engine/session helpers. The DB URL comes from Env.database_url so the same
code runs on SQLite (default) or Postgres by changing one env var."""
from contextlib import contextmanager

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import sessionmaker

from Env import database_url
from src.db.models import Base

_engine = create_engine(database_url, future=True)
SessionLocal = sessionmaker(bind=_engine, future=True, expire_on_commit=False)

# Columns added after the table first shipped. create_all() never ALTERs an
# existing table, so we add any missing ones in place (keeps the local-scan /
# iTunes / MBID-cache data instead of forcing a rebuild).
_ADDED_COLUMNS = {
    "track": [
        ("liked", "BOOLEAN DEFAULT 0"),
        ("from_playlist", "BOOLEAN DEFAULT 0"),
        ("from_top", "BOOLEAN DEFAULT 0"),
    ],
}


def get_engine():
    return _engine


def _migrate(engine):
    inspector = inspect(engine)
    existing_tables = set(inspector.get_table_names())
    with engine.begin() as conn:
        for table, columns in _ADDED_COLUMNS.items():
            if table not in existing_tables:
                continue
            have = {c["name"] for c in inspector.get_columns(table)}
            for name, ddl in columns:
                if name not in have:
                    conn.execute(text(f'ALTER TABLE "{table}" ADD COLUMN {name} {ddl}'))


def init_db():
    """Create the storage dir + tables if missing (safe to call every run)."""
    from Env import data_dir

    data_dir.mkdir(parents=True, exist_ok=True)
    _migrate(_engine)
    Base.metadata.create_all(_engine)


@contextmanager
def session_scope():
    """Transactional session: commit on success, roll back on error."""
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()
