fastapi_app/lib/core/locking.py

"""
File locking system for FastAPI.

Ported from server/lib/locking.py with FastAPI adaptations:
- Removed Flask dependencies (current_app, ApiError)
- Accept db_dir and logger as parameters
- Use stable_id-based file identification instead of paths
- Keep SQLite-based implementation (migrated schema)
"""

import sqlite3
import threading
from datetime import datetime, timezone, timedelta
from contextlib import contextmanager
from pathlib import Path
from typing import Dict, Optional, List
import logging

LOCK_TIMEOUT_SECONDS = 90

# Track if locks database has been initialized (to avoid redundant init calls)
_locks_db_initialized: set[str] = set()
_locks_db_init_lock = threading.Lock()


@contextmanager
def get_db_connection(db_dir: Path, logger: logging.Logger):
    """
    Context manager for database connections with proper error handling.

    Uses DELETE journal mode instead of WAL for the locks database because:
    - It's a small database with infrequent writes
    - Locks are short-lived and don't benefit from WAL's read concurrency
    - DELETE mode avoids WAL file corruption under rapid concurrent access

    Args:
        db_dir: Directory containing locks.db
        logger: Logger instance

    Yields:
        sqlite3.Connection: Database connection with row factory enabled
    """
    db_path = db_dir / "locks.db"
    conn = None
    try:
        conn = sqlite3.connect(
            str(db_path),
            timeout=30.0,
            check_same_thread=False
        )
        conn.row_factory = sqlite3.Row
        # Use DELETE journal mode (simpler, avoids WAL corruption issues)
        conn.execute("PRAGMA journal_mode = DELETE")
        # Set busy timeout to wait for locks
        conn.execute("PRAGMA busy_timeout = 30000")
        yield conn
        conn.commit()
    except sqlite3.Error as e:
        logger.error(f"Database error: {e}")
        raise RuntimeError(f"Database error: {e}")
    finally:
        if conn:
            conn.close()


def init_locks_db(db_dir: Path, logger: logging.Logger, force: bool = False) -> None:
    """
    Initialize the locks database with the required schema.
    Creates tables and indexes, and runs any pending migrations.

    This function tracks initialization state to avoid redundant calls
    during concurrent request handling. The database is initialized once
    at application startup via database_init.initialize_all_databases().

    Uses DELETE journal mode instead of WAL to avoid corruption issues.

    Args:
        db_dir: Directory containing locks.db
        logger: Logger instance
        force: If True, force re-initialization (used at startup)
    """
    db_path = db_dir / "locks.db"
    db_key = str(db_path.resolve())

    # Quick check without lock - skip if already initialized
    if not force:
        with _locks_db_init_lock:
            if db_key in _locks_db_initialized:
                return

    # Acquire lock for actual initialization
    with _locks_db_init_lock:
        # Double-check after acquiring lock
        if not force and db_key in _locks_db_initialized:
            return

        db_path.parent.mkdir(parents=True, exist_ok=True)

        # Create database and schema if it doesn't exist
        if not db_path.exists():
            with get_db_connection(db_dir, logger) as conn:
                conn.execute("""
                    CREATE TABLE IF NOT EXISTS locks (
                        file_hash TEXT PRIMARY KEY,
                        session_id TEXT NOT NULL,
                        acquired_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                """)

        # Run migrations using centralized runner
        from fastapi_app.lib.core.migration_runner import run_migrations_if_needed
        from fastapi_app.lib.core.migrations.versions import LOCKS_MIGRATIONS

        try:
            run_migrations_if_needed(
                db_path=db_path,
                migrations=LOCKS_MIGRATIONS,
                logger=logger
            )
        except Exception as e:
            logger.error(f"Failed to run migrations for locks.db: {e}")
            raise

        # Create/update indexes with current structure
        try:
            with get_db_connection(db_dir, logger) as conn:
                # Ensure indexes exist with current schema
                conn.execute("""
                    CREATE INDEX IF NOT EXISTS idx_file_id
                    ON locks(file_id)
                """)
                conn.execute("""
                    CREATE INDEX IF NOT EXISTS idx_session
                    ON locks(session_id)
                """)
                conn.execute("""
                    CREATE INDEX IF NOT EXISTS idx_updated
                    ON locks(updated_at)
                """)
                logger.debug(f"Locks database initialized at {db_path}")
        except sqlite3.Error as e:
            logger.error(f"Failed to create indexes: {e}")
            raise RuntimeError(f"Database error: {e}")

        # Mark as initialized
        _locks_db_initialized.add(db_key)


def acquire_lock(file_id: str, session_id: str, db_dir: Path, logger: logging.Logger, max_retries: int = 3) -> bool:
    """
    Tries to acquire a lock for a given file. Returns True on success, False on failure.

    This operation is atomic and handles:
    - Creating new locks
    - Refreshing existing locks owned by the same session
    - Taking over stale locks

    Args:
        file_id: The file's stable_id to lock
        session_id: The session ID requesting the lock
        db_dir: Directory containing locks.db
        logger: Logger instance
        max_retries: Maximum number of retry attempts for transient errors

    Returns:
        bool: True if lock was acquired/refreshed, False if held by another active session

    Raises:
        RuntimeError: If database operations fail after all retries
    """
    import time

    logger.debug(f"[LOCK] Session {session_id[:8]}... attempting to acquire lock for file {file_id[:8]}...")

    # Ensure database is initialized
    init_locks_db(db_dir, logger)

    last_error = None
    for attempt in range(max_retries):
        try:
            return _acquire_lock_impl(file_id, session_id, db_dir, logger)
        except (sqlite3.OperationalError, RuntimeError) as e:
            last_error = e
            error_msg = str(e).lower()
            # Retry on transient errors like disk I/O, busy, or locked
            if any(err in error_msg for err in ['disk i/o', 'busy', 'locked', 'database is malformed']):
                if attempt < max_retries - 1:
                    delay = 0.1 * (2 ** attempt)  # Exponential backoff: 0.1, 0.2, 0.4 seconds
                    logger.warning(
                        f"[LOCK] Transient error on attempt {attempt + 1}/{max_retries}: {e}. "
                        f"Retrying in {delay:.1f}s..."
                    )
                    time.sleep(delay)
                    continue
            # Non-transient error or max retries exceeded
            raise

    # Should not reach here, but handle it
    logger.error(f"[LOCK] Failed to acquire lock after {max_retries} attempts: {last_error}")
    raise RuntimeError(f"Database error after {max_retries} retries: {last_error}")


def _acquire_lock_impl(file_id: str, session_id: str, db_dir: Path, logger: logging.Logger) -> bool:
    """Internal implementation of acquire_lock without retry logic."""
    with get_db_connection(db_dir, logger) as conn:
        # Use IMMEDIATE transaction to get write lock upfront, preventing race conditions
        # This ensures only one transaction can modify locks at a time
        conn.isolation_level = None  # Auto-commit off for explicit transaction control
        conn.execute("BEGIN IMMEDIATE")

        try:
            cursor = conn.cursor()

            # Calculate staleness threshold
            stale_threshold = datetime.now(timezone.utc) - timedelta(seconds=LOCK_TIMEOUT_SECONDS)

            # Check if there's an existing lock
            cursor.execute("""
                SELECT session_id, updated_at
                FROM locks
                WHERE file_id = ?
            """, (file_id,))

            existing = cursor.fetchone()

            if existing:
                existing_session = existing['session_id']
                updated_at = datetime.fromisoformat(existing['updated_at'])
                updated_at_utc = updated_at.replace(tzinfo=timezone.utc)
                is_stale = updated_at_utc < stale_threshold
                age_seconds = (datetime.now(timezone.utc) - updated_at_utc).total_seconds()

                logger.debug(
                    f"[LOCK] Existing lock found: owner={existing_session[:8]}..., "
                    f"age={age_seconds:.1f}s, stale={is_stale}, threshold={LOCK_TIMEOUT_SECONDS}s"
                )

                if existing_session == session_id:
                    # It's our lock, refresh it
                    cursor.execute("""
                        UPDATE locks
                        SET updated_at = CURRENT_TIMESTAMP
                        WHERE file_id = ?
                    """, (file_id,))
                    conn.commit()
                    logger.debug(f"[LOCK] Session {session_id[:8]}... refreshed own lock for file {file_id[:8]}...")
                    return True
                elif is_stale:
                    # Lock is stale, take it over
                    cursor.execute("""
                        UPDATE locks
                        SET session_id = ?, updated_at = CURRENT_TIMESTAMP
                        WHERE file_id = ?
                    """, (session_id, file_id))
                    conn.commit()
                    logger.warning(
                        f"[LOCK] Session {session_id[:8]}... took over stale lock for file {file_id[:8]}... "
                        f"from session {existing_session[:8]}... (was {age_seconds:.1f}s old)"
                    )
                    return True
                else:
                    # Lock is held by another active session - DENY
                    conn.rollback()
                    logger.warning(
                        f"[LOCK] Session {session_id[:8]}... DENIED lock for file {file_id[:8]}.... "
                        f"Held by {existing_session[:8]}... (age={age_seconds:.1f}s, still fresh)"
                    )
                    return False
            else:
                # No local lock — check remote lock cache before granting
                if _is_remotely_locked(file_id, db_dir, logger):
                    conn.rollback()
                    logger.warning(
                        f"[LOCK] Session {session_id[:8]}... DENIED lock for file {file_id[:8]}...: "
                        f"file is locked on a remote instance"
                    )
                    return False

                # No lock at all, create new one
                logger.debug(f"[LOCK] No existing lock found for file {file_id[:8]}...")
                cursor.execute("""
                    INSERT INTO locks (file_id, session_id, updated_at)
                    VALUES (?, ?, CURRENT_TIMESTAMP)
                """, (file_id, session_id))
                conn.commit()
                logger.info(f"[LOCK] Session {session_id[:8]}... acquired NEW lock for file {file_id[:8]}...")
                return True
        except Exception as e:
            conn.rollback()
            logger.error(f"[LOCK] Error during lock acquisition: {e}")
            raise


_REMOTE_LOCK_TTL_SECONDS = 90 + 360  # local timeout + one sync-cycle buffer


def _is_remotely_locked(file_id: str, db_dir: Path, logger: logging.Logger) -> bool:
    """Check the remote lock cache to see if another instance holds this file.

    Args:
        file_id: The file's stable_id.
        db_dir: Application db directory (contains metadata.db).
        logger: Logger instance.

    Returns:
        True if a non-expired remote lock exists for this file.
    """
    try:
        from fastapi_app.lib.core.database import DatabaseManager
        from fastapi_app.lib.repository.file_repository import FileRepository
        repo = FileRepository(DatabaseManager(db_dir / "metadata.db"))
        remote_locks = repo.get_remote_locks()
        if file_id not in remote_locks:
            return False
        lock_info = remote_locks[file_id]
        updated_at_str = lock_info.get("updated_at") or lock_info.get("acquired_at", "")
        if not updated_at_str:
            return False
        updated_at = datetime.fromisoformat(updated_at_str)
        if updated_at.tzinfo is None:
            updated_at = updated_at.replace(tzinfo=timezone.utc)
        age_seconds = (datetime.now(timezone.utc) - updated_at).total_seconds()
        return age_seconds < _REMOTE_LOCK_TTL_SECONDS
    except Exception as exc:
        logger.debug(f"[LOCK] Remote lock check failed (non-fatal): {exc}")
        return False


def release_lock(file_id: str, session_id: str, db_dir: Path, logger: logging.Logger) -> Dict[str, str]:
    """
    Releases the lock for a given file if it is held by the current session.

    Args:
        file_id: The file's stable_id to unlock
        session_id: The session ID releasing the lock
        db_dir: Directory containing locks.db
        logger: Logger instance

    Returns:
        dict: Structured response with status, action, and message
            - status: "success" or "error"
            - action: "released", "already_released", "not_owned"
            - message: Human-readable description

    Raises:
        RuntimeError: If attempting to release a lock owned by another session or database error
    """
    logger.debug(f"[LOCK] Session {session_id[:8]}... attempting to release lock for file {file_id[:8]}...")

    # Ensure database is initialized
    init_locks_db(db_dir, logger)

    with get_db_connection(db_dir, logger) as conn:
        cursor = conn.cursor()

        # Check if lock exists and who owns it
        cursor.execute("""
            SELECT session_id
            FROM locks
            WHERE file_id = ?
        """, (file_id,))

        existing = cursor.fetchone()

        if existing:
            existing_session = existing['session_id']
            if existing_session == session_id:
                # Delete the lock
                cursor.execute("""
                    DELETE FROM locks
                    WHERE file_id = ?
                """, (file_id,))
                logger.info(f"[LOCK] Session {session_id[:8]}... released lock for file {file_id[:8]}...")
                return {
                    "status": "success",
                    "action": "released",
                    "message": f"Lock successfully released for file {file_id}"
                }
            else:
                # Attempting to release someone else's lock
                logger.warning(
                    f"[LOCK] Session {session_id[:8]}... DENIED release of lock for file {file_id[:8]}.... "
                    f"Owned by {existing_session[:8]}..."
                )
                raise RuntimeError(
                    f"Session {session_id} attempted to release a lock owned by {existing_session}"
                )
        else:
            # Lock doesn't exist - idempotent success
            logger.info(
                f"[LOCK] Session {session_id[:8]}... attempted to release lock for file {file_id[:8]}..., "
                f"but no lock exists (idempotent success)"
            )
            return {
                "status": "success",
                "action": "already_released",
                "message": f"Lock was already released for file {file_id}"
            }


def cleanup_stale_locks(db_dir: Path, logger: logging.Logger, timeout_seconds: int = LOCK_TIMEOUT_SECONDS) -> int:
    """
    Removes all stale locks from the database.

    Args:
        db_dir: Directory containing locks.db
        logger: Logger instance
        timeout_seconds: Lock timeout in seconds

    Returns:
        int: Number of stale locks purged
    """
    # Ensure database is initialized
    try:
        init_locks_db(db_dir, logger)
    except Exception as e:
        logger.error(f"Could not initialize locks database: {e}")
        return 0

    try:
        with get_db_connection(db_dir, logger) as conn:
            cursor = conn.cursor()

            # Calculate the staleness threshold
            stale_threshold = datetime.now(timezone.utc) - timedelta(seconds=timeout_seconds)

            # Delete stale locks
            cursor.execute("""
                DELETE FROM locks
                WHERE datetime(updated_at) < datetime(?)
            """, (stale_threshold.isoformat(),))

            purged_count = cursor.rowcount

            if purged_count > 0:
                logger.info(f"Purged {purged_count} stale locks")

            return purged_count
    except Exception as e:
        logger.error(f"Error purging stale locks: {e}")
        return 0


def get_all_active_locks(db_dir: Path, logger: logging.Logger, timeout_seconds: int = LOCK_TIMEOUT_SECONDS) -> Dict[str, str]:
    """
    Fetches all non-stale locks and returns a map of file_id -> session_id.

    Args:
        db_dir: Directory containing locks.db
        logger: Logger instance
        timeout_seconds: Lock timeout in seconds

    Returns:
        dict: Dictionary mapping file stable_ids to session IDs for all active locks
    """
    # Ensure database is initialized
    try:
        init_locks_db(db_dir, logger)
    except Exception as e:
        logger.error(f"Could not initialize locks database: {e}")
        return {}

    try:
        with get_db_connection(db_dir, logger) as conn:
            cursor = conn.cursor()

            # Calculate the staleness threshold
            stale_threshold = datetime.now(timezone.utc) - timedelta(seconds=timeout_seconds)

            # Get all non-stale locks
            cursor.execute("""
                SELECT file_id, session_id
                FROM locks
                WHERE datetime(updated_at) >= datetime(?)
            """, (stale_threshold.isoformat(),))

            active_locks = {row['file_id']: row['session_id'] for row in cursor.fetchall()}

            return active_locks
    except Exception as e:
        logger.error(f"Error fetching active locks: {e}")
        return {}


def get_locked_file_ids(db_dir: Path, logger: logging.Logger, session_id: Optional[str] = None, repo=None) -> List[str]:
    """
    Returns a list of file stable_ids which are currently locked.

    Optionally filters by session_id.

    Args:
        db_dir: Directory containing locks.db
        logger: Logger instance
        session_id: If provided, only return locks for this session
        repo: Unused (kept for backward compatibility)

    Returns:
        list: List of file stable_ids that are locked
    """
    active_locks = get_all_active_locks(db_dir, logger)
    locked_file_ids = []

    for file_id, lock_session_id in active_locks.items():
        if session_id and lock_session_id != session_id:
            continue

        locked_file_ids.append(file_id)

    return locked_file_ids


def check_lock(file_id: str, session_id: str, db_dir: Path, logger: logging.Logger) -> Dict[str, any]:
    """
    Checks if a single file is locked by another session.

    Args:
        file_id: The file's stable_id to check
        session_id: The session ID making the check
        db_dir: Directory containing locks.db
        logger: Logger instance

    Returns:
        dict: {"is_locked": bool, "locked_by": Optional[str]}
    """
    # Ensure database is initialized
    try:
        init_locks_db(db_dir, logger)
    except Exception as e:
        logger.error(f"Could not initialize locks database: {e}")
        return {"is_locked": False, "locked_by": None}

    active_locks = get_all_active_locks(db_dir, logger)

    if file_id in active_locks:
        lock_owner = active_locks[file_id]
        if lock_owner != session_id:
            logger.debug(f"File is locked by another session: {lock_owner}")
            return {"is_locked": True, "locked_by": lock_owner}

    return {"is_locked": False, "locked_by": None}


def reset_locks_db_initialized() -> None:
    """
    Reset the locks database initialization tracking.

    This is primarily for testing purposes, to allow re-initialization
    after database files are deleted/recreated between tests.
    """
    with _locks_db_init_lock:
        _locks_db_initialized.clear()