package com.atlassian.stash.internal.spring;

import org.springframework.transaction.support.TransactionSynchronizationAdapter;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.util.HashMap;
import java.util.Map;

import static com.google.common.base.Preconditions.checkState;

/**
 * {@link TransactionalState} implementation that stores the state as a transactional resource
 * through Spring's {@link TransactionSynchronizationManager}.
 */
@SuppressWarnings("unchecked")
class DefaultTransactionalState<K, V> implements TransactionalState<K, V> {
    private final Object namespace;

    DefaultTransactionalState(Object namespace) {
        this.namespace = namespace;
    }

    @Override
    public V get(K key) {
        Object resource = TransactionSynchronizationManager.getResource(namespace);
        if (resource instanceof Map) {
            Map<K, V> txMap = (Map<K, V>) resource;
            return txMap.get(key);
        }
        return null;
    }

    @Override
    public V put(K key, V value) {
        checkState(TransactionSynchronizationManager.isSynchronizationActive(), "There is no active transaction");
        Object resource = TransactionSynchronizationManager.getResource(namespace);
        Map<K, V> txMap;
        if (resource == null) {
            // does not have to be a concurrent map, is only going to be used on the current thread anyhow
            txMap = new HashMap<>();
            TransactionSynchronizationManager.bindResource(namespace, txMap);
            TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronizationAdapter() {
                @Override
                public void afterCompletion(int status) {
                    clear();
                }
            });
        } else if (resource instanceof Map) {
            txMap = (Map<K, V>) resource;
        } else {
            throw new IllegalStateException("The transactional resource is not a map! (" + resource.getClass().getName() + ")");
        }

        return txMap.put(key, value);
    }

    public void clear() {
        TransactionSynchronizationManager.unbindResourceIfPossible(namespace);
    }
}
