from __future__ import annotations

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.models.task import Task


class TaskRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def add(self, task: Task) -> Task:
        self.db.add(task)
        self.db.commit()
        self.db.refresh(task)
        return task

    def update(self, task: Task) -> Task:
        self.db.add(task)
        self.db.commit()
        self.db.refresh(task)
        return task

    def get_by_task_id(self, task_id: str) -> Task | None:
        statement = select(Task).where(Task.task_id == task_id)
        return self.db.execute(statement).scalar_one_or_none()

    def get_by_request_id(self, request_id: str) -> Task | None:
        statement = select(Task).where(Task.request_id == request_id)
        return self.db.execute(statement).scalar_one_or_none()