Untitled

from in_memory_db import InMemoryDB
from collections import defaultdict, deque


class InMemoryDBImpl(InMemoryDB):
    def __init__(self):
        self.db = {}
        self.modification_counts = {}
        self.locks = {}  # Track which user has locked which key
        self.lock_queues = defaultdict(deque)  # Queues for waiting lock requests

    # ... (Existing methods: set_or_inc, get, delete, top_n_keys remain the same)

    def set_or_inc_by_caller(self, key: str, field: str, value: int, caller_id: str) -> int | None:
        if key in self.db and (key not in self.locks or self.locks[key] == caller_id):
            result = self.set_or_inc(key, field, value)
            return result
        elif key in self.db and field in self.db[key]:  # Key exists but locked by someone else
            return self.db[key][field]
        else:
            return None

    def delete_by_caller(self, key: str, field: str, caller_id: str) -> bool:
        if key in self.db and (key not in self.locks or self.locks[key] == caller_id):
            return self.delete(key, field)
        return False

    def lock(self, caller_id: str, key: str) -> str | None:
        if key not in self.db:
            return "invalid_request"
        if key in self.locks:
            if self.locks[key] == caller_id:
                return None  # Already locked by the same user
            else:
                if caller_id not in self.lock_queues[key]:  # Avoid duplicate requests
                    self.lock_queues[key].append(caller_id)
                return "wait"
        else:
            self.locks[key] = caller_id
            return "acquired"

    def unlock(self, key: str) -> str | None:
        if key not in self.db:
            if key in self.lock_queues:
                del self.lock_queues[key]  # Clean up queue for deleted key
            return "invalid_request"
        if key in self.locks:
            del self.locks[key]
            if self.lock_queues[key]:
                next_caller = self.lock_queues[key].popleft()
                self.locks[key] = next_caller
            return "released"
        return None