from typing import Any, Dict, List

from uuid import UUID

from sqlalchemy.orm import selectinload
from sqlmodel import select
from sqlmodel.ext.asyncio.session import AsyncSession

from service.crud.base import CRUDBase, ModelType
from service.models.models import Topic, TopicBase, TopicCreate, TopicModel, Word


class CRUDTopic(CRUDBase[Topic, TopicCreate, TopicBase]):
    async def get_model_topics(
        self, db: AsyncSession, *, model_id: UUID, version: int, with_words: bool = False
    ) -> List[ModelType]:
        statement = (
            select(self.model)
            .join(TopicModel)
            .filter(TopicModel.model_id == model_id, TopicModel.version == version)
            .order_by(-self.model.count)
        )
        if with_words:
            statement = statement.options(selectinload(self.model.top_words))
        return (await db.execute(statement)).scalars().all()

    async def save_topics(
        self, db: AsyncSession, *, topics: List[Dict[str, Any]], model: TopicModel
    ) -> None:
        for topic in topics:
            db_obj = self.model.from_orm(
                TopicCreate.parse_obj({**topic, "topic_model_id": model.id})
            )
            db.add(db_obj)
            for word in topic["top_words"]:
                db.add(Word.parse_obj({**word, "topic": db_obj}))

        await db.commit()


topic = CRUDTopic(Topic)

sqlmodel.select sqlalchemy.orm.selectinload service.models.models.Word.parse_obj service.models.models.TopicCreate.parse_obj
[((805, 839), 'sqlalchemy.orm.selectinload', 'selectinload', (['self.model.top_words'], {}), '(self.model.top_words)\n', (817, 839), False, 'from sqlalchemy.orm import selectinload\n'), ((1115, 1175), 'service.models.models.TopicCreate.parse_obj', 'TopicCreate.parse_obj', (["{**topic, 'topic_model_id': model.id}"], {}), "({**topic, 'topic_model_id': model.id})\n", (1136, 1175), False, 'from service.models.models import Topic, TopicBase, TopicCreate, TopicModel, Word\n'), ((1284, 1325), 'service.models.models.Word.parse_obj', 'Word.parse_obj', (["{**word, 'topic': db_obj}"], {}), "({**word, 'topic': db_obj})\n", (1298, 1325), False, 'from service.models.models import Topic, TopicBase, TopicCreate, TopicModel, Word\n'), ((556, 574), 'sqlmodel.select', 'select', (['self.model'], {}), '(self.model)\n', (562, 574), False, 'from sqlmodel import select\n')]