from collections.abc import Generator

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

DATABASE_URL = "sqlite:///./data/clawfort.db"

engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False)


class Base(DeclarativeBase):
    pass


def get_db() -> Generator[Session, None, None]:
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def init_db() -> None:
    from backend.models import NewsItem, NewsTranslation  # noqa: F401

    Base.metadata.create_all(bind=engine)

    migration_sql = {
        "news_items": {
            "tldr_points": "ALTER TABLE news_items ADD COLUMN tldr_points TEXT",
            "summary_body": "ALTER TABLE news_items ADD COLUMN summary_body TEXT",
            "source_citation": "ALTER TABLE news_items ADD COLUMN source_citation TEXT",
            "summary_image_url": "ALTER TABLE news_items ADD COLUMN summary_image_url VARCHAR(2000)",
            "summary_image_credit": "ALTER TABLE news_items ADD COLUMN summary_image_credit VARCHAR(500)",
        },
        "news_translations": {
            "tldr_points": "ALTER TABLE news_translations ADD COLUMN tldr_points TEXT",
            "summary_body": "ALTER TABLE news_translations ADD COLUMN summary_body TEXT",
            "source_citation": "ALTER TABLE news_translations ADD COLUMN source_citation TEXT",
        },
    }

    with engine.begin() as conn:
        for table, columns in migration_sql.items():
            existing_cols = {
                row[1] for row in conn.execute(text(f"PRAGMA table_info({table})")).fetchall()
            }
            for col_name, ddl in columns.items():
                if col_name not in existing_cols:
                    conn.execute(text(ddl))