"""
WorkClaimer — Product-agnostic SELECT FOR UPDATE SKIP LOCKED utility.

Generates SQLAlchemy queries that atomically claim unclaimed rows for
a specific worker.  Products provide their own ORM models; WorkClaimer
never imports product-specific code.

Pattern:
    Worker process polls for pending work → claims rows with SKIP LOCKED →
    processes them → clears the claim.  If the worker crashes, the sweeper
    reaps expired locks and resets rows for another worker.

Requires PostgreSQL (SKIP LOCKED is a PG extension to SQL FOR UPDATE).
"""

from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from typing import Any, List, Optional

from sqlalchemy import select, update

logger = logging.getLogger(__name__)


class WorkClaimer:
    """Generate SKIP LOCKED queries for distributed DAG workers.

    Usage (product side):

        from pgdag.work import WorkClaimer

        claimer = WorkClaimer()

        # Claim pending rows
        query = claimer.claim_query(
            model=DAGRun,
            status_col=DAGRun.status,
            status_value="pending",
            worker_col=DAGRun.worker_id,
            locked_at_col=DAGRun.locked_at,
            batch_size=5,
        )
        rows = db.scalars(query).all()

        # Stamp ownership
        for row in rows:
            claimer.stamp(row, worker_id="w-abc",
                          worker_col_name="worker_id",
                          locked_at_col_name="locked_at",
                          lock_timeout_col_name="lock_timeout_at",
                          timeout_minutes=5)
        db.commit()

        # After execution, release
        claimer.release(row,
                        worker_col_name="worker_id",
                        locked_at_col_name="locked_at",
                        lock_timeout_col_name="lock_timeout_at")
    """

    def claim_query(
        self,
        *,
        model: Any,
        status_col: Any,
        status_value: Any,
        worker_col: Any,
        locked_at_col: Any,
        batch_size: int = 5,
        order_col: Any = None,
        extra_filters: Optional[List[Any]] = None,
    ) -> Any:
        """Build a SELECT FOR UPDATE SKIP LOCKED query.

        Args:
            model:        SQLAlchemy model class (e.g. DAGRun).
            status_col:   Column for status filtering (e.g. DAGRun.status).
            status_value: Value to match (e.g. "pending" or DAGRunStatus.pending).
            worker_col:   Column for worker ownership (must be NULL for unclaimed).
            locked_at_col: Column for lock timestamp (informational, not filtered).
            batch_size:   Max rows to claim per query.
            order_col:    Column to order by (default: model.created_at).
            extra_filters: Additional SQLAlchemy filter expressions.

        Returns:
            A SQLAlchemy Select that can be passed to db.scalars().
        """
        order = order_col if order_col is not None else model.created_at

        q = (
            select(model)
            .where(status_col == status_value, worker_col.is_(None))
            .order_by(order)
            .limit(batch_size)
            .with_for_update(skip_locked=True)
        )

        if extra_filters:
            for f in extra_filters:
                q = q.where(f)

        return q

    def stamp(
        self,
        row: Any,
        *,
        worker_id: str,
        worker_col_name: str = "worker_id",
        locked_at_col_name: str = "locked_at",
        lock_timeout_col_name: str = "lock_timeout_at",
        timeout_minutes: int = 5,
    ) -> None:
        """Stamp ownership on a claimed row.

        Sets the worker_id, locked_at (now), and lock_timeout_at
        (now + timeout_minutes) on the row object.  Caller must commit.
        """
        now = datetime.now(timezone.utc)
        setattr(row, worker_col_name, worker_id)
        setattr(row, locked_at_col_name, now)
        setattr(row, lock_timeout_col_name, now + timedelta(minutes=timeout_minutes))

    def release(
        self,
        row: Any,
        *,
        worker_col_name: str = "worker_id",
        locked_at_col_name: str = "locked_at",
        lock_timeout_col_name: str = "lock_timeout_at",
    ) -> None:
        """Clear worker claim on a row after execution.  Caller must commit."""
        setattr(row, worker_col_name, None)
        setattr(row, locked_at_col_name, None)
        setattr(row, lock_timeout_col_name, None)

    def reap_expired_query(
        self,
        *,
        model: Any,
        status_col: Any,
        running_status: Any,
        pending_status: Any,
        lock_timeout_col: Any,
        worker_col: Any,
        locked_at_col: Any,
        batch_size: int = 20,
    ) -> Any:
        """Find rows with expired locks for sweeper reaping.

        Returns rows that are in running status but whose lock_timeout_at
        has passed — the worker likely crashed.

        The sweeper should reset these to pending status with cleared
        worker fields so another worker can pick them up.
        """
        now = datetime.now(timezone.utc)
        return (
            select(model)
            .where(
                status_col == running_status,
                lock_timeout_col.isnot(None),
                lock_timeout_col < now,
            )
            .limit(batch_size)
        )
