import asyncio
from contextlib import asynccontextmanager
from collections.abc import AsyncIterator
from sqlalchemy import event
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from .models import (
    provider as provider_models,
    agent as agent_models,
    workspace as workspace_models,
    toolset as toolset_models,
    skill as skill_models,
)
# this unused import is necessary to alembic
from . import models
from src.common import DATA_DIR

db_path = DATA_DIR / "sqlite.db"

DB_ASYNC_URL = f"sqlite+aiosqlite:///{db_path}"
DB_SYNC_URL = f"sqlite:///{db_path}"

engine = create_async_engine(DB_ASYNC_URL)
AsyncSessionLocal = async_sessionmaker(engine,
                                       autoflush=False,
                                       autocommit=False,
                                       expire_on_commit=False)

async def get_db_session() -> AsyncIterator[AsyncSession]:
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except asyncio.CancelledError:
            pass
        except Exception:
            await session.rollback()
            raise
db_context = asynccontextmanager(get_db_session)

@event.listens_for(engine.sync_engine, "connect")
def on_connect(dbapi_conn, _):
    dbapi_conn.execute("PRAGMA foreign_keys=ON")
    dbapi_conn.execute("PRAGMA journal_mode=WAL")
    dbapi_conn.execute("PRAGMA busy_timeout=3000")

async def init_initial_data() -> None:
    async with AsyncSessionLocal.begin() as session:
        await provider_models.init(session)
        await toolset_models.init(session)
        await agent_models.init(session)
        await skill_models.init(session)
        await workspace_models.init(session)

async def migrate_db() -> None:
    from alembic.config import Config
    from alembic import command
    from src.common import PROJECT_ROOT

    alembic_cfg = Config("alembic.ini")
    alembic_cfg.set_main_option("script_location", f"{PROJECT_ROOT}/src/db/alembic")
    alembic_cfg.set_main_option("sqlalchemy.url", DB_SYNC_URL)
    command.upgrade(alembic_cfg, "head")
    await init_initial_data()
