"""SQLModel database setup."""
from __future__ import annotations

import logging
import os
from collections.abc import Generator

from sqlalchemy import inspect, text
from sqlmodel import Session, SQLModel, create_engine

from config import settings

logger = logging.getLogger("stackpilot.database")

os.makedirs(settings.DATA_DIR, exist_ok=True)
_DB_PATH = os.path.join(settings.DATA_DIR, "stackpilot.db")
_DB_URL = f"sqlite:///{_DB_PATH}"

engine = create_engine(
    _DB_URL,
    echo=False,
    connect_args={"check_same_thread": False},
)


def _ensure_model_columns() -> None:
    """Add columns that models define but a pre-existing table is missing.

    ``SQLModel.create_all`` creates missing *tables* but never ALTERs an
    existing one, so installs that predate a newly-added column keep the old
    schema — and every ORM query that names the column fails with
    ``OperationalError: no such column``. For each mapped table we diff the
    model's columns against the live table and ``ADD COLUMN`` the safe
    (nullable, or defaulted) ones. Idempotent: on a fresh DB create_all already
    made every column, so this is a no-op.
    """
    insp = inspect(engine)
    live_tables = set(insp.get_table_names())
    with engine.begin() as conn:
        for table_name, table in SQLModel.metadata.tables.items():
            if table_name not in live_tables:
                continue
            existing = {c["name"] for c in insp.get_columns(table_name)}
            for col in table.columns:
                if col.name in existing:
                    continue
                # SQLite can only ADD a NOT NULL column if it has a default to
                # backfill existing rows; skip the rest rather than crash.
                if not col.nullable and col.default is None and col.server_default is None:
                    logger.warning(
                        "Cannot auto-add non-nullable column %s.%s (no default); "
                        "manual migration needed", table_name, col.name
                    )
                    continue
                ddl_type = col.type.compile(dialect=engine.dialect)
                conn.execute(text(f'ALTER TABLE "{table_name}" ADD COLUMN "{col.name}" {ddl_type}'))
                logger.info("Schema migration: added column %s.%s", table_name, col.name)


def init_db() -> None:
    # Import models so they are registered on SQLModel.metadata.
    import models  # noqa: F401

    SQLModel.metadata.create_all(engine)
    _ensure_model_columns()


def get_session() -> Generator[Session, None, None]:
    with Session(engine) as session:
        yield session