package simpledb;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

public class LockManager {

    private Map<PageId, Set<TransactionId>> pageShLocks;
    private Map<PageId, TransactionId> pageExLocks;
    private Map<TransactionId, Set<PageId>> txLocks;
    private Map<TransactionId, Set<TransactionId>> deadlockMap;

    public LockManager() {
        this.pageExLocks = new ConcurrentHashMap<>();
        this.pageShLocks = new ConcurrentHashMap();
        this.txLocks = new ConcurrentHashMap<>();
        this.deadlockMap = new ConcurrentHashMap<>();
    }

    /**
     * Acquires a shared/exclusive lock on a page for a given transaction
     * Adds an entry mapping the tid -> {pid} reflecting the tx holds a lock on the page
     * @param tid the transaction id
     * @param pid the page id for which to lock
     * @param perm the permissions to determine if the lock should be shared or exlcusive
     */
    public void acquirePageLock(TransactionId tid, PageId pid, Permissions perm) throws TransactionAbortedException, InterruptedException {
        System.out.println("Tx" + tid + " acquiring " + perm + " lock");
        if (perm == Permissions.READ_ONLY) {
            acquireSharedLock(tid, pid);
        } else {
            acquireExclusiveLock(tid, pid);
        }

        System.out.println("Tx" + tid + " acquired " + perm + " lock");
        // acquired lock
        this.deadlockMap.remove(tid);

        // Maintain the which tx -> page lock entry
        if (!this.txLocks.containsKey(tid)) {
            this.txLocks.put(tid, new HashSet<>());
        }
        this.txLocks.get(tid).add(pid);
    }

    public synchronized void acquireSharedLock(TransactionId tid, PageId pid) throws InterruptedException, TransactionAbortedException {
        if (this.pageShLocks.containsKey(pid) && this.pageShLocks.get(pid).contains(tid)) {
            return;
        }
        while (this.pageExLocks.containsKey(pid) && !this.pageExLocks.get(pid).equals(tid)) {
            Set<TransactionId> dependencies = new HashSet<>();
            dependencies.add(this.pageExLocks.get(pid));
            mapDependency(tid, pid, dependencies);
            wait();
        }
        grabSharedLock(tid, pid);
    }

    public synchronized void acquireExclusiveLock(TransactionId tid, PageId pid) throws TransactionAbortedException, InterruptedException {
        if (this.pageExLocks.containsKey(pid) && this.pageExLocks.get(pid).equals(tid)) {
            return;
        }
        Set<TransactionId> dependencies = new HashSet<>();

        while ((this.pageExLocks.containsKey(pid) && !this.pageExLocks.get(pid).equals(tid)) ||
                (this.pageShLocks.containsKey(pid) &&
                        ((this.pageShLocks.get(pid).contains(tid) && this.pageShLocks.get(pid).size() > 1) ||
                        (!this.pageShLocks.get(pid).contains(tid) && this.pageShLocks.get(pid).size() > 0)))) {

            if (this.pageExLocks.containsKey(pid) && !this.pageExLocks.get(pid).equals(tid)) {
                dependencies.add(this.pageExLocks.get(pid));
            }
            if (this.pageShLocks.containsKey(pid)) {
                for (TransactionId depIds : this.pageShLocks.get(pid)) {
                    dependencies.add(depIds);
                }
            }
            mapDependency(tid, pid, dependencies);
            wait();
        }
        grabExLock(tid, pid);

    }

    private void grabSharedLock(TransactionId tid, PageId pid) {
        if (!this.pageShLocks.containsKey(pid)) {
            this.pageShLocks.put(pid, new HashSet<>());
        }
        this.pageShLocks.get(pid).add(tid);
    }

    private void grabExLock(TransactionId tid, PageId pid) {
        if (!this.pageExLocks.containsKey(pid)) {
            this.pageExLocks.put(pid, tid);
        }
    }
    /**
     * Releases a shared/exclusive lock on a page for a given transaction
     * Removes the entry tid -> {pid} reflecting the tx no longer holds a lock on the page
     * @param tid the transaction id
     * @param pid the page id for which to release the lock
     */
    public synchronized void releasePageLock(TransactionId tid, PageId pid) {
        if (this.pageExLocks.containsKey(pid) && this.pageExLocks.get(pid).equals(tid)) {
            // release write lock
            this.pageExLocks.remove(pid);
        }
        if (this.pageShLocks.containsKey(pid) && this.pageShLocks.get(pid).contains(tid)) {
            this.pageShLocks.get(pid).remove(tid);
        }

        // Remove the entry if empty
        if (this.pageShLocks.containsKey(pid) && this.pageShLocks.get(pid).isEmpty()) {
            this.pageShLocks.remove(pid);
        }

        this.removeDependency(tid);

        // Remove the tx -> page entry
        this.txLocks.get(tid).remove(pid);

        // If tx has no more locks, remove its entry
        if (this.txLocks.get(tid).size() == 0) {
            this.txLocks.remove(tid);
        }

        System.out.println("Tx" + tid + " release " + " lock");
        notifyAll();
    }

    /**
     * Returns true if the given transaction id holds a lock on the page associated with the specified page id
     * @param tid the transaction id
     * @param pid the page id
     * @return true if given transaction id holds a lock on the page
     */
    public synchronized boolean txHoldsLock(TransactionId tid, PageId pid) {
        if (this.pageExLocks.containsKey(pid) && this.pageExLocks.get(pid).equals(tid)) {
            return true;
        }
        if (this.pageShLocks.containsKey(pid) && this.pageShLocks.get(pid).contains(tid)) {
            return true;
        }
        return false;
    }

    /**
     * Returns a set of page ids that the given transaction has a lock on
     * @param tid
     * @return
     */
    public synchronized Set<PageId> txPages(TransactionId tid) {
        HashSet<PageId> s = new HashSet<>();
        if (this.txLocks.containsKey(tid)) {
            s.addAll(this.txLocks.get(tid));
        }
        return s;
    }

    /**
     * Returns true if a deadlock is detected
     * @param tid the tid of the acquiring transaction
     * @return true if a deadlock exists, false otherwise
     */
    private synchronized boolean findDeadlock(TransactionId tid, PageId pid) {
        Queue<TransactionId> q = new LinkedList<>();
        Set<TransactionId> seen = new HashSet<>();
        if (this.deadlockMap.containsKey(tid)) {
            q.addAll(this.deadlockMap.get(tid));
        }
        while (!q.isEmpty()) {
            TransactionId tx = q.remove();
            if (!seen.contains(tx)) {
                seen.add(tx);
                if (tx.equals(tid)) {
                    return true;
                }
                if (this.deadlockMap.containsKey(tx)) {
                    q.addAll(this.deadlockMap.get(tx));
                }
            }
        }
        return false;
    }

    // Maps a Tx dependency such that TxA -> {TxB, TxC} TxA is waiting on B and C
    // Will never map TxA -> itself
    private synchronized void mapDependency(TransactionId tid, PageId pid, Set<TransactionId> txHolders) throws TransactionAbortedException {
        if (!this.deadlockMap.containsKey(tid)) {
            this.deadlockMap.put(tid, new HashSet<>());
        }
        txHolders.remove(tid);
        this.deadlockMap.get(tid).addAll(txHolders);
        if (findDeadlock(tid, pid)) {
            System.out.println("Tx" + tid + " aborting due to deadlock ");
            throw new TransactionAbortedException();
        }
    }


    // Removes a tid from all dependencies so no transaction is waiting for it
    private synchronized void removeDependency(TransactionId tid) {
        this.deadlockMap.remove(tid);
        for (TransactionId tx : this.deadlockMap.keySet()) {
            this.deadlockMap.get(tx).remove(tid);
        }
    }
}