from collections.abc import AsyncGenerator

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase

from app.config import settings


class Base(DeclarativeBase):
    pass


settings.data_dir.mkdir(parents=True, exist_ok=True)
_db_path = settings.data_dir / "app.db"
DATABASE_URL = f"sqlite+aiosqlite:///{_db_path}"


engine = create_async_engine(
    DATABASE_URL,
    echo=False,
    future=True,
)
AsyncSessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)


async def get_session() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSessionLocal() as session:
        yield session


async def init_db() -> None:
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
        res = await conn.execute(text("PRAGMA table_info(jobs)"))
        cols = {row[1] for row in res.fetchall()}
        if "pipeline_timings_json" not in cols:
            await conn.execute(text("ALTER TABLE jobs ADD COLUMN pipeline_timings_json TEXT"))
        if "activity_log_json" not in cols:
            await conn.execute(text("ALTER TABLE jobs ADD COLUMN activity_log_json TEXT"))
        if "whisper_language" not in cols:
            await conn.execute(text("ALTER TABLE jobs ADD COLUMN whisper_language VARCHAR(32)"))
        if "whisper_task" not in cols:
            await conn.execute(text("ALTER TABLE jobs ADD COLUMN whisper_task VARCHAR(16) DEFAULT 'transcribe'"))
        if "title" not in cols:
            await conn.execute(text("ALTER TABLE jobs ADD COLUMN title VARCHAR(256)"))
        if "subject" not in cols:
            await conn.execute(text("ALTER TABLE jobs ADD COLUMN subject VARCHAR(128)"))
        if "thumbnail" not in cols:
            await conn.execute(text("ALTER TABLE jobs ADD COLUMN thumbnail VARCHAR(512)"))