from __future__ import annotations

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.models.chat_message import ChatMessage
from app.models.chat_session import ChatSession


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

    def add(self, item: ChatSession) -> ChatSession:
        self.db.add(item)
        self.db.commit()
        self.db.refresh(item)
        return item

    def update(self, item: ChatSession) -> ChatSession:
        self.db.add(item)
        self.db.commit()
        self.db.refresh(item)
        return item

    def get_by_session_id(self, session_id: str) -> ChatSession | None:
        statement = select(ChatSession).where(ChatSession.session_id == session_id)
        return self.db.execute(statement).scalar_one_or_none()


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

    def add(self, item: ChatMessage) -> ChatMessage:
        self.db.add(item)
        self.db.commit()
        self.db.refresh(item)
        return item

    def list_by_session_id(self, session_id: str) -> list[ChatMessage]:
        statement = select(ChatMessage).where(ChatMessage.session_id == session_id).order_by(ChatMessage.created_at.asc())
        return list(self.db.execute(statement).scalars())