netx/netx_api/biz_state/collect_runner.py

"""Collect runner: expand task items → CLI → match → parse → batch rows."""

from __future__ import annotations

import json
import logging
import re
import threading
from collections.abc import Iterable
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from typing import Any
from uuid import uuid4

from fastapi import HTTPException
from sqlalchemy.exc import DBAPIError, OperationalError

from ..cli_creds import cli_creds_skip_reason
from ..cli_resolve import resolve_cli_target
from ..cli_timeout import run_cli_with_timeout
from ..config import settings
from ..db import SessionLocal
from ..lldp_shared import resolve_vendor_key
from ..models import (
    BizStateBatch,
    BizStateBatchCommand,
    BizStateEvent,
    BizStateLldpNeighbor,
    BizStateMetricRow,
    BizStateTask,
    BizStateTaskItem,
    BizStateTaskItemBinding,
)
from ..ne_netmiko import disable_target_paging, send_show_command
from ..ne_session_factory import close_netmiko_connection, open_netmiko_connection
from ..timeutil import utcnow_naive
from .collect_session import (
    CachedCommand,
    CollectSession,
    ResolvedAux,
    build_parse_bundle,
    resolve_aux_command,
    run_primary_with_bundle,
)
from .command_match import (
    EXPAND_ALL_COMMAND,
    expand_bindings_from_discover_records,
    expand_from_bindings,
    match_command,
    normalize_command,
)
from .parsers import get_parser, run_parser
from .profiles import get_profile
from .collect_stop import (
    STOP_USER_MESSAGE,
    clear_stop_requested,
    is_stop_requested,
    register_lane_holder,
    unregister_lane_holder,
)

_log = logging.getLogger("netx.biz_state.runner")

# concrete, params, profile_id, item_id, mode
WorkItem = tuple[str, dict[str, str], str, str, str]

_heavy_pool_lock = threading.Lock()
_heavy_pool: ThreadPoolExecutor | None = None


def _heavy_cli_pool() -> ThreadPoolExecutor:
    global _heavy_pool
    with _heavy_pool_lock:
        if _heavy_pool is None:
            n = max(1, int(getattr(settings, "biz_state_heavy_workers", 4) or 4))
            _heavy_pool = ThreadPoolExecutor(max_workers=n, thread_name_prefix="biz-heavy")
        return _heavy_pool


def work_item_lane(profile_id: str, mode: str) -> str:
    """Return collect lane for a work item (custom_raw always light)."""
    if str(mode or "").strip() == "custom":
        return "light"
    profile = get_profile(profile_id) or _resolve_collect_profile(profile_id)
    lane = str(getattr(profile, "collect_lane", "light") or "light").strip().lower() if profile else "light"
    return "heavy" if lane == "heavy" else "light"


def partition_work(work: list[WorkItem]) -> tuple[list[WorkItem], list[WorkItem]]:
    """Split work into (light, heavy) lanes."""
    light: list[WorkItem] = []
    heavy: list[WorkItem] = []
    for item in work:
        if work_item_lane(item[2], item[4]) == "heavy":
            heavy.append(item)
        else:
            light.append(item)
    return light, heavy


def _utcnow() -> datetime:
    return utcnow_naive()


def _format_error(exc: BaseException) -> str:
    return f"{type(exc).__name__}: {exc}"[:1020]


def _append_event(db, *, task_id: str, message: str, level: str = "error") -> None:
    msg = str(message or "").strip()
    if not msg or not task_id:
        return
    db.add(
        BizStateEvent(
            id=uuid4().hex,
            task_id=task_id,
            level=str(level or "error")[:16],
            message=msg[:4000],
            created_at=_utcnow(),
        )
    )


def _bindings_for_item(db, item_id: str) -> list[dict[str, str]]:
    rows = (
        db.query(BizStateTaskItemBinding)
        .filter(BizStateTaskItemBinding.item_id == item_id)
        .all()
    )
    # Keep one row per binding value (same placeholder may appear many times).
    return [
        {"placeholder": str(r.placeholder or "").strip(), "value": str(r.value or "").strip()}
        for r in rows
        if str(r.placeholder or "").strip() and str(r.value or "").strip()
    ]


def _resolve_collect_profile(profile_id: str):
    """Resolve task-item profile; remap disabled if_intf → config_interface."""
    pid = str(profile_id or "").strip()
    if not pid:
        return None
    profile = get_profile(pid)
    if profile is None:
        return None
    if profile.enabled:
        return profile
    # Legacy IF VRF check merged into Config Interface Intent — VRF is a subset.
    if profile.metric_id == "if_intf" or pid.endswith(".if_intf"):
        vk = str(profile.vendor_key or "zte").strip() or "zte"
        remapped = get_profile(f"{vk}.config_interface") or get_profile("zte.config_interface")
        if remapped and remapped.enabled:
            return remapped
    return None


_METRIC_CHUNK = 2000


def _persist_lldp_rows(
    db,
    *,
    batch: BizStateBatch,
    cmd_row: BizStateBatchCommand,
    records: Iterable[dict[str, Any]],
) -> int:
    """Insert LLDP rows from an iterable (chunk-friendly)."""
    n = 0
    seen: set[tuple[str, str, str]] = set()
    buf: list[BizStateLldpNeighbor] = []
    for rec in records:
        if not isinstance(rec, dict):
            continue
        local_if = str(rec.get("local_if") or "").strip()[:128]
        remote_sys = str(rec.get("remote_sys") or "").strip()[:256]
        remote_if = str(rec.get("remote_if") or "").strip()[:128]
        if not local_if and not remote_sys and not remote_if:
            continue
        key = (local_if, remote_sys, remote_if)
        if key in seen:
            continue
        seen.add(key)
        buf.append(
            BizStateLldpNeighbor(
                id=uuid4().hex,
                batch_id=batch.id,
                batch_command_id=cmd_row.id,
                task_id=batch.task_id,
                ne_id=batch.ne_id,
                local_if=local_if,
                remote_sys=remote_sys,
                remote_if=remote_if,
                remote_ip=str(rec.get("remote_ip") or "")[:128],
                protocol=str(rec.get("protocol") or "lldp")[:32],
                collected_at=_utcnow(),
            )
        )
        n += 1
        if len(buf) >= _METRIC_CHUNK:
            db.add_all(buf)
            db.flush()
            buf.clear()
    if buf:
        db.add_all(buf)
        db.flush()
    return n


def _db_is_postgres(db) -> bool:
    try:
        bind = db.get_bind() if hasattr(db, "get_bind") else getattr(db, "bind", None)
        name = str(getattr(getattr(bind, "dialect", None), "name", "") or "").lower()
        return name in ("postgresql", "postgres")
    except Exception:
        return False


def _persist_metric_chunk_bulk(
    db,
    *,
    batch: BizStateBatch,
    cmd_row: BizStateBatchCommand,
    metric_id: str,
    chunk: list[dict[str, Any]],
    seq_start: int,
) -> int:
    if not chunk:
        return 0
    now = _utcnow()
    buf = [
        {
            "id": uuid4().hex,
            "batch_id": batch.id,
            "batch_command_id": cmd_row.id,
            "task_id": batch.task_id,
            "ne_id": batch.ne_id,
            "metric_id": metric_id,
            "seq": seq_start + i,
            "data_json": dict(rec),
            "collected_at": now,
        }
        for i, rec in enumerate(chunk)
        if isinstance(rec, dict) and rec
    ]
    if not buf:
        return 0
    db.bulk_insert_mappings(BizStateMetricRow, buf)
    db.flush()
    return len(buf)


def _persist_metric_chunk_copy(
    db,
    *,
    batch: BizStateBatch,
    cmd_row: BizStateBatchCommand,
    metric_id: str,
    chunk: list[dict[str, Any]],
    seq_start: int,
) -> int:
    """Postgres fast path via execute_values; falls back to bulk on error."""
    if not chunk:
        return 0
    now = _utcnow()
    rows: list[tuple[Any, ...]] = []
    for i, rec in enumerate(chunk):
        if not isinstance(rec, dict) or not rec:
            continue
        rows.append(
            (
                uuid4().hex,
                batch.id,
                cmd_row.id,
                batch.task_id or "",
                batch.ne_id or "",
                metric_id,
                seq_start + i,
                json.dumps(rec, ensure_ascii=False, default=str, separators=(",", ":")),
                now,
            )
        )
    if not rows:
        return 0
    try:
        from psycopg2.extras import execute_values  # type: ignore
    except ImportError:
        return _persist_metric_chunk_bulk(
            db,
            batch=batch,
            cmd_row=cmd_row,
            metric_id=metric_id,
            chunk=chunk,
            seq_start=seq_start,
        )
    # Unwrap SQLAlchemy connection → DBAPI (psycopg2) connection
    sa_conn = db.connection()
    dbapi = sa_conn.connection
    driver = getattr(dbapi, "dbapi_connection", None) or getattr(
        dbapi, "driver_connection", None
    ) or dbapi
    sql = (
        "INSERT INTO biz_state_metric_row "
        "(id,batch_id,batch_command_id,task_id,ne_id,metric_id,seq,data_json,collected_at) "
        "VALUES %s"
    )
    with driver.cursor() as cur:
        execute_values(
            cur,
            sql,
            rows,
            template="(%s,%s,%s,%s,%s,%s,%s,%s::jsonb,%s)",
            page_size=len(rows),
        )
    db.flush()
    return len(rows)


def _persist_metric_rows_from_spool(
    db,
    *,
    batch: BizStateBatch,
    cmd_row: BizStateBatchCommand,
    metric_id: str,
    records_rel_path: str,
) -> int:
    """Stream JSONL → DB in chunks (never loads full table into memory)."""
    from .spool import iter_record_chunks

    mid = str(metric_id or "").strip()
    if not mid or not records_rel_path:
        return 0
    use_copy = _db_is_postgres(db)
    n = 0
    for chunk in iter_record_chunks(records_rel_path, chunk_size=_METRIC_CHUNK):
        if use_copy:
            try:
                added = _persist_metric_chunk_copy(
                    db,
                    batch=batch,
                    cmd_row=cmd_row,
                    metric_id=mid,
                    chunk=chunk,
                    seq_start=n,
                )
            except Exception:
                _log.exception("biz_state COPY failed; falling back to bulk")
                use_copy = False
                added = _persist_metric_chunk_bulk(
                    db,
                    batch=batch,
                    cmd_row=cmd_row,
                    metric_id=mid,
                    chunk=chunk,
                    seq_start=n,
                )
        else:
            added = _persist_metric_chunk_bulk(
                db,
                batch=batch,
                cmd_row=cmd_row,
                metric_id=mid,
                chunk=chunk,
                seq_start=n,
            )
        n += added
    return n


def _persist_lldp_rows_from_spool(
    db,
    *,
    batch: BizStateBatch,
    cmd_row: BizStateBatchCommand,
    records_rel_path: str,
) -> int:
    from .spool import iter_record_chunks

    if not records_rel_path:
        return 0
    n = 0
    for chunk in iter_record_chunks(records_rel_path, chunk_size=_METRIC_CHUNK):
        n += _persist_lldp_rows(db, batch=batch, cmd_row=cmd_row, records=chunk)
    return n


_GENERIC_METRICS = {
    "isis_adjacency",
    "interface_brief",
    "interface_detail",
    "arp",
    "if_intf",
    "nd6_cache",
    "bgp_peer",
    "ospf_neighbor",
    "vrrp",
    "optical_brief",
    "bgp_route",
    "ip_route",
    "ipv6_route",
    "l2vpn_pw",
    "l2vpn_pw_detail",
    "l2vpn_mac",
    "evpn_mac",
    "config_vrf",
    "config_interface",
    "config_bgp_peer",
    "config_l2vpn_pw",
    "config_static_route",
    "config_ospf",
    "config_isis",
}


def _run_primary_parse_job(job: Any) -> tuple[bool, bool]:
    """Parse primary+aux raws, write spool, submit persist. Returns (any_ok, any_fail)."""
    from .parse_pool import AuxRawCapture, PrimaryParseJob
    from .persist_pool import get_persist_pool
    from .spool import SpooledCommand, count_text_lines, write_meta, write_records

    if not isinstance(job, PrimaryParseJob):
        return False, True

    batch_id = job.batch_id
    cmd_id = job.cmd_id
    pending: list[SpooledCommand] = []
    any_ok = False
    any_fail = False

    def _declared_total(raw: str) -> int:
        text = (raw or "").replace("\u00a0", " ")
        m = re.search(r"(?i)total\s+number\s+of\s+routes\s*:\s*(\d+)", text)
        return int(m.group(1)) if m else 0

    def _flush_item(item: SpooledCommand, *, records: Any = None) -> None:
        if records is not None and item.persist_kind:
            rel, n = write_records(batch_id, item.id, records)
            item.records_rel_path = rel
            item.row_count = int(n or 0)
        try:
            write_meta(batch_id, item.id, item.to_meta())
        except Exception:
            _log.exception("biz_state write meta failed cmd=%s", item.id)
        pending.append(item)

    aux_results: dict[str, CachedCommand] = {}
    resolved_aux: list[ResolvedAux] = []

    for cap in list(job.aux_captures or []):
        if not isinstance(cap, AuxRawCapture):
            continue
        from .profiles import get_profile as _gp

        prof = _gp(cap.profile_id)
        entry = CachedCommand(
            raw=cap.raw,
            records=list(cap.records or []),
            fsm_tables=dict(cap.fsm_tables or {}),
            ok=bool(cap.ok),
            error=str(cap.error or ""),
            cmd_row_id=cap.aux_id,
        )
        if entry.ok and not entry.records and cap.parser_id and get_parser(cap.parser_id):
            try:
                records, fsm_tables, _keys = run_parser(
                    cap.parser_id,
                    raw_text=cap.raw,
                    vendor=job.vendor,
                    device_type=job.device_type,
                    command=cap.textfsm_command or cap.command,
                    textfsm_command=cap.textfsm_command or "",
                    params=dict(job.merged_params or {}),
                )
                entry.records = list(records or [])
                entry.fsm_tables = dict(fsm_tables or {})
            except Exception as exc:
                entry.ok = False
                entry.error = f"parse: {type(exc).__name__}: {exc}"

        aux_results[cap.key] = entry
        if prof is not None:
            resolved_aux.append(
                ResolvedAux(
                    key=cap.key,
                    profile_id=cap.profile_id,
                    command=cap.command,
                    textfsm_command=cap.textfsm_command or cap.command,
                    parser_id=cap.parser_id,
                    rule_keys=tuple(cap.rule_keys or ()),
                    profile=prof,
                )
            )

        aux_sp = SpooledCommand(
            id=cap.aux_id,
            batch_id=batch_id,
            task_item_id=job.task_item_id,
            profile_id=cap.profile_id,
            parser_id=cap.parser_id,
            metric_id=cap.metric_id,
            raw_command=cap.command[:512],
            params_json={},
            raw_rel_path=cap.raw_rel_path,
            raw_line_count=count_text_lines(cap.raw),
        )
        if cap.cache_hit and entry.ok and entry.records:
            aux_sp.parse_status = "aux_cached"
            aux_sp.message = (
                f"aux_for={cmd_id};cache_hit;src={entry.cmd_row_id}"
            )[:1020]
            aux_sp.row_count = len(entry.records or [])
        elif not entry.ok:
            aux_sp.parse_status = "aux_failed"
            aux_sp.message = f"aux_for={cmd_id};{entry.error}"[:1020]
            any_fail = True
        else:
            aux_sp.parse_status = "aux"
            aux_sp.message = f"aux_for={cmd_id}"[:1020]
            aux_sp.row_count = len(entry.records or [])

        persist_recs: list[dict[str, Any]] | None = None
        if entry.ok and entry.records and cap.metric_id in _GENERIC_METRICS:
            persist_key = (normalize_command(cap.command), cap.metric_id)
            do_persist = False
            persisted = job.persisted
            if persisted is not None:
                if job.cache_lock is not None:
                    with job.cache_lock:
                        if persist_key not in persisted:
                            persisted.add(persist_key)
                            do_persist = True
                elif persist_key not in persisted:
                    persisted.add(persist_key)
                    do_persist = True
            if do_persist:
                aux_sp.persist_kind = "metric"
                persist_recs = list(entry.records)
        _flush_item(aux_sp, records=persist_recs)

    bundle = build_parse_bundle(
        primary_raw=job.raw_text,
        primary_parser_id=job.parser_id,
        aux_results=aux_results,
        resolved_aux=resolved_aux,
    )
    # Include aux raws that lacked a profile (still needed for multi-raw parsers).
    for cap in list(job.aux_captures or []):
        if not isinstance(cap, AuxRawCapture):
            continue
        if cap.key in bundle.raws:
            continue
        entry = aux_results.get(cap.key) or CachedCommand(ok=False)
        bundle.raws[cap.key] = entry.raw
        bundle.command_rules[cap.key] = list(cap.rule_keys or [])
        if entry.records:
            bundle.aux_records[cap.key] = list(entry.records)
        bundle.fsm_extra.update(entry.fsm_tables or {})
    primary = SpooledCommand(
        id=cmd_id,
        batch_id=batch_id,
        task_item_id=job.task_item_id,
        profile_id=job.profile_id,
        parser_id=job.parser_id,
        metric_id=job.metric_id,
        raw_command=job.concrete[:512],
        params_json=dict(job.merged_params or {}),
        raw_rel_path=job.raw_rel_path,
        raw_line_count=int(getattr(job, "raw_line_count", 0) or 0)
        or count_text_lines(job.raw_text),
        declared_total=_declared_total(job.raw_text),
    )
    try:
        records, fsm_tables, rule_keys = run_primary_with_bundle(
            job.parser_id,
            bundle=bundle,
            vendor=job.vendor,
            device_type=job.device_type,
            command=job.textfsm_command or job.concrete,
            textfsm_command=job.textfsm_command or "",
            params=dict(job.merged_params or {}),
            enrich_joins=list(job.enrich_joins or []),
        )
        hints = []
        if rule_keys:
            nonempty = [k for k in rule_keys if fsm_tables.get(k)]
            hints.append(f"fsm_keys={','.join(rule_keys)};hit={','.join(nonempty)}")
        if bundle.aux_records:
            hints.append(
                "aux="
                + ",".join(f"{k}:{len(v)}" for k, v in bundle.aux_records.items())
            )
        if job.enrich_joins:
            hints.append(
                "enrich=" + ",".join(getattr(j, "from_aux", "") for j in job.enrich_joins)
            )
        declared = int(primary.declared_total or 0)
        primary.parse_status = "ok"
        any_ok = True
        persist_recs = None
        if job.metric_id == "lldp_neighbor":
            primary.persist_kind = "lldp"
            persist_recs = records
        elif job.metric_id in _GENERIC_METRICS:
            primary.persist_kind = "metric"
            persist_recs = records
        _flush_item(primary, records=persist_recs)
        nrec = int(primary.row_count or 0)
        if declared > 0:
            hints.append(f"declared={declared};parsed={nrec}")
        if hints:
            primary.message = ";".join(hints)[:1020]
            try:
                write_meta(batch_id, primary.id, primary.to_meta())
            except Exception:
                pass
        _ = fsm_tables  # kept for hints above
    except Exception as exc:
        any_fail = True
        primary.parse_status = "failed"
        primary.message = f"parse: {_format_error(exc)}"
        _flush_item(primary)

    if pending:
        get_persist_pool().submit(batch_id, pending)
    return any_ok, any_fail


def _emit_task_event(*, task_id: str, message: str, level: str = "error") -> None:
    """Short-lived session for lane events (no long-held DB during CLI)."""
    if not task_id or not str(message or "").strip():
        return
    db = SessionLocal()
    try:
        _append_event(db, task_id=task_id, message=message, level=level)
        db.commit()
    except Exception:
        _log.exception("biz_state emit event failed task=%s", task_id)
        try:
            db.rollback()
        except Exception:
            pass
    finally:
        db.close()


def _flush_spooled_commands(
    batch_id: str,
    pending: list[Any],
) -> tuple[int, int]:
    """Insert SpooledCommand rows (+ metric/lldp) in one transaction. Returns (cmds, rows)."""
    from .spool import SpooledCommand, raw_max_bytes, read_raw_text, spool_file_size

    if not pending:
        return 0, 0
    items: list[SpooledCommand] = list(pending)
    pending.clear()

    def _write(db) -> tuple[int, int]:
        batch = db.get(BizStateBatch, batch_id)
        if not batch:
            return 0, 0
        rows_n = 0
        max_raw = raw_max_bytes()
        for item in items:
            raw = ""
            truncated = False
            line_count = int(getattr(item, "raw_line_count", 0) or 0)
            raw_size = 0
            if item.raw_rel_path:
                from .spool import count_file_lines

                raw_size = spool_file_size(item.raw_rel_path)
                if line_count <= 0:
                    try:
                        line_count = count_file_lines(item.raw_rel_path)
                    except Exception:
                        line_count = 0
                # Mega outputs stay on spool only — do not load full CLI into Postgres.
                if max_raw > 0 and raw_size > max_raw:
                    truncated = True
                    raw = (
                        f"[raw_on_spool={item.raw_rel_path}; size={raw_size}B; "
                        f"cap={max_raw}B; omitted_from_db]\n"
                    )
                else:
                    raw = read_raw_text(item.raw_rel_path, max_bytes=max_raw)
                    if max_raw > 0 and "[truncated" in raw:
                        truncated = True
            if line_count <= 0 and raw and not truncated:
                from .spool import count_text_lines

                line_count = count_text_lines(raw)
            msg = str(item.message or "").strip()
            if truncated:
                note = (
                    f"raw_on_spool@{raw_size}B"
                    if raw_size > max_raw > 0
                    else f"raw_truncated@{max_raw}B"
                )
                msg = f"{msg}; {note}" if msg else note
            cmd_row = BizStateBatchCommand(
                id=item.id,
                batch_id=batch_id,
                task_item_id=item.task_item_id,
                profile_id=item.profile_id,
                parser_id=item.parser_id,
                metric_id=item.metric_id,
                raw_command=str(item.raw_command or "")[:512],
                params_json=dict(item.params_json or {}),
                parse_status=item.parse_status,
                message=msg[:1020],
                raw_text=raw,
                row_count=int(item.row_count or 0),
                raw_line_count=int(line_count or 0),
                declared_total=int(getattr(item, "declared_total", 0) or 0),
                created_at=_utcnow(),
            )
            db.add(cmd_row)
            if item.persist_kind == "metric" and item.records_rel_path:
                mid = str(item.metric_id or "").strip()
                if mid:
                    n = _persist_metric_rows_from_spool(
                        db,
                        batch=batch,
                        cmd_row=cmd_row,
                        metric_id=mid,
                        records_rel_path=item.records_rel_path,
                    )
                    cmd_row.row_count = n
                    rows_n += n
            elif item.persist_kind == "lldp" and item.records_rel_path:
                n = _persist_lldp_rows_from_spool(
                    db,
                    batch=batch,
                    cmd_row=cmd_row,
                    records_rel_path=item.records_rel_path,
                )
                cmd_row.row_count = n
                rows_n += n
        db.commit()
        return len(items), rows_n

    try:
        cmds, rows = _run_db_with_reconnect(_write, label="biz_state_flush_spool")
        if cmds or rows:
            _bump_batch_progress(batch_id, add_cmds=cmds, add_rows=rows)
        return int(cmds or 0), int(rows or 0)
    except Exception:
        _log.exception("biz_state flush spool failed batch=%s n=%s", batch_id, len(items))
        # Put back so a later flush / finalize can retry.
        pending.extend(items)
        raise


def _finish_task(task_id: str, *, error: str = "") -> None:
    db = SessionLocal()
    try:
        task = db.get(BizStateTask, task_id)
        if not task:
            return
        task.collect_running = False
        if hasattr(task, "collect_queued_at"):
            task.collect_queued_at = None
        task.last_collect_ended_at = _utcnow()
        task.last_error = str(error or "")[:1020]
        task.updated_at = _utcnow()
        if error:
            _append_event(db, task_id=task_id, message=error, level="error")
        db.commit()
    finally:
        db.close()


def dispatch_collect(task_id: str, *, manual: bool = False) -> dict[str, Any]:
    """Enqueue a collect round; run inline when this process owns execution.

    Dedicated worker mode: only enqueue (claim loop runs the batch).
    Inline / non-dedicated: enqueue then atomically promote+execute (skip if
    another worker already claimed the batch).

    Returns the enqueue result dict (ok / queued / reason / batch_id / …).
    """
    from .claim import enqueue_collect

    result = enqueue_collect(task_id, manual=manual)
    if not result.get("queued"):
        return result
    batch_id = str(result.get("batch_id") or "")
    if not batch_id:
        return result
    execute_enqueued_batch(
        batch_id=batch_id,
        task_id=str(result.get("task_id") or task_id),
    )
    return result


def execute_enqueued_batch(*, batch_id: str, task_id: str) -> None:
    """Promote+run a queued batch when this process owns inline execution."""
    bid = str(batch_id or "").strip()
    if not bid:
        return
    if not _should_execute_inline():
        return
    # Atomic queued→running; if false, scheduler/worker already owns it.
    if not _try_claim_batch_for_execute(bid):
        return
    execute_claimed_batch(
        batch_id=bid,
        task_id=str(task_id or ""),
        source="",
        ne_id="",
        vendor="",
        device_type="",
    )


def _should_execute_inline() -> bool:
    """True when this process should run SSH after enqueue (not dedicated workers)."""
    if bool(getattr(settings, "run_inline_schedulers", True)):
        return True
    if not bool(getattr(settings, "biz_state_dedicated_workers", True)):
        return True
    return False


def _try_claim_batch_for_execute(batch_id: str) -> bool:
    """Promote queued→running only if still queued. Returns True iff we won the claim."""
    db = SessionLocal()
    try:
        batch = (
            db.query(BizStateBatch)
            .filter(BizStateBatch.id == batch_id)
            .with_for_update()
            .one_or_none()
        )
        if not batch or str(batch.status or "") != "queued":
            return False
        batch.status = "running"
        batch.message = ""
        db.commit()
        return True
    except Exception:
        _log.exception("biz_state claim-for-execute failed batch=%s", batch_id)
        try:
            db.rollback()
        except Exception:
            pass
        return False
    finally:
        db.close()


def _promote_queued_batch(batch_id: str) -> bool:
    """Backward-compatible alias for atomic claim. """
    return _try_claim_batch_for_execute(batch_id)

def execute_claimed_batch(
    *,
    batch_id: str,
    task_id: str,
    source: str = "",
    ne_id: str = "",
    vendor: str = "",
    device_type: str = "",
) -> None:
    """Run CLI+persist for an already-claimed (status=running) batch."""
    db = SessionLocal()
    try:
        batch = db.get(BizStateBatch, batch_id)
        task = db.get(BizStateTask, task_id) if task_id else None
        if not batch:
            return
        if not task_id:
            task_id = str(batch.task_id or "")
            task = db.get(BizStateTask, task_id) if task_id else None
        source = str(source or batch.source or (task.source if task else "") or "managed")
        ne_id = str(ne_id or batch.ne_id or (task.ne_id if task else "") or "")
        vendor = str(vendor or batch.vendor or (task.vendor if task else "") or "")
        device_type = str(device_type or (task.device_type if task else "") or "")
        if not task_id:
            return
    finally:
        db.close()

    error = ""
    try:
        if is_stop_requested(batch_id):
            _finalize_batch_status(
                batch_id=batch_id,
                task_id=task_id,
                cmd_count=0,
                total_rows=0,
                any_fail=False,
                any_ok=False,
                lane_errors=[],
                stopped=True,
            )
            error = STOP_USER_MESSAGE
        else:
            _run_collect_session(
                task_id=task_id,
                batch_id=batch_id,
                source=source,
                ne_id=ne_id,
                vendor=vendor,
                device_type=device_type,
            )
    except Exception as exc:
        _log.exception("biz_state collect failed task=%s batch=%s", task_id, batch_id)
        error = _format_error(exc)
        if "_stopped" in error or is_stop_requested(batch_id):
            error = STOP_USER_MESSAGE
        try:
            _fail_batch_status(batch_id, error)
        except Exception:
            _log.exception(
                "biz_state fail-batch after collect error failed batch=%s", batch_id
            )
    finally:
        _finish_task(task_id, error=error)


def _run_collect_lane(
    *,
    work: list[WorkItem],
    batch_id: str,
    creds: dict[str, Any],
    vendor_eff: str,
    device_type_eff: str,
    vendor_key: str,
    per_cmd: int,
    cap: int,
    label: str,
    shared_cache: dict[str, Any] | None = None,
    cache_lock: Any | None = None,
    cmd_locks: dict[str, Any] | None = None,
    aux_persisted: set[tuple[str, str]] | None = None,
) -> tuple[int, int, bool, bool]:
    """Run one SSH lane: collect+parse to spool, flush to DB in batches."""
    if not work:
        return 0, 0, False, False

    from .spool import (
        SpooledCommand,
        persist_every_cmds,
        write_meta,
        write_raw_text,
        write_records,
    )

    budget = min(int(cap), int(per_cmd) * max(1, len(work)) + 90)
    holder: dict[str, Any] = {}
    flush_every = persist_every_cmds()
    register_lane_holder(batch_id, holder)

    def _session() -> tuple[int, int, bool, bool]:
        from ..ne_netmiko import drain_read_channel

        if is_stop_requested(batch_id):
            raise TimeoutError(f"{label}_stopped")

        conn = open_netmiko_connection(creds, session_timeout=budget)
        holder["conn"] = conn
        total_rows = 0
        cmd_count = 0
        any_fail = False
        any_ok = False
        pending: list[SpooledCommand] = []
        task_id = ""
        from .parse_pool import (
            AuxRawCapture,
            PrimaryParseJob,
            get_parse_pool,
            parse_async_enabled,
        )
        from .persist_pool import get_persist_pool

        persist = get_persist_pool()
        parse_pool = get_parse_pool() if parse_async_enabled() else None
        parse_stats_lock = threading.Lock()
        parse_stats = {"ok": False, "fail": False, "pending": 0}

        def _on_parse_done(ok: bool, fail: bool) -> None:
            with parse_stats_lock:
                if ok:
                    parse_stats["ok"] = True
                if fail:
                    parse_stats["fail"] = True
                parse_stats["pending"] = max(0, int(parse_stats["pending"]) - 1)

        def _submit_pending() -> None:
            nonlocal pending
            if not pending:
                return
            chunk = list(pending)
            pending = []
            persist.submit(batch_id, chunk)

        def _queue(item: SpooledCommand, *, records: Any = None) -> None:
            nonlocal cmd_count
            if records is not None and item.persist_kind:
                rel, n = write_records(batch_id, item.id, records)
                item.records_rel_path = rel
                item.row_count = int(n or 0)
            try:
                write_meta(batch_id, item.id, item.to_meta())
            except Exception:
                _log.exception("biz_state write meta failed cmd=%s", item.id)
            pending.append(item)
            cmd_count += 1
            if len(pending) >= flush_every:
                _submit_pending()

        try:
            try:
                disable_target_paging(
                    conn,
                    vendor=str(creds.get("vendor") or vendor_eff or ""),
                    device_type=str(creds.get("device_type") or device_type_eff or ""),
                )
            except Exception:
                pass
            try:
                drain_read_channel(conn)
            except Exception:
                pass

            # Lightweight lookup for task_id / batch existence (no long hold).
            sdb = SessionLocal()
            try:
                batch_row = sdb.get(BizStateBatch, batch_id)
                if not batch_row:
                    return 0, 0, True, False
                task_id = str(batch_row.task_id or "")
            finally:
                sdb.close()

            session = CollectSession(
                conn,
                vendor=vendor_eff,
                device_type=device_type_eff,
                vendor_key=vendor_key,
                read_timeout=per_cmd,
                shared_cache=shared_cache,
                cache_lock=cache_lock,
                cmd_locks=cmd_locks,
            )

            # Resolve expand_all → concrete per-VRF commands via discover profile.
            flat_work: list[WorkItem] = []
            persisted = aux_persisted if aux_persisted is not None else set()
            for concrete, params, profile_id, item_id, mode in work:
                if mode != "expand_all":
                    flat_work.append((concrete, params, profile_id, item_id, mode))
                    continue
                profile = get_profile(profile_id)
                if profile is None or not profile.placeholders:
                    any_fail = True
                    _emit_task_event(
                        task_id=task_id,
                        message=f"expand_all missing profile {profile_id}",
                        level="error",
                    )
                    continue
                ph = profile.placeholders[0]
                disc = get_profile(str(ph.discover_profile_id or "").strip())
                if disc is None:
                    any_fail = True
                    _emit_task_event(
                        task_id=task_id,
                        message=f"expand_all discover profile missing for {profile_id}",
                        level="error",
                    )
                    continue
                disc_cmd = normalize_command(disc.command_template)
                entry, _ = session.fetch_and_parse(
                    disc_cmd,
                    parser_id=disc.parser_id,
                    textfsm_command=disc.textfsm_command or disc_cmd,
                    params={},
                )
                if not entry.ok:
                    any_fail = True
                    _emit_task_event(
                        task_id=task_id,
                        message=f"expand_all discover failed: {entry.error}",
                        level="error",
                    )
                    continue
                try:
                    pairs = expand_bindings_from_discover_records(
                        profile=profile,
                        records=entry.records,
                    )
                except ValueError as exc:
                    any_fail = True
                    _emit_task_event(task_id=task_id, message=str(exc), level="error")
                    continue
                for cmd, p in pairs:
                    flat_work.append((cmd, p, profile_id, item_id, "normal"))

            for concrete, params, profile_id, item_id, mode in flat_work:
                if holder.get("timed_out") or holder.get("stop_requested") or is_stop_requested(
                    batch_id
                ):
                    raise TimeoutError(f"{label}_stopped")

                cmd_id = uuid4().hex
                raw_text = ""
                cache_hit_primary = False
                try:
                    cached = session.get_cached(concrete)
                    if cached is not None and str(cached.raw or "").strip():
                        raw_text = str(cached.raw or "")
                        cache_hit_primary = True
                    else:
                        raw_text = str(
                            send_show_command(conn, concrete, read_timeout=per_cmd) or ""
                        )
                except Exception as exc:
                    any_fail = True
                    sp = SpooledCommand(
                        id=cmd_id,
                        batch_id=batch_id,
                        task_item_id=item_id,
                        profile_id=profile_id,
                        raw_command=concrete[:512],
                        params_json=dict(params or {}),
                        parse_status="failed",
                        message=_format_error(exc),
                    )
                    _queue(sp)
                    continue

                raw_rel = ""
                raw_lines = 0
                try:
                    from .spool import count_text_lines, write_raw_text

                    raw_rel = write_raw_text(batch_id, cmd_id, raw_text)
                    raw_lines = count_text_lines(raw_text)
                except Exception:
                    _log.exception("biz_state spool raw failed cmd=%s", cmd_id)

                if mode == "custom":
                    any_ok = True
                    _queue(
                        SpooledCommand(
                            id=cmd_id,
                            batch_id=batch_id,
                            task_item_id=item_id,
                            profile_id=profile_id,
                            raw_command=concrete[:512],
                            params_json=dict(params or {}),
                            parse_status="skipped_custom",
                            message="custom_raw",
                            raw_rel_path=raw_rel,
                            raw_line_count=raw_lines,
                        )
                    )
                    continue

                hit = match_command(vendor_key=vendor_key, command=concrete)
                if not hit:
                    any_fail = True
                    _queue(
                        SpooledCommand(
                            id=cmd_id,
                            batch_id=batch_id,
                            task_item_id=item_id,
                            profile_id=profile_id,
                            raw_command=concrete[:512],
                            params_json=dict(params or {}),
                            parse_status="unmatched",
                            message="no profile matched concrete command",
                            raw_rel_path=raw_rel,
                            raw_line_count=raw_lines,
                        )
                    )
                    continue

                merged = {**params, **hit.params}
                if not get_parser(hit.profile.parser_id):
                    any_fail = True
                    _queue(
                        SpooledCommand(
                            id=cmd_id,
                            batch_id=batch_id,
                            task_item_id=item_id,
                            profile_id=hit.profile.profile_id,
                            parser_id=hit.profile.parser_id,
                            metric_id=hit.profile.metric_id,
                            raw_command=concrete[:512],
                            params_json=merged,
                            parse_status="failed",
                            message=f"unknown parser {hit.profile.parser_id}",
                            raw_rel_path=raw_rel,
                            raw_line_count=raw_lines,
                        )
                    )
                    continue

                if not cache_hit_primary:
                    session.remember(
                        concrete,
                        raw=raw_text,
                        ok=True,
                        cmd_row_id=cmd_id,
                    )

                # Collect aux raws on the SSH thread (no TextFSM); parse overlaps next CLI.
                aux_captures: list[AuxRawCapture] = []
                for aux in list(hit.profile.aux_commands or []):
                    try:
                        ra = resolve_aux_command(aux, params=merged)
                    except ValueError as exc:
                        _queue(
                            SpooledCommand(
                                id=uuid4().hex,
                                batch_id=batch_id,
                                task_item_id=item_id,
                                profile_id=str(aux.profile_id or "")[:128],
                                raw_command=str(aux.key or "")[:512],
                                parse_status="aux_failed",
                                message=f"aux_for={cmd_id};resolve:{exc}"[:1020],
                            )
                        )
                        continue
                    aux_id = uuid4().hex
                    entry, cache_hit = session.fetch_raw(ra.command, cmd_row_id=aux_id)
                    aux_mid = str(getattr(ra.profile, "metric_id", "") or "").strip()
                    aux_rel = ""
                    if entry.raw:
                        try:
                            aux_rel = write_raw_text(batch_id, aux_id, entry.raw or "")
                        except Exception:
                            _log.exception(
                                "biz_state spool aux raw failed cmd=%s", aux_id
                            )
                    if not entry.ok:
                        any_fail = True
                        _queue(
                            SpooledCommand(
                                id=aux_id,
                                batch_id=batch_id,
                                task_item_id=item_id,
                                profile_id=ra.profile_id,
                                parser_id=ra.parser_id,
                                metric_id=aux_mid,
                                raw_command=ra.command[:512],
                                parse_status="aux_failed",
                                message=f"aux_for={cmd_id};{entry.error}"[:1020],
                                raw_rel_path=aux_rel,
                            )
                        )
                        continue
                    # If cache already has parsed records, pass them through.
                    aux_captures.append(
                        AuxRawCapture(
                            key=ra.key,
                            aux_id=aux_id,
                            profile_id=ra.profile_id,
                            parser_id=ra.parser_id,
                            metric_id=aux_mid,
                            command=ra.command,
                            textfsm_command=ra.textfsm_command,
                            rule_keys=tuple(ra.rule_keys or ()),
                            raw=entry.raw,
                            raw_rel_path=aux_rel,
                            cache_hit=bool(cache_hit and entry.records),
                            records=list(entry.records or []),
                            fsm_tables=dict(entry.fsm_tables or {}),
                            ok=True,
                        )
                    )

                job = PrimaryParseJob(
                    batch_id=batch_id,
                    cmd_id=cmd_id,
                    task_item_id=item_id,
                    profile_id=hit.profile.profile_id,
                    parser_id=hit.profile.parser_id,
                    metric_id=hit.profile.metric_id,
                    concrete=concrete,
                    merged_params=dict(merged or {}),
                    raw_text=raw_text,
                    raw_rel_path=raw_rel,
                    raw_line_count=raw_lines,
                    textfsm_command=hit.profile.textfsm_command or concrete,
                    vendor=vendor_eff,
                    device_type=device_type_eff,
                    enrich_joins=list(hit.profile.enrich_joins or []),
                    aux_captures=aux_captures,
                    persisted=persisted,
                    cache_lock=cache_lock,
                    on_done=_on_parse_done if parse_pool is not None else None,
                )
                if parse_pool is not None:
                    with parse_stats_lock:
                        parse_stats["pending"] += 1
                    parse_pool.submit(job)
                    # Count primary (+ aux will be counted in parse worker via persist).
                    # cmd_count: bump for primary + each aux capture so progress is visible.
                    cmd_count += 1 + len(aux_captures)
                else:
                    ok, fail = _run_primary_parse_job(job)
                    if ok:
                        any_ok = True
                    if fail:
                        any_fail = True
                    # Sync path: parse job already submitted persist; count cmds.
                    cmd_count += 1 + len(aux_captures)

            # Drain CLI spool leftovers, then wait parse + persist pools.
            _submit_pending()
            if parse_pool is not None:
                if not parse_pool.wait_idle(timeout=max(30.0, float(budget))):
                    _log.warning(
                        "biz_state parse barrier timed out batch=%s lane=%s",
                        batch_id,
                        label,
                    )
                    any_fail = True
                    _emit_task_event(
                        task_id=task_id,
                        message=f"{label}: parse_barrier_timeout",
                        level="error",
                    )
                with parse_stats_lock:
                    if parse_stats["ok"]:
                        any_ok = True
                    if parse_stats["fail"]:
                        any_fail = True
            if not persist.wait_idle(timeout=max(30.0, float(budget))):
                _log.warning(
                    "biz_state persist barrier timed out batch=%s lane=%s",
                    batch_id,
                    label,
                )
                any_fail = True
                _emit_task_event(
                    task_id=task_id,
                    message=f"{label}: persist_barrier_timeout",
                    level="error",
                )
            # Do NOT read batch.row_count here — dual light+heavy lanes would each
            # see the cumulative DB total and _absorb would double-count.
            return total_rows, cmd_count, any_fail, any_ok
        finally:
            # Best-effort: enqueue leftover spool before connection teardown.
            try:
                _submit_pending()
                if parse_pool is not None:
                    parse_pool.wait_idle(timeout=60.0)
                persist.wait_idle(timeout=60.0)
            except Exception:
                _log.exception(
                    "biz_state persist drain on lane exit failed batch=%s", batch_id
                )
            holder.pop("conn", None)
            close_netmiko_connection(conn)

    try:
        return run_cli_with_timeout(
            _session,
            timeout_sec=budget,
            conn_holder=holder,
            label=label,
            acquire_budget=True,
        )
    except TimeoutError as exc:
        raise RuntimeError(str(exc)[:1020]) from exc
    finally:
        unregister_lane_holder(batch_id, holder)


def _absorb_lane_result(
    result: tuple[int, int, bool, bool] | BaseException,
    *,
    total_rows: int,
    cmd_count: int,
    any_fail: bool,
    any_ok: bool,
    lane_errors: list[str],
) -> tuple[int, int, bool, bool]:
    if isinstance(result, BaseException):
        lane_errors.append(_format_error(result))
        return total_rows, cmd_count, True, any_ok
    rows, cmds, fail, ok = result
    return (
        total_rows + int(rows or 0),
        cmd_count + int(cmds or 0),
        any_fail or bool(fail),
        any_ok or bool(ok),
    )


def _bump_batch_progress(batch_id: str, *, add_cmds: int = 0, add_rows: int = 0) -> None:
    """Atomically bump batch counters so UI can show progress while lanes still run."""
    cmds = int(add_cmds or 0)
    rows = int(add_rows or 0)
    if not batch_id or (cmds <= 0 and rows <= 0):
        return
    db = SessionLocal()
    try:
        batch = (
            db.query(BizStateBatch)
            .filter(BizStateBatch.id == batch_id)
            .with_for_update()
            .one_or_none()
        )
        if not batch:
            return
        if cmds > 0:
            batch.command_count = int(batch.command_count or 0) + cmds
        if rows > 0:
            batch.row_count = int(batch.row_count or 0) + rows
        db.commit()
    except Exception:
        _log.exception("biz_state bump batch progress failed batch=%s", batch_id)
        try:
            db.rollback()
        except Exception:
            pass
    finally:
        db.close()


def _is_stale_db_connection(exc: BaseException) -> bool:
    """True when PG/middleware closed an idle connection mid-collect."""
    if isinstance(exc, OperationalError):
        return True
    if isinstance(exc, DBAPIError) and bool(getattr(exc, "connection_invalidated", False)):
        return True
    msg = str(exc or "").lower()
    return (
        "server closed the connection" in msg
        or "connection not open" in msg
        or "connection already closed" in msg
        or "ssl connection has been closed" in msg
    )


def _invalidate_session(db) -> None:
    try:
        conn = db.connection()
        conn.invalidate()
    except Exception:
        pass
    try:
        db.close()
    except Exception:
        pass


def _run_db_with_reconnect(fn, *, label: str = "biz_state_db"):
    """Run ``fn(db)`` on a fresh Session; retry once after disconnect/OperationalError."""
    last: BaseException | None = None
    for attempt in range(2):
        db = SessionLocal()
        try:
            result = fn(db)
            try:
                db.close()
            except Exception:
                pass
            return result
        except Exception as exc:
            last = exc
            try:
                db.rollback()
            except Exception:
                pass
            _invalidate_session(db)
            if attempt == 0 and _is_stale_db_connection(exc):
                _log.warning("%s reconnect after stale connection: %s", label, exc)
                continue
            raise
    assert last is not None
    raise last


def _batch_has_progress(batch_id: str) -> bool:
    """Whether progressive bumps already persisted cmds/rows for this batch."""

    def _read(db) -> bool:
        live = db.get(BizStateBatch, batch_id)
        if not live:
            return False
        return int(live.command_count or 0) > 0 or int(live.row_count or 0) > 0

    return bool(_run_db_with_reconnect(_read, label="biz_state_batch_progress"))


def _is_fail_parse_status(status: str | None) -> bool:
    st = str(status or "").strip().lower()
    return st in ("failed", "error", "fail", "aux_failed") or st.endswith("_failed")


def _is_skip_parse_status(status: str | None) -> bool:
    st = str(status or "").strip().lower()
    return st.startswith("skipped")


def _batch_issue_summary(
    db,
    batch_id: str,
    lane_errors: list[str] | None = None,
) -> str:
    """Human-readable reason for partial/failed batches (never leave message empty)."""
    parts: list[str] = []
    for err in lane_errors or []:
        e = str(err or "").strip()
        if e and e not in parts:
            parts.append(e)

    cmds = (
        db.query(BizStateBatchCommand)
        .filter(BizStateBatchCommand.batch_id == batch_id)
        .order_by(BizStateBatchCommand.created_at.asc())
        .all()
    )
    n_fail = 0
    n_skip = 0
    n_aux_fail = 0
    samples: list[str] = []
    for c in cmds:
        st = str(c.parse_status or "").strip().lower()
        msg = str(c.message or "").strip()
        if st == "aux_failed" or (st.startswith("aux") and "fail" in st):
            n_aux_fail += 1
            if msg and len(samples) < 3:
                samples.append(f"aux:{c.raw_command}: {msg}"[:160])
        elif _is_fail_parse_status(st):
            n_fail += 1
            if msg and len(samples) < 3:
                samples.append(f"{c.raw_command}: {msg}"[:160])
        elif _is_skip_parse_status(st):
            n_skip += 1
            if msg and len(samples) < 3:
                samples.append(f"skip:{c.profile_id or c.raw_command}: {msg}"[:160])

    if n_fail:
        parts.append(f"{n_fail} command(s) failed")
    if n_aux_fail:
        parts.append(f"{n_aux_fail} aux command(s) failed")
    if n_skip:
        parts.append(f"{n_skip} item(s) skipped (e.g. missing bindings)")
    for s in samples:
        if s not in parts:
            parts.append(s)

    text = "; ".join(parts).strip()
    if not text:
        text = "partial success (some steps failed or were skipped)"
    return text[:1020]


def _finalize_batch_status(
    *,
    batch_id: str,
    task_id: str,
    cmd_count: int,
    total_rows: int,
    any_fail: bool,
    any_ok: bool,
    lane_errors: list[str],
    stopped: bool = False,
) -> str:
    """Write terminal batch status on a fresh Session (retry once on disconnect)."""

    def _write(db) -> str:
        batch = db.get(BizStateBatch, batch_id)
        if not batch:
            return ""
        # Prefer progressive counters (survive lane timeout) over in-memory lane totals.
        batch.command_count = max(int(batch.command_count or 0), int(cmd_count or 0))
        batch.row_count = max(int(batch.row_count or 0), int(total_rows or 0))
        batch.ended_at = _utcnow()
        if stopped:
            if any_ok or int(batch.command_count or 0) > 0 or int(batch.row_count or 0) > 0:
                batch.status = "partial"
                batch.message = STOP_USER_MESSAGE
            else:
                batch.status = "cancelled"
                batch.message = STOP_USER_MESSAGE
        elif any_fail and any_ok:
            batch.status = "partial"
            batch.message = _batch_issue_summary(db, batch_id, lane_errors)
        elif any_fail and not any_ok:
            batch.status = "failed"
            batch.message = _batch_issue_summary(db, batch_id, lane_errors) or (
                "; ".join(lane_errors)[:1020] if lane_errors else "all commands failed"
            )
        else:
            # Still surface skipped-only rows as partial when some cmds ran ok.
            statuses = [
                str(c.parse_status or "")
                for c in db.query(BizStateBatchCommand)
                .filter(BizStateBatchCommand.batch_id == batch_id)
                .all()
            ]
            n_skip = sum(1 for st in statuses if _is_skip_parse_status(st))
            if n_skip > 0 and any_ok:
                batch.status = "partial"
                batch.message = _batch_issue_summary(db, batch_id, lane_errors)
            else:
                batch.status = "success"
                batch.message = ""
        status = str(batch.status or "")
        db.commit()
        # Only full success triggers auto compare; never block the collect thread.
        if status == "success" and not stopped:
            try:
                from .compare_service import schedule_auto_compare_for_task

                schedule_auto_compare_for_task(task_id, batch_id)
            except Exception:
                _log.exception("biz_state auto compare schedule failed task=%s", task_id)
        return status

    try:
        return str(
            _run_db_with_reconnect(_write, label="biz_state_finalize") or ""
        )
    finally:
        clear_stop_requested(batch_id)


def _fail_batch_status(batch_id: str, error: str) -> None:
    """Mark batch failed on a fresh Session (retry once on disconnect)."""
    msg = str(error or "")[:1020]
    stopped = "_stopped" in msg or is_stop_requested(batch_id)

    def _write(db) -> None:
        batch = db.get(BizStateBatch, batch_id)
        if not batch:
            return
        st = str(batch.status or "")
        if st not in ("running", "queued"):
            return
        if stopped:
            has_progress = int(batch.command_count or 0) > 0 or int(batch.row_count or 0) > 0
            batch.status = "partial" if has_progress else "cancelled"
            batch.message = STOP_USER_MESSAGE
        else:
            batch.status = "failed"
            batch.message = msg
        batch.ended_at = _utcnow()
        db.commit()

    try:
        _run_db_with_reconnect(_write, label="biz_state_fail_batch")
    finally:
        clear_stop_requested(batch_id)

def _run_collect_session(
    *,
    task_id: str,
    batch_id: str,
    source: str,
    ne_id: str,
    vendor: str,
    device_type: str,
) -> None:
    light_per = int(settings.ne_collect_read_timeout_sec or 120)
    light_cap = int(settings.ne_collect_run_timeout_cap_sec or 600)
    heavy_per = int(getattr(settings, "biz_state_heavy_read_timeout_sec", 300) or 300)
    heavy_cap = int(getattr(settings, "biz_state_heavy_run_timeout_cap_sec", 900) or 900)

    # Phase 1: resolve target + build work list, then release the DB connection.
    # Holding one Session across heavy CLI (up to ~2400s) lets PG/middleware close
    # the idle connection; finalize would then hit OperationalError.
    creds: dict[str, Any]
    vendor_eff: str
    device_type_eff: str
    vendor_key: str
    light_work: list[WorkItem]
    heavy_work: list[WorkItem]

    db = SessionLocal()
    try:
        task = db.get(BizStateTask, task_id)
        batch = db.get(BizStateBatch, batch_id)
        if not task or not batch:
            raise RuntimeError(f"batch_or_task_missing batch={batch_id} task={task_id}")

        try:
            if source == "managed":
                creds, info = resolve_cli_target(db, managed_ne_id=ne_id)
            elif source == "ume":
                creds, info = resolve_cli_target(db, ume_ne_id=ne_id)
            else:
                raise RuntimeError("invalid_source")
        except HTTPException as exc:
            raise RuntimeError(str(exc.detail or "resolve_failed")) from exc

        skip = cli_creds_skip_reason(creds, interactive=False)
        if skip:
            raise RuntimeError(skip)

        vendor_eff = str(info.get("vendor") or vendor or "")
        device_type_eff = str(info.get("device_type") or device_type or "")
        if vendor_eff and vendor_eff != task.vendor:
            task.vendor = vendor_eff
        if device_type_eff and device_type_eff != task.device_type:
            task.device_type = device_type_eff
        db.commit()

        vendor_key = resolve_vendor_key(vendor_eff, device_type_eff)

        items = (
            db.query(BizStateTaskItem)
            .filter(
                BizStateTaskItem.task_id == task_id,
                BizStateTaskItem.enabled.is_(True),
            )
            .order_by(BizStateTaskItem.sort_order.asc())
            .all()
        )

        # Build work list before opening session
        work: list[WorkItem] = []
        # Dedupe same CLI → same metric (e.g. legacy if_intf + config_interface).
        seen_work: set[tuple[str, str]] = set()
        for item in items:
            if item.kind == "custom_raw":
                cmd = normalize_command(item.command_override)
                if cmd:
                    key = (cmd, "__custom__")
                    if key in seen_work:
                        continue
                    seen_work.add(key)
                    work.append((cmd, {}, "", item.id, "custom"))
                continue
            profile = _resolve_collect_profile(item.source_profile_id)
            if profile is None:
                _append_event(
                    db,
                    task_id=task_id,
                    message=f"skip profile {item.source_profile_id} (missing or disabled)",
                    level="info",
                )
                continue
            binds = _bindings_for_item(db, item.id)
            try:
                pairs = expand_from_bindings(
                    profile=profile,
                    bindings=binds,
                    command_override=item.command_override,
                )
            except ValueError as exc:
                msg = str(exc)
                _append_event(db, task_id=task_id, message=msg, level="error")
                # Persist a visible skip row so UI / partial status can explain it.
                db.add(
                    BizStateBatchCommand(
                        id=uuid4().hex,
                        batch_id=batch_id,
                        task_item_id=item.id,
                        profile_id=profile.profile_id,
                        parser_id=str(profile.parser_id or ""),
                        metric_id=str(profile.metric_id or ""),
                        raw_command=str(profile.command_template or "")[:512],
                        params_json={},
                        parse_status="skipped",
                        message=msg[:1020],
                        row_count=0,
                        raw_text="",
                    )
                )
                batch.command_count = int(batch.command_count or 0) + 1
                continue
            for concrete, params in pairs:
                if concrete == EXPAND_ALL_COMMAND:
                    # Defer VRF list expansion until CollectSession is open.
                    work.append(("", dict(params or {}), profile.profile_id, item.id, "expand_all"))
                    continue
                cmd = normalize_command(concrete)
                # Prefer match_command metric so remapped if_intf shares key with config_interface
                hit = match_command(vendor_key=vendor_key, command=cmd)
                mid = str((hit.profile.metric_id if hit else profile.metric_id) or "").strip()
                pid = str((hit.profile.profile_id if hit else profile.profile_id) or "").strip()
                key = (cmd, mid or pid)
                if key in seen_work:
                    continue
                seen_work.add(key)
                work.append((cmd, params, pid or profile.profile_id, item.id, "normal"))

        if not work:
            batch.status = "failed"
            batch.message = "no commands to run"
            batch.ended_at = _utcnow()
            db.commit()
            raise RuntimeError("no commands to run")

        light_work, heavy_work = partition_work(work)
        db.commit()
    finally:
        db.close()

    # Fresh spool dir for this batch (collect → disk, then flush to DB).
    try:
        from .spool import clear_batch_spool

        clear_batch_spool(batch_id)
    except Exception:
        _log.exception("biz_state clear spool failed batch=%s", batch_id)

    # Phase 2: CLI lanes — no outer Session held across long timeouts.
    shared_cache: dict[str, Any] = {}
    cache_lock = threading.RLock()
    cmd_locks: dict[str, Any] = {}
    aux_persisted: set[tuple[str, str]] = set()
    lane_kwargs = dict(
        batch_id=batch_id,
        creds=creds,
        vendor_eff=vendor_eff,
        device_type_eff=device_type_eff,
        vendor_key=vendor_key,
        shared_cache=shared_cache,
        cache_lock=cache_lock,
        cmd_locks=cmd_locks,
        aux_persisted=aux_persisted,
    )

    def _run_light() -> tuple[int, int, bool, bool]:
        return _run_collect_lane(
            work=light_work,
            per_cmd=light_per,
            cap=light_cap,
            label="biz_state_light",
            **lane_kwargs,
        )

    def _run_heavy() -> tuple[int, int, bool, bool]:
        return _run_collect_lane(
            work=heavy_work,
            per_cmd=heavy_per,
            cap=heavy_cap,
            label="biz_state_heavy",
            **lane_kwargs,
        )

    total_rows = 0
    cmd_count = 0
    any_fail = False
    any_ok = False
    lane_errors: list[str] = []

    if light_work and heavy_work:
        heavy_fut = _heavy_cli_pool().submit(_run_heavy)
        light_res: tuple[int, int, bool, bool] | BaseException
        try:
            light_res = _run_light()
        except Exception as exc:
            light_res = exc
        heavy_res: tuple[int, int, bool, bool] | BaseException
        try:
            heavy_res = heavy_fut.result()
        except Exception as exc:
            heavy_res = exc
        total_rows, cmd_count, any_fail, any_ok = _absorb_lane_result(
            light_res,
            total_rows=total_rows,
            cmd_count=cmd_count,
            any_fail=any_fail,
            any_ok=any_ok,
            lane_errors=lane_errors,
        )
        total_rows, cmd_count, any_fail, any_ok = _absorb_lane_result(
            heavy_res,
            total_rows=total_rows,
            cmd_count=cmd_count,
            any_fail=any_fail,
            any_ok=any_ok,
            lane_errors=lane_errors,
        )
    elif heavy_work:
        try:
            total_rows, cmd_count, any_fail, any_ok = _run_heavy()
        except Exception as exc:
            total_rows, cmd_count, any_fail, any_ok = _absorb_lane_result(
                exc,
                total_rows=0,
                cmd_count=0,
                any_fail=False,
                any_ok=False,
                lane_errors=lane_errors,
            )
    else:
        try:
            total_rows, cmd_count, any_fail, any_ok = _run_light()
        except Exception as exc:
            total_rows, cmd_count, any_fail, any_ok = _absorb_lane_result(
                exc,
                total_rows=0,
                cmd_count=0,
                any_fail=False,
                any_ok=False,
                lane_errors=lane_errors,
            )

    if lane_errors and not any_ok and cmd_count == 0:
        # Progressive bumps may already have cmds; only hard-fail if nothing landed.
        if not _batch_has_progress(batch_id):
            if is_stop_requested(batch_id) or any("_stopped" in e for e in lane_errors):
                _finalize_batch_status(
                    batch_id=batch_id,
                    task_id=task_id,
                    cmd_count=cmd_count,
                    total_rows=total_rows,
                    any_fail=False,
                    any_ok=False,
                    lane_errors=lane_errors,
                    stopped=True,
                )
                return
            raise RuntimeError("; ".join(lane_errors)[:1020])
        any_fail = True

    # Ensure parse + persist pools drained before terminal status write.
    try:
        from .parse_pool import get_parse_pool, parse_async_enabled
        from .persist_pool import get_persist_pool

        if parse_async_enabled():
            if not get_parse_pool().wait_idle(timeout=120.0):
                _log.warning(
                    "biz_state parse barrier before finalize timed out batch=%s",
                    batch_id,
                )
                any_fail = True
                if "parse_barrier_timeout" not in lane_errors:
                    lane_errors.append("RuntimeError: parse_barrier_timeout")
        if not get_persist_pool().wait_idle(timeout=120.0):
            _log.warning(
                "biz_state persist barrier before finalize timed out batch=%s",
                batch_id,
            )
            any_fail = True
            if "persist_barrier_timeout" not in lane_errors:
                lane_errors.append("RuntimeError: persist_barrier_timeout")
    except Exception:
        _log.exception("biz_state parse/persist barrier before finalize failed batch=%s", batch_id)

    stopped = is_stop_requested(batch_id) or any("_stopped" in e for e in lane_errors)
    _finalize_batch_status(
        batch_id=batch_id,
        task_id=task_id,
        cmd_count=cmd_count,
        total_rows=total_rows,
        any_fail=any_fail,
        any_ok=any_ok,
        lane_errors=lane_errors,
        stopped=stopped,
    )

    def _purge(db) -> None:
        _purge_task_retention(db, task_id=task_id)

    try:
        _run_db_with_reconnect(_purge, label="biz_state_purge")
    except Exception:
        _log.exception("biz_state retention purge wrapper failed task=%s", task_id)


def _purge_task_retention(db, *, task_id: str) -> None:
    from .retention import purge_task_batches

    task = db.get(BizStateTask, task_id)
    if not task:
        return
    try:
        info = purge_task_batches(db, task)
        if info.get("dropped"):
            _log.info(
                "biz_state retention purged task=%s dropped=%s days=%s daily=%s",
                task_id,
                info.get("dropped"),
                info.get("retention_days"),
                info.get("daily_keep_enabled"),
            )
    except Exception:
        _log.exception("biz_state retention purge failed task=%s", task_id)


def trigger_collect_now(task_id: str) -> dict[str, Any]:
    db = SessionLocal()
    try:
        task = db.get(BizStateTask, task_id)
        if not task:
            raise HTTPException(status_code=404, detail="task not found")
        if task.collect_running:
            raise HTTPException(status_code=409, detail="collect already running")
    finally:
        db.close()
    # Manual: allow even when schedule is on (status=running) or paused/stopped.
    dispatch_collect(task_id, manual=True)
    return {"ok": True, "task_id": task_id}