"""PostgreSQL-backed scheduled job operations."""

from contextlib import contextmanager
from datetime import datetime, timezone
import uuid

from psycopg2.extras import Json

from core import postgres


DEFAULT_MAX_ATTEMPTS = 3
DEFAULT_LEASE_SECONDS = 300
DEFAULT_RETRY_SECONDS = 30
MAX_RETRY_SECONDS = 900


@contextmanager
def _using_cursor(cursor=None):
    if cursor is not None:
        yield cursor
        return
    with postgres.get_cursor() as owned_cursor:
        yield owned_cursor


def _timestamp(value, field="timestamp"):
    if isinstance(value, str):
        value = datetime.fromisoformat(value.replace("Z", "+00:00"))
    if not isinstance(value, datetime):
        raise ValueError(f"{field} must be a datetime or ISO-8601 string")
    if value.tzinfo is None:
        value = value.replace(tzinfo=timezone.utc)
    return value.astimezone(timezone.utc)


def _positive(value, field, maximum=None):
    try:
        value = int(value)
    except (TypeError, ValueError) as error:
        raise ValueError(f"{field} must be a whole number") from error
    if value < 1 or (maximum is not None and value > maximum):
        suffix = f" and at most {maximum}" if maximum is not None else ""
        raise ValueError(f"{field} must be at least 1{suffix}")
    return value


def _row(cursor):
    record = cursor.fetchone()
    return dict(record) if record else None


def create_job(
    job_type,
    payload,
    run_at,
    user_uuid=None,
    max_attempts=DEFAULT_MAX_ATTEMPTS,
    idempotency_key=None,
    job_id=None,
    cursor=None,
):
    """Create a job, returning the existing row for a repeated idempotency key."""
    if not isinstance(job_type, str) or not job_type.strip():
        raise ValueError("job_type is required")
    max_attempts = _positive(max_attempts, "max_attempts", 100)
    values = {
        "id": str(job_id or uuid.uuid4()),
        "user_uuid": user_uuid,
        "job_type": job_type.strip(),
        "payload": Json(payload if payload is not None else {}),
        "run_at": _timestamp(run_at, "run_at"),
        "max_attempts": max_attempts,
        "idempotency_key": idempotency_key,
    }
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            INSERT INTO scheduled_jobs (
                id, user_uuid, job_type, payload, run_at,
                max_attempts, idempotency_key
            ) VALUES (
                %(id)s, %(user_uuid)s, %(job_type)s, %(payload)s, %(run_at)s,
                %(max_attempts)s, %(idempotency_key)s
            )
            ON CONFLICT (job_type, idempotency_key)
            DO UPDATE SET idempotency_key = EXCLUDED.idempotency_key
            WHERE scheduled_jobs.user_uuid IS NOT DISTINCT FROM EXCLUDED.user_uuid
            RETURNING *
            """,
            values,
        )
        return _row(active_cursor)


def get_job(job_id, cursor=None):
    """Return one job by UUID."""
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute("SELECT * FROM scheduled_jobs WHERE id = %s", (job_id,))
        return _row(active_cursor)


def list_jobs(user_uuid=None, status=None, job_type=None, limit=100, cursor=None):
    """List jobs newest first, optionally filtered by owner, status, or type."""
    limit = _positive(limit, "limit", 500)
    clauses = []
    params = []
    if user_uuid is not None:
        clauses.append("user_uuid = %s")
        params.append(user_uuid)
    if status is not None:
        clauses.append("status = %s")
        params.append(status)
    if job_type is not None:
        clauses.append("job_type = %s")
        params.append(job_type)
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    params.append(limit)

    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            f"""
            SELECT * FROM scheduled_jobs
            {where}
            ORDER BY created_at DESC
            LIMIT %s
            """,
            params,
        )
        return [dict(record) for record in active_cursor.fetchall()]


def claim_due_jobs(
    worker_id,
    limit=10,
    lease_seconds=DEFAULT_LEASE_SECONDS,
    job_types=None,
    cursor=None,
):
    """Atomically lease due jobs using row locks that skip other workers."""
    if not worker_id:
        raise ValueError("worker_id is required")
    limit = _positive(limit, "limit", 100)
    lease_seconds = _positive(lease_seconds, "lease_seconds")
    if isinstance(job_types, str):
        job_types = [job_types]
    elif job_types is not None:
        job_types = list(job_types)
        if not job_types:
            return []

    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            WITH exhausted AS (
                SELECT id
                FROM scheduled_jobs
                WHERE status = 'running'
                  AND (lease_until IS NULL OR lease_until <= NOW())
                  AND attempts >= max_attempts
                FOR UPDATE SKIP LOCKED
            )
            UPDATE scheduled_jobs AS job
            SET status = 'failed',
                lease_until = NULL,
                leased_by = NULL,
                last_error = COALESCE(last_error, 'lease expired'),
                updated_at = NOW()
            FROM exhausted
            WHERE job.id = exhausted.id
            """
        )
        active_cursor.execute(
            """
            WITH candidates AS (
                SELECT id
                FROM scheduled_jobs
                WHERE run_at <= NOW()
                  AND attempts < max_attempts
                  AND (
                      status = 'pending'
                      OR (
                          status = 'running'
                          AND (lease_until IS NULL OR lease_until <= NOW())
                      )
                  )
                  AND (%(job_types)s IS NULL OR job_type = ANY(%(job_types)s))
                ORDER BY run_at, created_at
                FOR UPDATE SKIP LOCKED
                LIMIT %(limit)s
            )
            UPDATE scheduled_jobs AS job
            SET status = 'running',
                attempts = job.attempts + 1,
                leased_by = %(worker_id)s,
                lease_until = NOW() + (%(lease_seconds)s * INTERVAL '1 second'),
                updated_at = NOW()
            FROM candidates
            WHERE job.id = candidates.id
            RETURNING job.*
            """,
            {
                "job_types": job_types,
                "limit": limit,
                "worker_id": worker_id,
                "lease_seconds": lease_seconds,
            },
        )
        return [dict(record) for record in active_cursor.fetchall()]


def renew_job_lease(job_id, worker_id, lease_seconds=DEFAULT_LEASE_SECONDS, cursor=None):
    """Extend a lease only while it is owned by the requesting worker."""
    lease_seconds = _positive(lease_seconds, "lease_seconds")
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            UPDATE scheduled_jobs
            SET lease_until = NOW() + (%s * INTERVAL '1 second'), updated_at = NOW()
            WHERE id = %s AND status = 'running' AND leased_by = %s
            RETURNING *
            """,
            (lease_seconds, job_id, worker_id),
        )
        return _row(active_cursor)


def complete_job(job_id, worker_id, cursor=None):
    """Mark a job complete when its lease is still owned by the worker."""
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            UPDATE scheduled_jobs
            SET status = 'completed', completed_at = NOW(), updated_at = NOW(),
                leased_by = NULL, lease_until = NULL, last_error = NULL
            WHERE id = %s AND status = 'running' AND leased_by = %s
            RETURNING *
            """,
            (job_id, worker_id),
        )
        return _row(active_cursor)


def fail_job(
    job_id,
    worker_id,
    error,
    retry_seconds=DEFAULT_RETRY_SECONDS,
    max_retry_seconds=MAX_RETRY_SECONDS,
    cursor=None,
):
    """Fail or reschedule an owned job using bounded exponential backoff."""
    retry_seconds = _positive(retry_seconds, "retry_seconds")
    max_retry_seconds = _positive(max_retry_seconds, "max_retry_seconds")
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            SELECT attempts, max_attempts
            FROM scheduled_jobs
            WHERE id = %s AND status = 'running' AND leased_by = %s
            FOR UPDATE
            """,
            (job_id, worker_id),
        )
        current = active_cursor.fetchone()
        if not current:
            return None

        exhausted = current["attempts"] >= current["max_attempts"]
        delay = min(
            max_retry_seconds,
            retry_seconds * (2 ** min(max(current["attempts"] - 1, 0), 30)),
        )
        active_cursor.execute(
            """
            UPDATE scheduled_jobs
            SET status = %(status)s,
                run_at = CASE
                    WHEN %(exhausted)s THEN run_at
                    ELSE NOW() + (%(delay)s * INTERVAL '1 second')
                END,
                leased_by = NULL,
                lease_until = NULL,
                last_error = %(error)s,
                updated_at = NOW()
            WHERE id = %(job_id)s
            RETURNING *
            """,
            {
                "status": "failed" if exhausted else "pending",
                "exhausted": exhausted,
                "delay": delay,
                "error": str(error)[:4000],
                "job_id": job_id,
            },
        )
        return _row(active_cursor)


def cancel_job(job_id, user_uuid=None, cursor=None):
    """Cancel one unfinished job, optionally enforcing its owner."""
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            UPDATE scheduled_jobs
            SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
                leased_by = NULL, lease_until = NULL
            WHERE id = %(job_id)s
              AND status IN ('pending', 'running')
              AND (%(user_uuid)s IS NULL OR user_uuid = %(user_uuid)s)
            RETURNING *
            """,
            {"job_id": job_id, "user_uuid": user_uuid},
        )
        return _row(active_cursor)


def cancel_jobs(
    user_uuid=None, job_type=None, idempotency_key=None, cursor=None
):
    """Cancel matching unfinished jobs; at least one filter is required."""
    filters = {
        "user_uuid": user_uuid,
        "job_type": job_type,
        "idempotency_key": idempotency_key,
    }
    clauses = [f"{name} = %({name})s" for name, value in filters.items() if value is not None]
    if not clauses:
        raise ValueError("at least one cancellation filter is required")

    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            f"""
            UPDATE scheduled_jobs
            SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
                leased_by = NULL, lease_until = NULL
            WHERE status IN ('pending', 'running')
              AND {' AND '.join(clauses)}
            RETURNING *
            """,
            filters,
        )
        return [dict(record) for record in active_cursor.fetchall()]