from collections.abc import Iterator
from contextlib import contextmanager

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

from app.config import settings
from app.db.models import Base

engine = create_engine(settings.db_url, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(bind=engine, expire_on_commit=False)

# Columns added after the table already existed in a deployed database.
# `create_all` only creates missing *tables*, never adds columns to an
# existing one -- there's no Alembic here, so new nullable columns are
# self-migrated by hand on startup instead. Only ever append to this list.
_ADDED_COLUMNS: list[tuple[str, str, str]] = [
    # (table, column, DDL type)
    ("interval_reading", "battery_soc_percent", "FLOAT"),
    ("interval_reading", "battery_charge_w", "FLOAT"),
    ("interval_reading", "battery_discharge_w", "FLOAT"),
]


def _apply_column_migrations() -> None:
    inspector = inspect(engine)
    if "interval_reading" not in inspector.get_table_names():
        return  # fresh DB -- create_all already created the table with every current column
    existing = {(table, col["name"]) for table in {t for t, _, _ in _ADDED_COLUMNS} for col in inspector.get_columns(table)}
    with engine.begin() as conn:
        for table, column, ddl_type in _ADDED_COLUMNS:
            if (table, column) not in existing:
                conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {column} {ddl_type}"))


def init_db() -> None:
    Base.metadata.create_all(engine)
    _apply_column_migrations()


@contextmanager
def get_session() -> Iterator[Session]:
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()