"""Worker entry point: `python -m aix.main`. Claims jobs, dispatches handlers (design §5.2)."""
from __future__ import annotations
import logging
import select
import signal
import time
import traceback

from . import jobs
from .db import connect
from .records import raise_flag
from .settings import settings

log = logging.getLogger("aix.worker")
_stop = False


def _handlers():
    # Imported lazily so a broken optional dependency doesn't stop unrelated job types.
    from .pipeline import rasterize, extract, assemble, score, pdf
    return {
        "noop": lambda conn, job: {"ok": True},
        "rasterize": rasterize.handle,
        "extract_page": extract.handle,
        "assemble": assemble.handle,
        "score": score.handle,
        "rescore": score.handle_rescore,
        "pdf": pdf.handle,
    }


def run_one(conn, handlers=None, worker_id: str | None = None) -> bool:
    """Claim and run one job. Returns False if the queue was empty. Used by the loop and tests."""
    handlers = handlers or _handlers()
    with conn.transaction():
        job = jobs.claim(conn, worker_id or settings.worker_id, list(handlers))
    if not job:
        return False
    log.info("job start", extra={"job_id": job["id"], "type": job["type"]})
    try:
        with conn.transaction():
            result = handlers[job["type"]](conn, job)
            jobs.complete(conn, job["id"], result)
    except jobs.RetryableError as e:
        with conn.transaction():
            status = jobs.fail(conn, job, str(e), retryable=True)
            if status == "failed":
                _flag_failure(conn, job, str(e))
    except Exception as e:  # FatalError and anything unexpected
        err = f"{type(e).__name__}: {e}\n{traceback.format_exc(limit=5)}"
        with conn.transaction():
            jobs.fail(conn, job, err, retryable=not isinstance(e, jobs.FatalError))
            if isinstance(e, jobs.FatalError) or job["attempts"] >= job["max_attempts"]:
                _flag_failure(conn, job, str(e))
    return True


def _flag_failure(conn, job, error: str) -> None:
    p = job["payload_json"]
    batch_id = p.get("batch_id")
    owned = batch_id and job["tenant_id"] and conn.execute(
        "SELECT 1 FROM batches WHERE id=%s AND tenant_id=%s", (batch_id, job["tenant_id"])).fetchone()
    if owned:
        raise_flag(conn, tenant_id=job["tenant_id"], batch_id=batch_id, type_="processing_failed",
                   script_id=p.get("script_id"), page_id=p.get("page_id"),
                   payload={"job_type": job["type"], "error": error[:500]})
        if p.get("script_id"):
            conn.execute("UPDATE scripts SET status='failed', failure_reason=%s WHERE id=%s AND tenant_id=%s",
                         (f"{job['type']}: {error[:200]}", p["script_id"], job["tenant_id"]))


def main() -> None:
    logging.basicConfig(level=logging.INFO, format='{"t":"%(asctime)s","lvl":"%(levelname)s","msg":"%(message)s"}')
    signal.signal(signal.SIGTERM, lambda *_: globals().__setitem__("_stop", True))
    conn = connect(autocommit=True)   # explicit transaction() blocks only
    listener = connect(autocommit=True)
    listener.execute("LISTEN aix_jobs")
    handlers = _handlers()
    last_reap = 0.0
    while not _stop:
        if time.monotonic() - last_reap > 60:
            with conn.transaction():
                jobs.reap_stale(conn)
            last_reap = time.monotonic()
        if run_one(conn, handlers):
            continue
        # Sleep until NOTIFY or poll interval.
        select.select([listener.fileno()], [], [], settings.poll_interval_s)
        for _ in listener.notifies(timeout=0):
            pass


if __name__ == "__main__":
    main()