Zero-Knowledge-Platform/backend/app/services/records.py

from datetime import datetime, timezone

from bson import ObjectId
from fastapi import HTTPException, status
from motor.motor_asyncio import AsyncIOMotorDatabase
from pymongo import ReturnDocument
from pymongo.errors import DuplicateKeyError

from app.schemas.records import (
    ChildCreateRequest,
    ChildRecordResponse,
    ChildUpdateRequest,
    ParentCreateRequest,
    ParentRecordResponse,
    ParentUpdateRequest,
)

PARENTS_COLLECTION = "parents"
CHILDREN_COLLECTION = "children"


class RecordNotFoundError(HTTPException):
    def __init__(self, detail: str = "Record not found.") -> None:
        super().__init__(status_code=status.HTTP_404_NOT_FOUND, detail=detail)


class BadRequestError(HTTPException):
    def __init__(self, detail: str = "Invalid request.") -> None:
        super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail)


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _to_object_id(value: str, field_name: str) -> ObjectId:
    if not ObjectId.is_valid(value):
        raise BadRequestError(f"Invalid {field_name}.")
    return ObjectId(value)


def _serialize_parent(document: dict) -> ParentRecordResponse:
    return ParentRecordResponse(
        id=str(document["_id"]),
        ownerId=document["ownerId"],
        iv=document["iv"],
        ciphertext=document["ciphertext"],
        createdAt=document["createdAt"],
        updatedAt=document["updatedAt"],
    )


def _serialize_child(document: dict) -> ChildRecordResponse:
    return ChildRecordResponse(
        id=str(document["_id"]),
        parentId=str(document["parentId"]),
        ownerId=document["ownerId"],
        iv=document["iv"],
        ciphertext=document["ciphertext"],
        createdAt=document["createdAt"],
        updatedAt=document["updatedAt"],
    )


async def ensure_record_indexes(database: AsyncIOMotorDatabase) -> None:
    await database[PARENTS_COLLECTION].create_index([("ownerId", 1)], name="ix_parents_ownerId")
    await database[CHILDREN_COLLECTION].create_index([("ownerId", 1)], name="ix_children_ownerId")
    await database[CHILDREN_COLLECTION].create_index([("parentId", 1)], name="ix_children_parentId")
    await database[CHILDREN_COLLECTION].create_index(
        [("ownerId", 1), ("parentId", 1)],
        name="ix_children_ownerId_parentId",
    )


async def create_parent(database: AsyncIOMotorDatabase, owner_id: str, payload: ParentCreateRequest) -> ParentRecordResponse:
    now = _now()
    document = {
        "ownerId": owner_id,
        "iv": payload.iv,
        "ciphertext": payload.ciphertext,
        "createdAt": now,
        "updatedAt": now,
    }
    result = await database[PARENTS_COLLECTION].insert_one(document)
    document["_id"] = result.inserted_id
    return _serialize_parent(document)


async def list_parents(database: AsyncIOMotorDatabase, owner_id: str) -> list[ParentRecordResponse]:
    parents: list[ParentRecordResponse] = []
    cursor = database[PARENTS_COLLECTION].find({"ownerId": owner_id}).sort("createdAt", -1)
    async for document in cursor:
        parents.append(_serialize_parent(document))
    return parents


async def get_parent(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str) -> ParentRecordResponse:
    document = await database[PARENTS_COLLECTION].find_one(
        {"_id": _to_object_id(parent_id, "parentId"), "ownerId": owner_id}
    )
    if document is None:
        raise RecordNotFoundError("Parent record not found.")
    return _serialize_parent(document)


async def update_parent(
    database: AsyncIOMotorDatabase,
    owner_id: str,
    parent_id: str,
    payload: ParentUpdateRequest,
) -> ParentRecordResponse:
    document = await database[PARENTS_COLLECTION].find_one_and_update(
        {"_id": _to_object_id(parent_id, "parentId"), "ownerId": owner_id},
        {
            "$set": {
                "iv": payload.iv,
                "ciphertext": payload.ciphertext,
                "updatedAt": _now(),
            }
        },
        return_document=ReturnDocument.AFTER,
    )
    if document is None:
        raise RecordNotFoundError("Parent record not found.")
    return _serialize_parent(document)


async def delete_parent(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str) -> None:
    client = database.client
    if client is None:
        raise RuntimeError("MongoDB client has not been initialized")

    parent_object_id = _to_object_id(parent_id, "parentId")
    async with await client.start_session() as session:
        async with session.start_transaction():
            parent = await database[PARENTS_COLLECTION].find_one(
                {"_id": parent_object_id, "ownerId": owner_id},
                session=session,
            )
            if parent is None:
                raise RecordNotFoundError("Parent record not found.")

            await database[CHILDREN_COLLECTION].delete_many(
                {"parentId": parent_object_id, "ownerId": owner_id},
                session=session,
            )
            await database[PARENTS_COLLECTION].delete_one(
                {"_id": parent_object_id, "ownerId": owner_id},
                session=session,
            )


async def create_child(
    database: AsyncIOMotorDatabase,
    owner_id: str,
    payload: ChildCreateRequest,
) -> ChildRecordResponse:
    parent_object_id = _to_object_id(payload.parentId, "parentId")
    parent = await database[PARENTS_COLLECTION].find_one({"_id": parent_object_id, "ownerId": owner_id})
    if parent is None:
        raise RecordNotFoundError("Parent record not found.")

    now = _now()
    document = {
        "ownerId": owner_id,
        "parentId": parent_object_id,
        "iv": payload.iv,
        "ciphertext": payload.ciphertext,
        "createdAt": now,
        "updatedAt": now,
    }
    result = await database[CHILDREN_COLLECTION].insert_one(document)
    document["_id"] = result.inserted_id
    return _serialize_child(document)


async def list_children(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str) -> list[ChildRecordResponse]:
    parent_object_id = _to_object_id(parent_id, "parentId")
    children: list[ChildRecordResponse] = []
    cursor = database[CHILDREN_COLLECTION].find({"ownerId": owner_id, "parentId": parent_object_id}).sort("createdAt", -1)
    async for document in cursor:
        children.append(_serialize_child(document))
    return children


async def get_child(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str, child_id: str) -> ChildRecordResponse:
    document = await database[CHILDREN_COLLECTION].find_one(
        {
            "_id": _to_object_id(child_id, "childId"),
            "parentId": _to_object_id(parent_id, "parentId"),
            "ownerId": owner_id,
        }
    )
    if document is None:
        raise RecordNotFoundError("Child record not found.")
    return _serialize_child(document)


async def update_child(
    database: AsyncIOMotorDatabase,
    owner_id: str,
    parent_id: str,
    child_id: str,
    payload: ChildUpdateRequest,
) -> ChildRecordResponse:
    document = await database[CHILDREN_COLLECTION].find_one_and_update(
        {
            "_id": _to_object_id(child_id, "childId"),
            "parentId": _to_object_id(parent_id, "parentId"),
            "ownerId": owner_id,
        },
        {
            "$set": {
                "iv": payload.iv,
                "ciphertext": payload.ciphertext,
                "updatedAt": _now(),
            }
        },
        return_document=ReturnDocument.AFTER,
    )
    if document is None:
        raise RecordNotFoundError("Child record not found.")
    return _serialize_child(document)


async def delete_child(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str, child_id: str) -> None:
    result = await database[CHILDREN_COLLECTION].delete_one(
        {
            "_id": _to_object_id(child_id, "childId"),
            "parentId": _to_object_id(parent_id, "parentId"),
            "ownerId": owner_id,
        }
    )
    if result.deleted_count == 0:
        raise RecordNotFoundError("Child record not found.")