from collections.abc import Generator

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

from app.config import settings

settings.data_dir.mkdir(parents=True, exist_ok=True)

engine = create_engine(
    settings.database_url,
    connect_args={"check_same_thread": False},
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


def ensure_database_schema() -> None:
    inspector = inspect(engine)
    if "html_files" not in inspector.get_table_names():
        return

    existing_columns = {
        column["name"]
        for column in inspector.get_columns("html_files")
    }
    column_migrations = {
        "title": "ALTER TABLE html_files ADD COLUMN title VARCHAR(120)",
        "source": "ALTER TABLE html_files ADD COLUMN source VARCHAR(80)",
        "request_id": "ALTER TABLE html_files ADD COLUMN request_id VARCHAR(120)",
        "size_bytes": "ALTER TABLE html_files ADD COLUMN size_bytes INTEGER",
        "expires_at": "ALTER TABLE html_files ADD COLUMN expires_at DATETIME",
    }

    with engine.begin() as connection:
        for column_name, statement in column_migrations.items():
            if column_name not in existing_columns:
                connection.execute(text(statement))

        connection.execute(
            text(
                "CREATE INDEX IF NOT EXISTS ix_html_files_expires_at "
                "ON html_files (expires_at)"
            )
        )


def get_db() -> Generator:
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()