netx/netx_api/biz_state/collect_session.py

"""Per-batch CLI session helpers: aux resolution, command cache, parse bundle."""

from __future__ import annotations

from contextlib import nullcontext
from dataclasses import dataclass, field
from typing import Any, Callable
import re
import threading

from .command_match import normalize_command
from .enrich import EnrichJoin, apply_enrich_joins
from .parsers import get_parser, get_parser_meta, run_parser
from .profiles import AuxCommand, ParseProfile, get_profile

SendFn = Callable[..., str]


@dataclass
class CachedCommand:
    raw: str = ""
    fsm_tables: dict[str, list[dict[str, Any]]] = field(default_factory=dict)
    records: list[dict[str, Any]] = field(default_factory=list)
    ok: bool = True
    error: str = ""
    cmd_row_id: str = ""


@dataclass(frozen=True)
class ResolvedAux:
    key: str
    profile_id: str
    command: str
    textfsm_command: str
    parser_id: str
    rule_keys: tuple[str, ...]
    profile: ParseProfile


def resolve_aux_command(
    aux: AuxCommand,
    params: dict[str, str] | None = None,
) -> ResolvedAux:
    """Resolve aux from ``profile_id``; render ``<placeholders>`` from primary params."""
    key = str(aux.key or "").strip()
    pid = str(aux.profile_id or "").strip()
    if not key or not pid:
        raise ValueError("AuxCommand requires key and profile_id")
    prof = get_profile(pid)
    if not prof:
        raise ValueError(f"aux profile not found: {pid}")
    tmpl = str(prof.command_template or "")
    rendered = tmpl
    for name, val in dict(params or {}).items():
        n = str(name or "").strip()
        if not n:
            continue
        rendered = rendered.replace(f"<{n}>", str(val or "").strip())
    cmd = normalize_command(rendered)
    if not cmd:
        raise ValueError(f"aux profile {pid} has empty command_template")
    if re.search(r"<[^>]+>", cmd):
        raise ValueError(f"aux profile {pid} unresolved placeholders: {cmd}")
    textfsm = str(prof.textfsm_command or cmd).strip()
    # Prefer concrete rendered command for FSM when aux had placeholders
    if params and any(f"<{k}>" in tmpl for k in (params or {})):
        textfsm = cmd
    parser_id = str(prof.parser_id or "").strip()
    meta = get_parser_meta(parser_id) if parser_id else None
    rule_keys = tuple((meta or {}).get("rule_keys") or ())
    return ResolvedAux(
        key=key,
        profile_id=pid,
        command=cmd,
        textfsm_command=textfsm,
        parser_id=parser_id,
        rule_keys=rule_keys,
        profile=prof,
    )


@dataclass
class ParseBundle:
    """Inputs for primary ``run_parser`` after primary + aux collection."""

    raws: dict[str, str] = field(default_factory=dict)
    command_rules: dict[str, list[str]] = field(default_factory=dict)
    aux_records: dict[str, list[dict[str, Any]]] = field(default_factory=dict)
    fsm_extra: dict[str, list[dict[str, Any]]] = field(default_factory=dict)


class CollectSession:
    """SSH session-scoped command cache + aux fetch/parse.

    Optional ``shared_cache`` / ``cache_lock`` / ``cmd_locks`` let light+heavy
    lanes reuse the same CLI results within one batch (config_vrf / FIB aux).
    """

    def __init__(
        self,
        conn: Any,
        *,
        vendor: str = "",
        device_type: str = "",
        vendor_key: str = "",
        read_timeout: int = 120,
        send_fn: SendFn | None = None,
        shared_cache: dict[str, CachedCommand] | None = None,
        cache_lock: threading.RLock | None = None,
        cmd_locks: dict[str, threading.Lock] | None = None,
    ) -> None:
        self.conn = conn
        self.vendor = vendor
        self.device_type = device_type
        self.vendor_key = vendor_key
        self.read_timeout = int(read_timeout or 120)
        self._send = send_fn
        self.cache: dict[str, CachedCommand] = (
            shared_cache if shared_cache is not None else {}
        )
        self._cache_lock = cache_lock
        self._cmd_locks = cmd_locks if cmd_locks is not None else {}

    def _meta_lock(self):
        return self._cache_lock if self._cache_lock is not None else nullcontext()

    def _command_lock(self, command: str):
        ck = normalize_command(command)
        if self._cache_lock is None:
            return nullcontext()
        with self._cache_lock:
            lock = self._cmd_locks.get(ck)
            if lock is None:
                lock = threading.Lock()
                self._cmd_locks[ck] = lock
        return lock

    def _send_show(self, command: str) -> str:
        if self._send is None:
            from ..ne_netmiko import send_show_command

            return str(send_show_command(self.conn, command, read_timeout=self.read_timeout) or "")
        return str(self._send(self.conn, command, read_timeout=self.read_timeout) or "")

    def remember(
        self,
        command: str,
        *,
        raw: str = "",
        fsm_tables: dict[str, list[dict[str, Any]]] | None = None,
        records: list[dict[str, Any]] | None = None,
        ok: bool = True,
        error: str = "",
        cmd_row_id: str = "",
    ) -> CachedCommand:
        ck = normalize_command(command)
        entry = CachedCommand(
            raw=str(raw or ""),
            fsm_tables=dict(fsm_tables or {}),
            records=list(records or []),
            ok=bool(ok),
            error=str(error or ""),
            cmd_row_id=str(cmd_row_id or ""),
        )
        with self._meta_lock():
            self.cache[ck] = entry
        return entry

    def get_cached(self, command: str) -> CachedCommand | None:
        ck = normalize_command(command)
        with self._meta_lock():
            hit = self.cache.get(ck)
            if hit and hit.ok:
                return hit
        return None

    def fetch_raw(
        self,
        command: str,
        *,
        cmd_row_id: str = "",
    ) -> tuple[CachedCommand, bool]:
        """CLI-only fetch (no TextFSM). Return ``(entry, cache_hit)``.

        Used so the SSH thread can overlap subsequent commands while a parse
        pool parses prior raws.
        """
        with self._command_lock(command):
            cached = self.get_cached(command)
            if cached is not None and str(cached.raw or "").strip():
                return cached, True
            try:
                raw = self._send_show(command)
            except Exception as exc:
                entry = self.remember(
                    command,
                    raw="",
                    ok=False,
                    error=f"{type(exc).__name__}: {exc}",
                    cmd_row_id=cmd_row_id,
                )
                return entry, False
            entry = self.remember(
                command,
                raw=raw,
                ok=True,
                cmd_row_id=cmd_row_id,
            )
            return entry, False

    def fetch_and_parse(
        self,
        command: str,
        *,
        parser_id: str = "",
        textfsm_command: str = "",
        params: dict[str, str] | None = None,
        cmd_row_id: str = "",
    ) -> tuple[CachedCommand, bool]:
        """Return ``(entry, cache_hit)``. On miss: CLI + optional parser.

        Same concrete CLI is serialized across shared-cache lanes so aux of
        one monitor item can be reused by the next without re-collecting.
        Prefer ``fetch_raw`` on the collect hot path when parse can be async.
        """
        with self._command_lock(command):
            cached = self.get_cached(command)
            if cached is not None:
                # Full parse hit (records or explicit prior parse failure with raw).
                if cached.records or (cached.ok and not parser_id):
                    return cached, True
                if cached.ok and str(cached.raw or "").strip() and parser_id and get_parser(parser_id):
                    # Raw-only cache (async collect): parse now without re-CLI.
                    try:
                        records, fsm_tables, _keys = run_parser(
                            parser_id,
                            raw_text=cached.raw,
                            vendor=self.vendor,
                            device_type=self.device_type,
                            command=textfsm_command or command,
                            textfsm_command=textfsm_command or "",
                            params=params or {},
                        )
                    except Exception as exc:
                        entry = self.remember(
                            command,
                            raw=cached.raw,
                            ok=False,
                            error=f"parse: {type(exc).__name__}: {exc}",
                            cmd_row_id=cmd_row_id or cached.cmd_row_id,
                        )
                        return entry, True
                    entry = self.remember(
                        command,
                        raw=cached.raw,
                        fsm_tables=fsm_tables,
                        records=records,
                        ok=True,
                        cmd_row_id=cmd_row_id or cached.cmd_row_id,
                    )
                    return entry, True
                if cached.ok:
                    return cached, True
            try:
                raw = self._send_show(command)
            except Exception as exc:
                entry = self.remember(
                    command,
                    raw="",
                    ok=False,
                    error=f"{type(exc).__name__}: {exc}",
                    cmd_row_id=cmd_row_id,
                )
                return entry, False
            records: list[dict[str, Any]] = []
            fsm_tables: dict[str, list[dict[str, Any]]] = {}
            if parser_id and get_parser(parser_id):
                try:
                    records, fsm_tables, _keys = run_parser(
                        parser_id,
                        raw_text=raw,
                        vendor=self.vendor,
                        device_type=self.device_type,
                        command=textfsm_command or command,
                        textfsm_command=textfsm_command or "",
                        params=params or {},
                    )
                except Exception as exc:
                    entry = self.remember(
                        command,
                        raw=raw,
                        ok=False,
                        error=f"parse: {type(exc).__name__}: {exc}",
                        cmd_row_id=cmd_row_id,
                    )
                    return entry, False
            entry = self.remember(
                command,
                raw=raw,
                fsm_tables=fsm_tables,
                records=records,
                ok=True,
                cmd_row_id=cmd_row_id,
            )
            return entry, False


def primary_rule_keys(parser_id: str) -> list[str]:
    meta = get_parser_meta(parser_id) or {}
    return list(meta.get("rule_keys") or ())


def build_parse_bundle(
    *,
    primary_raw: str,
    primary_parser_id: str,
    aux_results: dict[str, CachedCommand],
    resolved_aux: list[ResolvedAux],
) -> ParseBundle:
    bundle = ParseBundle(
        raws={"primary": str(primary_raw or "")},
        command_rules={"primary": primary_rule_keys(primary_parser_id)},
    )
    for ra in resolved_aux:
        entry = aux_results.get(ra.key) or CachedCommand(ok=False)
        bundle.raws[ra.key] = entry.raw
        bundle.command_rules[ra.key] = list(ra.rule_keys)
        if entry.records:
            bundle.aux_records[ra.key] = list(entry.records)
        bundle.fsm_extra.update(entry.fsm_tables or {})
    return bundle


def run_primary_with_bundle(
    parser_id: str,
    *,
    bundle: ParseBundle,
    vendor: str = "",
    device_type: str = "",
    command: str = "",
    textfsm_command: str = "",
    params: dict[str, str] | None = None,
    enrich_joins: list[EnrichJoin] | None = None,
) -> tuple[list[dict[str, Any]], dict[str, list[dict[str, Any]]], list[str]]:
    records, fsm_tables, keys = run_parser(
        parser_id,
        raw_text=bundle.raws.get("primary") or "",
        vendor=vendor,
        device_type=device_type,
        command=command,
        textfsm_command=textfsm_command,
        params=params,
        raws=bundle.raws,
        command_rules=bundle.command_rules,
        aux_records=bundle.aux_records,
        fsm_tables_extra=bundle.fsm_extra,
    )
    if enrich_joins:
        apply_enrich_joins(records, bundle.aux_records, enrich_joins)
    return records, fsm_tables, keys