/*
 * Copyright (C) 2020 Grakn Labs
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 */

package hypergraph.core;

import hypergraph.Hypergraph;
import hypergraph.common.exception.HypergraphException;
import hypergraph.graph.util.KeyGenerator;
import org.rocksdb.OptimisticTransactionDB;
import org.rocksdb.RocksDBException;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Comparator;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

public class CoreKeyspace implements Hypergraph.Keyspace {

    private final String name;
    private final CoreHypergraph core;
    private final OptimisticTransactionDB rocksDB;
    private final KeyGenerator.Persisted keyGenerator;
    private final Set<CoreSession> sessions;
    private final AtomicBoolean isOpen;

    private CoreKeyspace(CoreHypergraph core, String name) {
        this.name = name;
        this.core = core;
        keyGenerator = new KeyGenerator.Persisted();
        sessions = ConcurrentHashMap.newKeySet();
        isOpen = new AtomicBoolean(false);

        try {
            rocksDB = OptimisticTransactionDB.open(this.core.rocksOptions(), directory().toString());
        } catch (RocksDBException e) {
            throw new HypergraphException(e);
        }
    }

    static CoreKeyspace createNewAndOpen(CoreHypergraph core, String name) {
        return new CoreKeyspace(core, name).initialiseAndOpen();
    }

    static CoreKeyspace loadExistingAndOpen(CoreHypergraph core, String name) {
        return new CoreKeyspace(core, name).loadAndOpen();
    }

    private CoreKeyspace initialiseAndOpen() {
        try (CoreSession session = createAndOpenSession(Hypergraph.Session.Type.SCHEMA)) {
            try (CoreTransaction txn = session.transaction(Hypergraph.Transaction.Type.WRITE)) {
                if (txn.graph().isInitialised()) {
                    throw new HypergraphException("Invalid Keyspace Initialisation");
                }
                txn.graph().initialise();
                txn.commit();
            }
        }
        isOpen.set(true);
        return this;
    }

    private CoreKeyspace loadAndOpen() {
        try (CoreSession session = createAndOpenSession(Hypergraph.Session.Type.DATA)) {
            try (CoreTransaction txn = session.transaction(Hypergraph.Transaction.Type.READ)) {
                keyGenerator.sync(txn.storage());
            }
        }
        isOpen.set(true);
        return this;
    }

    CoreSession createAndOpenSession(Hypergraph.Session.Type type) {
        CoreSession session = new CoreSession(this, type);
        sessions.add(session);
        return session;
    }

    private Path directory() {
        return core.directory().resolve(name);
    }

    OptimisticTransactionDB rocks() {
        return rocksDB;
    }

    KeyGenerator keyGenerator() {
        return keyGenerator;
    }

    public void remove(CoreSession session) {
        sessions.remove(session);
    }

    void close() {
        if (isOpen.compareAndSet(true, false)) {
            sessions.parallelStream().forEach(CoreSession::close);
            rocksDB.close();
        }
    }

    @Override
    public String name() {
        return name;
    }

    @Override
    public void delete() {
        close();
        core.keyspaces().remove(this);
        try {
            Files.walk(directory()).sorted(Comparator.reverseOrder())
                    .map(Path::toFile).forEach(File::delete);
        } catch (IOException e) {
            throw new HypergraphException(e);
        }
    }
}

@@ -764,16 +764,74 @@
 rgraph;%0A
+import hypergraph.common.concurrent.ManagedReadWriteLock;%0A
 import h
@@ -1205,16 +1205,59 @@
 ashMap;%0A
+import java.util.concurrent.ConcurrentMap;%0A
 import j
@@ -1580,16 +1580,67 @@
 ssions;%0A
+    private final ManagedReadWriteLock schemaLock;%0A
     priv
@@ -1930,16 +1930,65 @@
 (false);
+%0A        schemaLock = new ManagedReadWriteLock();
 %0A%0A      
@@ -3453,24 +3453,317 @@
 ype type) %7B%0A
+        try %7B%0A            if (type.equals(Hypergraph.Session.Type.DATA)) %7B%0A                schemaLock.lockRead();%0A            %7D else if (type.equals(Hypergraph.Session.Type.SCHEMA)) %7B%0A                schemaLock.lockWrite();%0A            %7D else %7B%0A                assert false;%0A            %7D%0A    
         Core
@@ -3809,32 +3809,36 @@
  type);%0A        
+ 
+   
 sessions.add(ses
@@ -3848,24 +3848,28 @@
 n);%0A        
+    
 return sessi
@@ -3872,16 +3872,115 @@
 ession;%0A
+        %7D catch (InterruptedException e) %7B%0A            throw new HypergraphException(e);%0A        %7D%0A
     %7D%0A%0A 
@@ -4278,24 +4278,295 @@
 e(session);%0A
+        if (session.type().equals(Hypergraph.Session.Type.DATA)) %7B%0A            schemaLock.unlockRead();%0A        %7D else if (session.type().equals(Hypergraph.Session.Type.SCHEMA)) %7B%0A            schemaLock.unlockWrite();%0A        %7D else %7B%0A            assert false;%0A        %7D%0A
     %7D%0A%0A    v

Introduce SchemaLock in CoreKeyspace