from typing import Awaitable, Callable
from sqlalchemy.ext.asyncio import AsyncSession
from app.models.data_ingestion import DataIngestionJob

# Handler signature:
#   handler(job, job_session, data_session) -> dict (becomes job.meta on success)
HandlerFn = Callable[
    [DataIngestionJob, AsyncSession, AsyncSession],
    Awaitable[dict],
]

_REGISTRY: dict[str, HandlerFn] = {}

def register(job_type: str):
    """Decorator to register a handler for a job_type."""
    def decorator(fn: HandlerFn) -> HandlerFn:
        if job_type in _REGISTRY:
            raise ValueError(f"job_type {job_type!r} already registered")
        _REGISTRY[job_type] = fn
        return fn
    return decorator

def get_handler(job_type: str) -> HandlerFn:
    handler = _REGISTRY.get(job_type)
    if handler is None:
        raise ValueError(f"No handler registered for job_type={job_type!r}")
    return handler

import asyncio
from app.db import SessionLocal
from app.core.logging import get_logger
from app.models.data_ingestion import IngestionResult, IngestionState
from app.repositories.data_ingestion import DataIngestionRepository
from app.tasks._pod_id import POD_ID
from app.tasks.registry import get_handler

logger = get_logger(__name__)

async def run_job(job_id: int) -> None:
    """
    Single dispatch path for every job_type. Used by:
      - endpoints (fire_and_forget(run_job(id)) after creating a job)
      - chain_job (fire_and_forget(run_job(child_id)) after a parent commits)
      - the safety poller (Plan A)
    """
    async with SessionLocal() as job_session, SessionLocal() as data_session:
        repo = DataIngestionRepository(job_session)
        job = await repo.get_job_by_id(job_id)
        if job is None:
            logger.error(f"run_job: job {job_id} not found")
            return
        if job.job_type is None:
            logger.error(f"run_job: job {job_id} has no job_type — refusing to dispatch")
            return

        # Capture job_type while it's narrowed to ``str`` by the check above;
        # the post-claim re-fetch widens it back to Optional[str].
        job_type: str = job.job_type

        if not await repo.claim_job(job_id, POD_ID):
            return  # another pod claimed it, or attempts exhausted, or finished

        # claim_job ran a raw SQL UPDATE (state=RUNNING, attempts++,
        # locked_by, locked_at, AND started_at via func.coalesce — atomic
        # with the RUNNING transition, see PR #1026).  The in-memory `job`
        # still reflects the pre-claim row; re-fetch so handlers see the
        # authoritative post-claim state.  A vanished row here = preempted
        # in the gap (treat as such).
        job = await repo.get_job_by_id(job_id)
        if job is None:
            logger.warning(f"run_job: job {job_id} disappeared after claim — exiting")
            return

        # Plain asyncio.create_task (NOT fire_and_forget): cancellation in
        # the finally block is the expected shutdown path, and
        # fire_and_forget's deliberate cancellation-WARNING (kept loud
        # for diagnosing the 310-B incident) would fire on every successful
        # run, drowning out genuine cancellations elsewhere.
        heartbeat_task = asyncio.create_task(
            _heartbeat_loop(job_id), name=f"heartbeat-{job_id}"
        )

        try:
            try:
                handler = get_handler(job_type)
                meta = await handler(job, job_session, data_session)
                status_message = str(meta.get("status_message", "Success"))
                metadata = dict(meta)
                result = meta.get("result", IngestionResult.SUCCESS)
                handler_succeeded = True
            except Exception as exc:
                logger.exception(
                    f"run_job: handler for job_type={job_type!r} failed (job {job_id})"
                )
                status_message = str(exc)
                metadata = {}
                result = IngestionResult.ERROR
                handler_succeeded = False

            # Preemption check covers BOTH success AND error paths.  If a
            # stale-lock sweep recovered this row mid-handler, a different
            # pod may now own it; our writes — successful or error — must
            # NOT race with the new owner.  Roll back data and skip the
            # state update; the new owner closes out.
            current = await repo.get_job_by_id(job_id)
            if current is None or current.locked_by != POD_ID:
                logger.warning(
                    f"run_job: job {job_id} preempted "
                    f"(locked_by={current and current.locked_by!r}); "
                    "rolling back data writes and exiting without updating job state"
                )
                await data_session.rollback()
                return

            if handler_succeeded:
                await data_session.commit()
            else:
                await data_session.rollback()
            await repo.update_ingestion_job(
                job_id,
                status_message=status_message,
                metadata=metadata,
                state=IngestionState.FINISHED,
                result=result,
                # NOTE: no finished_at parameter — PR #1026 dropped that
                # opt-in flag and made FINISHED auto-stamp the column.
            )
            await job_session.commit()
        finally:
            heartbeat_task.cancel()
            try:
                await heartbeat_task
            except asyncio.CancelledError:
                pass

asyncio.create_task(run_job(created.id))

# repo helper
async def heartbeat(self, job_id: int) -> None:
    await self.session.execute(
        update(DataIngestionJob)
        .where(col(DataIngestionJob.id) == job_id,
               col(DataIngestionJob.locked_by) == POD_ID,  # only OUR job
               col(DataIngestionJob.state) == IngestionState.RUNNING)
        .values(locked_at=func.now())
    )
    await self.session.commit()

current = await repo.get_job_by_id(job_id)
if current is None or current.locked_by != POD_ID:
    logger.warning(
        f"Job {job_id} was preempted (locked_by={current and current.locked_by!r}); "
        "rolling back our work and exiting"
    )
    await data_session.rollback()
    return  # do NOT update job state — the new owner will

# backend/app/tasks/runner.py

async def chain_job(
    parent: DataIngestionJob,
    *,
    job_type: str,
    session: AsyncSession,
    module_type_id: Optional[int] = None,
    data_entry_type_id: Optional[int] = None,
    year: Optional[int] = None,
    config: Optional[dict] = None,
    target_type: TargetType = TargetType.DATA_ENTRIES,
    ingestion_method: IngestionMethod = IngestionMethod.computed,
    entity_type: EntityType = EntityType.MODULE_PER_YEAR,
) -> int:
    """
    Create a child job that inherits parent's pipeline_id and fire it via
    fire_and_forget. Safety poller picks up if pod crashes between commit
    and fire_and_forget.

    ``module_type_id`` and ``year`` inherit from the parent when the
    caller passes None.  ``data_entry_type_id`` is intentionally NOT
    inherited: a multi-det parent (e.g. a FACTORS ingest spanning
    several dets) fans out to one child per det, so the caller MUST
    pass the specific det per call.

    Returns child job_id.  If ``parent.pipeline_id`` is None, generates
    a fresh UUID and persists it on the parent BEFORE creating the
    child — so a pod-crash-then-recovery of the parent doesn't
    generate a different UUID and orphan the child.
    """
    repo = DataIngestionRepository(session)

    pipeline_id = parent.pipeline_id
    if pipeline_id is None:
        pipeline_id = uuid4()
        parent.pipeline_id = pipeline_id
        session.add(parent)
        await session.commit()

    child = DataIngestionJob(
        job_type           = job_type,
        module_type_id     = module_type_id if module_type_id is not None else parent.module_type_id,
        data_entry_type_id = data_entry_type_id,  # NO inheritance — see docstring
        year               = year if year is not None else parent.year,
        target_type        = target_type,
        ingestion_method   = ingestion_method,
        entity_type        = entity_type,
        state              = IngestionState.NOT_STARTED,
        is_current         = False,
        pipeline_id        = pipeline_id,
        # NULL means "runnable immediately" — claim_job's WHERE treats
        # NULL run_after as eligible.  Matches the existing
        # ingestion_tasks.py recalc-job creation pattern.
        run_after          = None,
        meta               = {"config": config or {}, "parent_job_id": parent.id},
    )
    created = await repo.create_ingestion_job(child)
    await session.commit()
    fire_and_forget(run_job(created.id), name=f"run_job-{created.id}")
    return created.id

@register("factor_ingest")
async def factor_ingest_handler(job, job_session, data_session):
    # ... existing factor upsert logic (Plan B) ...
    # On success, fan out to one emission_recalc per stale (module, det):
    if final_result != IngestionResult.ERROR:
        rows = await DataIngestionRepository(job_session).get_recalculation_status_by_year(job.year)
        for row in rows:
            if not row["needs_recalculation"]:
                continue
            if job.module_type_id is not None and row["module_type_id"] != job.module_type_id:
                continue
            if job.data_entry_type_id is not None and row["data_entry_type_id"] != job.data_entry_type_id:
                continue
            await chain_job(
                job,
                job_type="emission_recalc",
                module_type_id=row["module_type_id"],
                data_entry_type_id=row["data_entry_type_id"],
                year=job.year,
                config={"data_entry_type_id": row["data_entry_type_id"]},
                session=job_session,
            )
    return {"upsert_count": ..., "recalc_jobs_chained": len(...)}

# backend/app/tasks/_poller.py — after Plan C
from app.tasks.runner import run_job


async def dispatch_job(job: DataIngestionJob, pod_id: str) -> None:
    """Plan 310-C cutover: every job_type funnels through ``run_job`` —
    the registry resolves the handler from the row's job_type."""
    jid = job.id
    if jid is None:
        logger.warning("Job has no ID — skipping")
        return
    await run_job(jid)


def _pending_runner_jobs_query(limit: int = 10):
    """Same predicates as Plan-A's ``_pending_jobs_query`` plus a
    ``job_type IS NOT NULL`` filter so legacy in-flight jobs (created
    pre-Plan-C with a NULL job_type) don't get dispatched through a
    runner that has no registered handler for them."""
    return (
        select(DataIngestionJob)
        .where(
            col(DataIngestionJob.state) == IngestionState.NOT_STARTED,
            col(DataIngestionJob.job_type).is_not(None),  # NEW
            or_(
                col(DataIngestionJob.run_after).is_(None),
                col(DataIngestionJob.run_after) <= func.now(),
            ),
            col(DataIngestionJob.locked_by).is_(None),
            col(DataIngestionJob.attempts) < col(DataIngestionJob.max_attempts),
        )
        .with_for_update(skip_locked=True)
        .limit(limit)
    )

ALTER TABLE data_ingestion_jobs
  ADD COLUMN started_at  TIMESTAMPTZ,
  ADD COLUMN finished_at TIMESTAMPTZ;

async def set_started_at(self, job_id: int) -> None:
    await self.session.execute(
        update(DataIngestionJob)
        .where(
            DataIngestionJob.id == job_id,
            DataIngestionJob.started_at.is_(None),
        )
        .values(started_at=func.now())
    )

# update_ingestion_job extended with optional finished_at: bool = False arg that
# sets finished_at=now() when state transitions to FINISHED.

SELECT
    job_type,
    state,
    result,
    count(*)                                                          AS jobs,
    avg(extract(epoch from (finished_at - started_at)))               AS avg_duration_s,
    percentile_cont(0.95) within group (order by extract(epoch from (finished_at - started_at)))
                                                                      AS p95_duration_s,
    sum(case when attempts > 1 then 1 else 0 end)                     AS retried_jobs
FROM data_ingestion_jobs
WHERE created_at > now() - interval '7 days'
  AND job_type IS NOT NULL
GROUP BY 1, 2, 3
ORDER BY 1, 2;

SELECT id, job_type, state, result, started_at, finished_at, status_message
  FROM data_ingestion_jobs
 WHERE pipeline_id = :pipeline_id
 ORDER BY id;

current_user: User = Depends(
    require_permission("backoffice.data_management", "view")
),

fired_children = await chain_job(parent, ...)
await update_ingestion_job(
    parent.id,
    extra_metadata={"children_fired": len(fired_children),
                    "child_pipeline_id": str(pipeline_id)},
)