from dataclasses import dataclass
from uuid import uuid4

from app.modules.chat.repository import ChatRepository


@dataclass
class DialogSession:
    dialog_session_id: str
    rag_session_id: str


class DialogSessionStore:
    def __init__(self, repository: ChatRepository) -> None:
        self._repo = repository

    def create(self, rag_session_id: str) -> DialogSession:
        session = DialogSession(dialog_session_id=str(uuid4()), rag_session_id=rag_session_id)
        self._repo.create_dialog(session.dialog_session_id, session.rag_session_id)
        return session

    def get(self, dialog_session_id: str) -> DialogSession | None:
        row = self._repo.get_dialog(dialog_session_id)
        if not row:
            return None
        return DialogSession(
            dialog_session_id=str(row["dialog_session_id"]),
            rag_session_id=str(row["rag_session_id"]),
        )