from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from typing import Optional

from db.models import Order, OrderStatus


class OrderRepository:
    async def get_order_by_id(
        self, session: AsyncSession, order_id: int
    ) -> Optional[Order]:
        stmt = select(Order).where(Order.id == order_id)
        return await session.scalar(stmt)

    async def get_orders_by_user(
        self, session: AsyncSession, customer: int
    ) -> list[Order]:
        stmt = select(Order).where(Order.customer == customer)
        result = await session.scalars(stmt)

        return list(result)

    async def get_draft_order_by_user(
        self, session: AsyncSession, customer: int
    ) -> Optional[Order]:
        stmt = (
            select(Order)
            .where(Order.customer == customer)
            .where(Order.status == OrderStatus.DRAFT)
            .limit(1)
        )
        return await session.scalar(stmt)

    async def create_order(
        self,
        session: AsyncSession,
        customer: int,
        status: OrderStatus = OrderStatus.DRAFT,
    ):
        order = Order(customer=customer, status=status)

        session.add(order)
        await session.commit()

        return order