from uuid import UUID

from domain.workspace import Workspace
from usecase.interface import WorkspaceRepository


class InMemoryWorkspaceRepository(WorkspaceRepository):
    def __init__(self) -> None:
        self._by_id: dict[UUID, Workspace] = {}
        self._user_id_to_workspace_id: dict[UUID, UUID] = {}

    def get(self, workspace_id: UUID) -> Workspace | None:
        return self._by_id.get(workspace_id)

    def get_by_user_id(self, user_id: UUID) -> Workspace | None:
        wid = self._user_id_to_workspace_id.get(user_id)
        if wid is None:
            return None
        return self._by_id.get(wid)

    def save(self, workspace: Workspace) -> None:
        existing_wid = self._user_id_to_workspace_id.get(workspace.user_id)
        if existing_wid is not None and existing_wid != workspace.workspace_id:
            self._by_id.pop(existing_wid, None)
        self._by_id[workspace.workspace_id] = workspace
        self._user_id_to_workspace_id[workspace.user_id] = workspace.workspace_id