from typing import List

from sqlalchemy import create_engine
from sqlalchemy.orm import Session

from models.orm import Prompt


class DB:
    def __init__(self, db: str = "sqlite:///example.db"):
        self.engine = create_engine(db, connect_args={"check_same_thread": False})

    def get_prompts(self) -> List[Prompt]:
        with Session(self.engine) as session:
            return session.query(Prompt).all()

    def get_prompt_by_id(self, prompt_id: int) -> Prompt | None:
        with Session(self.engine) as session:
            return session.query(Prompt).filter(Prompt.id == prompt_id).first()