foundation/src/Core/Database/UnitOfWork.php

<?php

declare(strict_types=1);

namespace Foundation\Core\Database;

use PDO;
use PDOException;
use RuntimeException;
use SplObjectStorage;

class UnitOfWork implements UnitOfWorkInterface
{
    private SplObjectStorage $entityStates;
    private array $newEntities = [];
    private array $dirtyEntities = [];
    private array $removedEntities = [];
    private bool $inTransaction = false;
    /** @var EntityPersisterInterface[] */
    private array $persisters = [];

    public function __construct(private readonly PDO $pdo)
    {
        $this->entityStates = new SplObjectStorage();
    }

    public function registerNew(object $entity): void
    {
        $this->entityStates[$entity] = EntityState::NEW;
        $this->newEntities[spl_object_id($entity)] = $entity;
        $this->removeFromOtherCollections($entity, 'new');
    }

    public function registerDirty(object $entity): void
    {
        if (!isset($this->entityStates[$entity]) || $this->entityStates[$entity] === EntityState::CLEAN) {
            $this->entityStates[$entity] = EntityState::DIRTY;
            $this->dirtyEntities[spl_object_id($entity)] = $entity;
            $this->removeFromOtherCollections($entity, 'dirty');
        }
    }

    public function registerRemoved(object $entity): void
    {
        $this->entityStates[$entity] = EntityState::REMOVED;
        $this->removedEntities[spl_object_id($entity)] = $entity;
        $this->removeFromOtherCollections($entity, 'removed');
    }

    public function registerClean(object $entity): void
    {
        $this->entityStates[$entity] = EntityState::CLEAN;
        $this->removeFromOtherCollections($entity, 'clean');
    }

    public function commit(): void
    {
        if ($this->newEntities === [] && $this->dirtyEntities === [] && $this->removedEntities === []) {
            return;
        }

        try {
            $this->pdo->beginTransaction();
            $this->inTransaction = true;

            $this->processRemovedEntities();
            $this->processNewEntities();
            $this->processDirtyEntities();

            $this->pdo->commit();
            $this->inTransaction = false;

            $this->markAllAsClean();
            $this->clear();
        } catch (PDOException $e) {
            $this->rollback();
            throw $e;
        }
    }

    public function rollback(): void
    {
        if ($this->inTransaction) {
            $this->pdo->rollBack();
            $this->inTransaction = false;
        }
    }

    public function clear(): void
    {
        $this->newEntities = [];
        $this->dirtyEntities = [];
        $this->removedEntities = [];
    }

    public function getEntityState(object $entity): EntityState
    {
        return $this->entityStates[$entity] ?? EntityState::CLEAN;
    }

    public function isInTransaction(): bool
    {
        return $this->inTransaction;
    }

    public function registerPersister(EntityPersisterInterface $persister): void
    {
        $this->persisters[] = $persister;
    }

    private function removeFromOtherCollections(object $entity, string $except): void
    {
        $objectId = spl_object_id($entity);

        if ($except !== 'new') {
            unset($this->newEntities[$objectId]);
        }
        if ($except !== 'dirty') {
            unset($this->dirtyEntities[$objectId]);
        }
        if ($except !== 'removed') {
            unset($this->removedEntities[$objectId]);
        }
    }

    private function processNewEntities(): void
    {
        foreach ($this->newEntities as $entity) {
            $this->persistNewEntity($entity);
        }
    }

    private function processDirtyEntities(): void
    {
        foreach ($this->dirtyEntities as $entity) {
            $this->updateDirtyEntity($entity);
        }
    }

    private function processRemovedEntities(): void
    {
        foreach ($this->removedEntities as $entity) {
            $this->removeEntity($entity);
        }
    }

    private function persistNewEntity(object $entity): void
    {
        $persister = $this->getPersisterForEntity($entity);
        $persister->insert($entity, $this->pdo);
    }

    private function updateDirtyEntity(object $entity): void
    {
        $persister = $this->getPersisterForEntity($entity);
        $persister->update($entity, $this->pdo);
    }

    private function removeEntity(object $entity): void
    {
        $persister = $this->getPersisterForEntity($entity);
        $persister->delete($entity, $this->pdo);
    }

    private function getPersisterForEntity(object $entity): EntityPersisterInterface
    {
        foreach ($this->persisters as $persister) {
            if ($persister->supports($entity)) {
                return $persister;
            }
        }

        throw new RuntimeException(
            sprintf('No persister found for entity of type %s', $entity::class)
        );
    }

    private function markAllAsClean(): void
    {
        foreach ($this->newEntities as $entity) {
            $this->entityStates[$entity] = EntityState::CLEAN;
        }
        foreach ($this->dirtyEntities as $entity) {
            $this->entityStates[$entity] = EntityState::CLEAN;
        }
        foreach ($this->removedEntities as $entity) {
            unset($this->entityStates[$entity]);
        }
    }
}