import arrow
from bloop import Column, UUID, DateTime, ConstraintViolation, NotModified
from dvox.app import engine, atomic_engine, config
from dvox.exceptions import InUse, RetryOperation
from dvox.models.types import Position


class ChunkLock(engine.model):
    world = Column(UUID, hash_key=True, name='w')
    chunk = Column(Position, range_key=True, name="c")
    worker = Column(UUID, name='r')
    acquire = Column(DateTime, name='a')

    @classmethod
    def acquire(cls, world, chunk, worker):
        now = arrow.now()
        lock = ChunkLock(world=world, chunk=chunk,
                         worker=worker, acquire=now)
        not_exist = ChunkLock.world.is_(None) & ChunkLock.chunk.is_(None)
        try:
            engine.save(lock, condition=not_exist)
        except ConstraintViolation:
            # This chunk is already locked.  If the given worker matches
            # the existing lock's worker, we can update the acquire time to
            # renew the existing lock.

            # Otherwise, the lock is held by a different worker; if the
            # acquire time is older than the configurable update interval,
            # that worker may have died.  We can forcefully acquire the
            # lock at that point.
            existing_lock = ChunkLock(world=world, chunk=chunk)
            try:
                engine.load(existing_lock)
            except NotModified:
                # The lock that existed when we tried to save above is
                # no longer there.  A periodic cleanup may have swept it
                # up before we could compare acquire time.

                # At this point we should simply retry the acquire.
                raise RetryOperation()
            elapsed = (existing_lock.acquire - now).total_seconds()

            # The worker that's trying to acquire the lock is
            # the same one that previously had the lock.  Update
            # acquire time and attempt to save
            if existing_lock.worker == lock.worker:
                existing_lock.acquire = now
                try:
                    atomic_engine.save(existing_lock)
                except ConstraintViolation:
                    # The worker or acquire time on the existing lock
                    # has changed.  Signal that the acquire can be
                    # retried.
                    raise RetryOperation()
                else:
                    return existing_lock
            # The worker is different, but hasn't checked in within the
            # CHUNK_LOCK_TIMEOUT_SECONDS, and can be re-acquired by a new
            # worker.
            elif elapsed >= config["CHUNK_LOCK_TIMEOUT_SECONDS"]:
                existing_lock.worker = lock.worker
                try:
                    atomic_engine.save(existing_lock)
                except ConstraintViolation:
                    # The worker or acquire time on the existing lock
                    # has changed.  Signal that the acquire can be
                    # retried.
                    raise RetryOperation()
                else:
                    return existing_lock
            # Existing lock is owned by a different worker, and was last
            # acquired within the lock timeout.  Can't be acquired right now.
            else:
                raise InUse()
        # No previous lock existed, save to return.
        else:
            return lock

@@ -445,193 +445,100 @@
     
-@classmethod%0A    def acquire(cls, world, chunk, worker):%0A        now = arrow.now()%0A        lock = ChunkLock(world=world, chunk=chunk,%0A                         worker=worker, acquire=
+def acquire(self):%0A        previous_acquire = self.acquire%0A        self.acquire = arrow.
 now
+(
 )%0A  
@@ -646,20 +646,20 @@
 ne.save(
-lock
+self
 , condit
@@ -725,95 +725,16 @@
     
-# This chunk is already locked.  If the given worker matches%0A            # the 
 existing
  loc
@@ -733,506 +733,32 @@
 ting
- 
+_
 lock
-'s worker, we can update the acquire time to%0A            # renew the existing lock.%0A%0A            # Otherwise, the lock is held by a different worker; if the%0A            # acquire time is older than the configurable update interval,%0A            # that worker may have died.  We can forcefully acquire the%0A            # lock at that point.%0A            existing_lock = ChunkLock(world=world, chunk=chunk)%0A            try:%0A                engine.load(existing_lock)%0A            except NotModified:
+ = self._current()%0A
 %0A   
@@ -762,28 +762,24 @@
             
-    
 # The lock t
@@ -833,20 +833,16 @@
         
-   
- 
 # no lon
@@ -902,20 +902,16 @@
         
-    
 # up bef
@@ -947,15 +947,30 @@
 time
-.%0A%0A    
+, or it was released.%0A
     
@@ -1029,16 +1029,102 @@
 cquire.%0A
+            if existing_lock is None:%0A                self.acquire = previous_acquire%0A
         
@@ -1146,32 +1146,33 @@
 etryOperation()%0A
+%0A
             elap
@@ -1202,19 +1202,28 @@
 quire - 
-now
+self.acquire
 ).total_
@@ -1446,20 +1446,20 @@
 rker == 
-lock
+self
 .worker:
@@ -1499,19 +1499,28 @@
 quire = 
-now
+self.acquire
 %0A       
@@ -1823,52 +1823,39 @@
     
-raise RetryOperation()%0A                else:
+self.acquire = previous_acquire
 %0A   
@@ -1872,35 +1872,37 @@
        r
-eturn existing_lock
+aise RetryOperation()
 %0A       
@@ -2061,16 +2061,61 @@
  worker.
+  Try to take it over and update acquire time
 %0A       
@@ -2216,20 +2216,20 @@
 r = 
-lock
+self
 .worker%0A
     
@@ -2224,16 +2224,69 @@
 .worker%0A
+                existing_lock.acquire = self.acquire%0A
         
@@ -2588,52 +2588,39 @@
     
-raise RetryOperation()%0A                else:
+self.acquire = previous_acquire
 %0A   
@@ -2641,27 +2641,29 @@
    r
-eturn existing_lock
+aise RetryOperation()
 %0A   
@@ -2848,73 +2848,411 @@
     
-raise InUse()%0A        # No previous lock existed, save to
+self.acquire = previous_acquire%0A                raise InUse()%0A        # Success%0A%0A    def _current(self):%0A        %22%22%22%0A        Returns the current lock for the lock's world/position,%0A        or None if one does not exist.%0A        %22%22%22%0A        current_lock = ChunkLock(world=self.world, chunk=self.chunk)%0A        try:%0A            engine.load(current_lock)%0A        except NotModified:%0A           
  return
-.
+ None
 %0A   
@@ -3281,13 +3281,21 @@
  return 
+current_
 lock%0A

Split out ChunkLock._current