from typing import List, Optional

from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.mysql import ArchiveMediaFiles


class ArchiveMediaFilesRepository:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def get_voice_files(self, limit: int = 100, offset: int = 0) -> List[ArchiveMediaFiles]:
        """
        获取语音文件列表

        Args:
            limit: 返回数量限制
            offset: 偏移量

        Returns:
            ArchiveMediaFiles 对象列表
        """
        query = select(ArchiveMediaFiles).where(
            ArchiveMediaFiles.file_type == "voice"
        ).limit(limit).offset(offset)

        result = await self.session.execute(query)
        return list(result.scalars().all())

    async def count_voice_files(self) -> int:
        """
        统计语音文件总数

        Returns:
            总数
        """
        query = select(func.count()).select_from(ArchiveMediaFiles).where(
            ArchiveMediaFiles.file_type == "voice"
        )

        result = await self.session.execute(query)
        return result.scalar() or 0

    async def get_by_archive_message_id(self, archive_message_id: int, file_type: str = "voice") -> Optional[ArchiveMediaFiles]:
        """
        根据 archive_message_id 查询媒体文件

        Args:
            archive_message_id: 归档消息ID
            file_type: 文件类型过滤，默认 voice

        Returns:
            ArchiveMediaFiles 对象或 None
        """
        query = select(ArchiveMediaFiles).where(
            ArchiveMediaFiles.archive_message_id == archive_message_id
        )

        if file_type:
            query = query.where(ArchiveMediaFiles.file_type == file_type)

        result = await self.session.execute(query)
        return result.scalar_one_or_none()