from sqlalchemy import select

from db.models.invoice import Invoice
from db.session import UnitOfWork
from schemas.invoices import InvoiceStatus


class InvoiceRepository:
    def __init__(self, uow: UnitOfWork) -> None:
        self.uow = uow
        self.session = uow.session

    async def get_by_id(self, id: int) -> Invoice | None:
        stmt = select(Invoice).where(Invoice.id == id)
        r = await self.session.execute(stmt)

        return r.scalar_one_or_none()

    async def get_by_user_id(self, user_id: int) -> list[Invoice]:
        stmt = select(Invoice).where(Invoice.creator_id == user_id)
        r = await self.session.execute(stmt)

        return list(r.scalars().all())

    async def create(
        self, creator_id: int, order_id: int, amount: int | float, status: InvoiceStatus
    ) -> Invoice:
        obj = Invoice(
            creator_id=creator_id,
            order_id=order_id,
            amount=amount,
            status=status,
        )

        self.session.add(obj)
        return obj

    async def update_status_by_id(self, invoice_id: int, status: InvoiceStatus) -> Invoice | None:
        invoice = await self.get_by_id(invoice_id)
        invoice.status = status
        return invoice