"""Extraction repository for database operations."""


from sqlalchemy import desc
from sqlmodel import select
from sqlmodel.ext.asyncio.session import AsyncSession

from app.models.extraction import Extraction
from app.repositories.base import BaseRepository


class ExtractionRepository(BaseRepository[Extraction]):
    """Repository for extraction database operations."""

    def __init__(self, session: AsyncSession) -> None:
        """Initialize the extraction repository."""
        super().__init__(Extraction, session)

    async def get_by_service_and_id(
        self, service: str, service_id: str,
    ) -> Extraction | None:
        """Get an extraction by service and service_id."""
        result = await self.session.exec(
            select(Extraction).where(
                Extraction.service == service, Extraction.service_id == service_id,
            ),
        )
        return result.first()

    async def get_by_user(self, user_id: int) -> list[Extraction]:
        """Get all extractions for a user."""
        result = await self.session.exec(
            select(Extraction)
            .where(Extraction.user_id == user_id)
            .order_by(desc(Extraction.created_at)),
        )
        return list(result.all())

    async def get_pending_extractions(self) -> list[Extraction]:
        """Get all pending extractions."""
        result = await self.session.exec(
            select(Extraction)
            .where(Extraction.status == "pending")
            .order_by(Extraction.created_at),
        )
        return list(result.all())

    async def get_extractions_by_status(self, status: str) -> list[Extraction]:
        """Get extractions by status."""
        result = await self.session.exec(
            select(Extraction)
            .where(Extraction.status == status)
            .order_by(desc(Extraction.created_at)),
        )
        return list(result.all())