"""Durable background jobs, kept in Postgres (the `jobs` table).

    await jobs.enqueue(db, "storage.build", {"source_id": ...}, tenant_id=...)

enqueue() writes the job in the caller's own transaction: it exists exactly
when the change that asked for it committed — a request that rolls back
leaves nothing behind, and nothing runs before the data it needs is visible.

Workers (run_worker, started by the API's lifespan, or `python -m app.worker`
on its own) claim due jobs with FOR UPDATE SKIP LOCKED, so any number of
them can run side by side and each job runs once at a time. A claimed job is
held under a lease the worker renews while it runs; a worker that dies stops
renewing, and the job is handed out again once the lease runs out.

A handler that raises is retried with exponential backoff (and jitter) until
its kind's max_attempts; Permanent fails it at once, Retry picks the delay.
A job that fails for good runs its kind's on_failure hook, if it has one,
and stays in the table with its error for an administrator to see and
retry. Finished jobs are cleared after a while (see _RETENTION).

Handlers register with @handler("kind"); the modules holding them are
imported by load_handlers(). A worker claims only the kinds it has a
handler for (an older worker mid-deploy leaves newer kinds alone), and only
from its own queue (settings.job_queue) — which is also where this
process's enqueue() puts jobs.
"""

import asyncio
import contextlib
import logging
import os
import random
import socket
import time
import uuid
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession

from app.core import observability
from app.core.config import settings
from app.core.database import AsyncSessionLocal

log = logging.getLogger("flowfabric.jobs")

WORKER_ID = f"{socket.gethostname()}:{os.getpid()}"
LEASE_SECONDS = 60
POLL_SECONDS = 1.0
CONCURRENCY = 4
_BACKOFF_BASE = 10  # seconds before the first retry; doubles each time
_BACKOFF_CAP = 3600
_RETENTION = {"done": timedelta(days=7), "cancelled": timedelta(days=7), "failed": timedelta(days=30)}

# The modules whose @handler registrations the worker needs.
_HANDLER_MODULES = ("app.crud.field_indexes", "app.crud.entity_storage", "app.core.scheduler", "app.crud.notifications", "app.crud.account_emails", "app.crud.form_documents", "app.crud.workflow_engine", "app.crud.record_import", "app.crud.summary_cache", "app.crud.task_history", "app.crud.application_review")


class Permanent(Exception):
    """Fails the job now — retrying can't help."""


class Retry(Exception):
    """Retries the job after `delay` seconds (counts as an attempt)."""

    def __init__(self, message: str, delay: float):
        super().__init__(message)
        self.delay = delay


@dataclass
class Job:
    id: uuid.UUID
    kind: str
    payload: dict
    tenant_id: uuid.UUID | None
    attempt: int
    max_attempts: int

    @property
    def last_attempt(self) -> bool:
        return self.attempt >= self.max_attempts


Handler = Callable[[Job], Awaitable[str | None]]
FailureHook = Callable[[Job, BaseException], Awaitable[None]]


@dataclass
class _Kind:
    run: Handler
    max_attempts: int
    on_failure: FailureHook | None = None


_KINDS: dict[str, _Kind] = {}


def handler(kind: str, *, max_attempts: int = 5):
    """Registers the function that runs jobs of `kind`. It returns a short
    result text (or None) — kept on the job row."""

    def register(fn: Handler) -> Handler:
        _KINDS[kind] = _Kind(fn, max_attempts)
        return fn

    return register


def on_failure(kind: str):
    """Registers what to do once a job of `kind` has failed for good."""

    def register(fn: FailureHook) -> FailureHook:
        _KINDS[kind].on_failure = fn
        return fn

    return register


def load_handlers() -> None:
    import importlib

    for module in _HANDLER_MODULES:
        importlib.import_module(module)


# ---- enqueueing ---------------------------------------------------------------------------------------------


async def enqueue(
    db: AsyncSession,
    kind: str,
    payload: dict | None = None,
    *,
    tenant_id: uuid.UUID | None = None,
    dedupe_key: str | None = None,
    delay: float = 0,
    max_attempts: int | None = None,
) -> uuid.UUID:
    """Adds a job in `db`'s transaction. With a dedupe_key, a job already
    queued under that key is updated instead — its payload replaced (the
    latest one wins) and it runs no later than it would have."""
    if max_attempts is None:
        load_handlers()
        max_attempts = _KINDS[kind].max_attempts if kind in _KINDS else 5
    row = await db.execute(
        text(
            "INSERT INTO jobs (queue, tenant_id, kind, payload, dedupe_key, max_attempts, run_after) "
            "VALUES (:queue, :tenant, :kind, CAST(:payload AS jsonb), :key, :max, now() + make_interval(secs => :delay)) "
            "ON CONFLICT (dedupe_key) WHERE status = 'queued' DO UPDATE SET payload = EXCLUDED.payload, "
            "run_after = LEAST(jobs.run_after, EXCLUDED.run_after), max_attempts = EXCLUDED.max_attempts "
            "RETURNING id"
        ),
        {"queue": settings.job_queue, "tenant": tenant_id, "kind": kind, "payload": _json(payload or {}), "key": dedupe_key, "max": max_attempts, "delay": float(delay)},
    )
    return row.scalar_one()


def _json(value: Any) -> str:
    import json

    return json.dumps(value, default=str)


@contextlib.asynccontextmanager
async def tenant_session(tenant_id: uuid.UUID | None):
    """A session in its own transaction, scoped to a tenant the way a request is (get_db)."""
    async with AsyncSessionLocal() as db:
        async with db.begin():
            if tenant_id is not None:
                await db.execute(text("SELECT set_config('app.current_tenant_id', :v, true)"), {"v": str(tenant_id)})
            yield db


async def save_payload(db: AsyncSession, job: Job) -> None:
    """Keeps a handler's progress (job.payload, changed in place) for its next attempt."""
    await db.execute(text("UPDATE jobs SET payload = CAST(:payload AS jsonb) WHERE id = :id"), {"payload": _json(job.payload), "id": job.id})


# ---- running ------------------------------------------------------------------------------------------------


async def _claim(limit: int, kinds: list[str] | None = None, tenant_id: uuid.UUID | None = None) -> list[Job]:
    where = ["queue = :queue", "status = 'queued'", "run_after <= now()", "kind = ANY(:kinds)"]
    known = [k for k in (kinds or _KINDS) if k in _KINDS]
    params: dict[str, Any] = {"limit": limit, "worker": WORKER_ID, "lease": float(LEASE_SECONDS), "queue": settings.job_queue, "kinds": known}
    if tenant_id is not None:
        where.append("tenant_id = :tenant")
        params["tenant"] = tenant_id
    async with AsyncSessionLocal() as db:
        async with db.begin():
            rows = await db.execute(
                text(
                    "UPDATE jobs SET status = 'running', attempts = attempts + 1, locked_by = :worker, "
                    "locked_until = now() + make_interval(secs => :lease), started_at = now() "
                    f"WHERE id IN (SELECT id FROM jobs WHERE {' AND '.join(where)} "
                    "ORDER BY run_after, created_at FOR UPDATE SKIP LOCKED LIMIT :limit) "
                    "RETURNING id, kind, payload, tenant_id, attempts, max_attempts"
                ),
                params,
            )
            return [Job(r.id, r.kind, r.payload, r.tenant_id, r.attempts, r.max_attempts) for r in rows]


def backoff(attempt: int) -> float:
    """Seconds before retry number `attempt`: 10s, 20s, 40s … capped at an hour, ±20%."""
    delay = min(_BACKOFF_CAP, _BACKOFF_BASE * 2 ** (attempt - 1))
    return delay * random.uniform(0.8, 1.2)


async def _finish(job: Job, status: str, *, result: str | None = None, error: str | None = None, retry_in: float | None = None) -> None:
    async with AsyncSessionLocal() as db:
        async with db.begin():
            # Only while this worker still holds the job — one whose lease ran
            # out may have been handed to another worker meanwhile.
            await db.execute(
                text(
                    "UPDATE jobs SET status = CAST(:status AS text), result = COALESCE(:result, result), last_error = :error, "
                    "locked_by = NULL, locked_until = NULL, "
                    "run_after = CASE WHEN CAST(:status AS text) = 'queued' THEN now() + make_interval(secs => :retry_in) ELSE run_after END, "
                    "finished_at = CASE WHEN CAST(:status AS text) IN ('done', 'failed') THEN now() END "
                    "WHERE id = :id AND locked_by = :worker"
                ),
                {"status": status, "result": result, "error": error, "retry_in": float(retry_in or 0), "id": job.id, "worker": WORKER_ID},
            )


async def _heartbeat(job: Job) -> None:
    while True:
        await asyncio.sleep(LEASE_SECONDS / 3)
        async with AsyncSessionLocal() as db:
            async with db.begin():
                await db.execute(
                    text("UPDATE jobs SET locked_until = now() + make_interval(secs => :lease) WHERE id = :id AND locked_by = :worker"),
                    {"lease": float(LEASE_SECONDS), "id": job.id, "worker": WORKER_ID},
                )


def _error_text(exc: BaseException) -> str:
    return (f"{type(exc).__name__}: {exc}" if str(exc) else type(exc).__name__)[:2000]


async def run_job(job: Job) -> str:
    """Runs one claimed job to its next state. Returns that state."""
    with observability.context(job_id=job.id, job_kind=job.kind, tenant_id=job.tenant_id, attempt=job.attempt):
        started = time.perf_counter()
        outcome = await _run_job(job)
        observability.JOB_SECONDS.labels(job.kind).observe(time.perf_counter() - started)
        observability.JOB_RUNS.labels(job.kind, "retry" if outcome == "queued" else outcome).inc()
        return outcome


async def _run_job(job: Job) -> str:
    kind = _KINDS.get(job.kind)
    if kind is None:
        await _finish(job, "failed", error=f"No handler for job kind {job.kind!r}")
        return "failed"
    beat = asyncio.create_task(_heartbeat(job))
    try:
        result = await kind.run(job)
    except Exception as exc:  # noqa: BLE001 - every failure is recorded on the job
        retry_in = exc.delay if isinstance(exc, Retry) else backoff(job.attempt)
        final = isinstance(exc, Permanent) or job.last_attempt
        if final:
            log.warning("Job %s (%s) failed for good after %s attempt(s): %s", job.id, job.kind, job.attempt, exc)
            if kind.on_failure is not None:
                try:
                    await kind.on_failure(job, exc)
                except Exception:  # noqa: BLE001
                    log.exception("on_failure hook for job %s failed", job.id)
        else:
            log.info("Job %s (%s) attempt %s failed, retrying in %.0fs: %s", job.id, job.kind, job.attempt, retry_in, exc)
        await _finish(job, "failed" if final else "queued", error=_error_text(exc), retry_in=retry_in)
        return "failed" if final else "queued"
    finally:
        beat.cancel()
    await _finish(job, "done", result=(result or "")[:2000] or None)
    return "done"


async def reclaim_expired() -> int:
    """Jobs whose worker stopped renewing its lease go back in the queue
    (or fail, when that was their last attempt)."""
    async with AsyncSessionLocal() as db:
        async with db.begin():
            rows = await db.execute(
                text(
                    "UPDATE jobs SET status = CASE WHEN attempts >= max_attempts THEN 'failed' ELSE 'queued' END, "
                    "last_error = 'The worker running it stopped responding', locked_by = NULL, locked_until = NULL, "
                    "finished_at = CASE WHEN attempts >= max_attempts THEN now() END "
                    "WHERE status = 'running' AND locked_until < now() RETURNING id"
                )
            )
            return len(rows.all())


async def purge_finished() -> None:
    async with AsyncSessionLocal() as db:
        async with db.begin():
            for status, age in _RETENTION.items():
                await db.execute(
                    text("DELETE FROM jobs WHERE status = :status AND COALESCE(finished_at, created_at) < now() - make_interval(secs => :age)"),
                    {"status": status, "age": age.total_seconds()},
                )


async def run_due(kinds: list[str] | None = None, *, tenant_id: uuid.UUID | None = None, limit: int = 1000) -> int:
    """Runs every job that's due now, one after another, in this process —
    for tests and scripts (the worker does this continuously). Returns how
    many ran."""
    load_handlers()
    ran = 0
    while ran < limit:
        claimed = await _claim(1, kinds, tenant_id)
        if not claimed:
            return ran
        await run_job(claimed[0])
        ran += 1
    return ran


async def _release_own() -> None:
    """On shutdown: this worker's unfinished jobs go straight back in the
    queue, without the attempt counting against them."""
    async with AsyncSessionLocal() as db:
        async with db.begin():
            await db.execute(
                text(
                    "UPDATE jobs SET status = 'queued', attempts = GREATEST(attempts - 1, 0), locked_by = NULL, locked_until = NULL "
                    "WHERE status = 'running' AND locked_by = :worker"
                ),
                {"worker": WORKER_ID},
            )


async def run_worker() -> None:
    """Claims and runs jobs until cancelled, up to CONCURRENCY at a time."""
    load_handlers()
    log.info("job worker %s started (up to %s at a time)", WORKER_ID, CONCURRENCY)
    running: set[asyncio.Task] = set()
    last_upkeep = datetime.min.replace(tzinfo=timezone.utc)
    try:
        while True:
            try:
                now = datetime.now(timezone.utc)
                if now - last_upkeep > timedelta(seconds=LEASE_SECONDS):
                    last_upkeep = now
                    await reclaim_expired()
                    await purge_finished()
                free = CONCURRENCY - len(running)
                claimed = await _claim(free) if free > 0 else []
                for job in claimed:
                    task = asyncio.create_task(run_job(job))
                    running.add(task)
                    task.add_done_callback(running.discard)
            except Exception:  # noqa: BLE001 - a database hiccup mustn't stop the worker
                log.exception("job worker loop failed; carrying on")
                claimed = []
            if not claimed:
                await asyncio.sleep(POLL_SECONDS)
    finally:
        for task in running:
            task.cancel()
        await asyncio.gather(*running, return_exceptions=True)
        with contextlib.suppress(Exception):
            await _release_own()


# ---- for the API ------------------------------------------------------------------------------------------


async def retry_now(db: AsyncSession, tenant_id: uuid.UUID, job_id: uuid.UUID) -> bool:
    """A failed job back in the queue, due now, with a fresh set of attempts."""
    row = await db.execute(
        text(
            "UPDATE jobs SET status = 'queued', attempts = 0, run_after = now(), last_error = NULL, finished_at = NULL "
            "WHERE id = :id AND tenant_id = :tenant AND status IN ('failed', 'cancelled') "
            # Not while a newer job under the same dedupe key is waiting — that one supersedes it.
            "AND NOT EXISTS (SELECT 1 FROM jobs q WHERE q.dedupe_key = jobs.dedupe_key AND q.status = 'queued') RETURNING id"
        ),
        {"id": job_id, "tenant": tenant_id},
    )
    return row.first() is not None


async def cancel(db: AsyncSession, tenant_id: uuid.UUID, job_id: uuid.UUID) -> bool:
    """A queued job that hasn't started won't run."""
    row = await db.execute(
        text("UPDATE jobs SET status = 'cancelled', finished_at = now() WHERE id = :id AND tenant_id = :tenant AND status = 'queued' RETURNING id"),
        {"id": job_id, "tenant": tenant_id},
    )
    return row.first() is not None