Medios-Macina/SYS/pipeline_progress.py

from __future__ import annotations

import sys
from contextlib import contextmanager
from typing import Any, Iterator, Optional, Sequence, Tuple


class PipelineProgress:
    """Small adapter around PipelineLiveProgress.

    This centralizes the boilerplate used across cmdlets:
    - locating the active Live UI (if any)
    - resolving the current pipe_index from stage context
    - step-based progress (begin_pipe_steps/advance_pipe_step)
    - optional pipe percent/status updates
    - optional byte transfer bars
    - optional local Live panel when a cmdlet runs standalone

    The class is intentionally defensive: all UI operations are best-effort.
    """

    def __init__(self, pipeline_module: Any):
        self._ctx = pipeline_module
        self._local_ui: Optional[Any] = None
        self._local_attached: bool = False

    def ui_and_pipe_index(self) -> Tuple[Optional[Any], int]:
        ui = None
        try:
            ui = self._ctx.get_live_progress(
            ) if hasattr(self._ctx,
                         "get_live_progress") else None
        except Exception:
            ui = None

        pipe_idx: int = 0
        try:
            stage_ctx = (
                self._ctx.get_stage_context()
                if hasattr(self._ctx,
                           "get_stage_context") else None
            )
            maybe_idx = getattr(
                stage_ctx,
                "pipe_index",
                None
            ) if stage_ctx is not None else None
            if isinstance(maybe_idx, int):
                pipe_idx = int(maybe_idx)
        except Exception:
            pipe_idx = 0

        return ui, pipe_idx

    def begin_steps(self, total_steps: int) -> None:
        ui, pipe_idx = self.ui_and_pipe_index()
        if ui is None:
            return
        try:
            begin = getattr(ui, "begin_pipe_steps", None)
            if callable(begin):
                begin(int(pipe_idx), total_steps=int(total_steps))
        except Exception:
            return

    def step(self, text: str) -> None:
        ui, pipe_idx = self.ui_and_pipe_index()
        if ui is None:
            return
        try:
            adv = getattr(ui, "advance_pipe_step", None)
            if callable(adv):
                adv(int(pipe_idx), str(text))
        except Exception:
            return

    def set_percent(self, percent: int) -> None:
        ui, pipe_idx = self.ui_and_pipe_index()
        if ui is None:
            return
        try:
            set_pct = getattr(ui, "set_pipe_percent", None)
            if callable(set_pct):
                set_pct(int(pipe_idx), int(percent))
        except Exception:
            return

    def set_status(self, text: str) -> None:
        ui, pipe_idx = self.ui_and_pipe_index()
        if ui is None:
            return
        try:
            setter = getattr(ui, "set_pipe_status_text", None)
            if callable(setter):
                setter(int(pipe_idx), str(text))
        except Exception:
            return

    def clear_status(self) -> None:
        ui, pipe_idx = self.ui_and_pipe_index()
        if ui is None:
            return
        try:
            clr = getattr(ui, "clear_pipe_status_text", None)
            if callable(clr):
                clr(int(pipe_idx))
        except Exception:
            return

    def begin_transfer(self, *, label: str, total: Optional[int] = None) -> None:
        ui, _ = self.ui_and_pipe_index()
        if ui is None:
            return
        try:
            fn = getattr(ui, "begin_transfer", None)
            if callable(fn):
                fn(label=str(label or "transfer"), total=total)
        except Exception:
            return

    def update_transfer(
        self,
        *,
        label: str,
        completed: Optional[int],
        total: Optional[int] = None
    ) -> None:
        ui, _ = self.ui_and_pipe_index()
        if ui is None:
            return
        try:
            fn = getattr(ui, "update_transfer", None)
            if callable(fn):
                fn(label=str(label or "transfer"), completed=completed, total=total)
        except Exception:
            return

    def finish_transfer(self, *, label: str) -> None:
        ui, _ = self.ui_and_pipe_index()
        if ui is None:
            return
        try:
            fn = getattr(ui, "finish_transfer", None)
            if callable(fn):
                fn(label=str(label or "transfer"))
        except Exception:
            return

    def on_emit(self, emitted: Any) -> None:
        """Advance local pipe progress after pipeline_context.emit().

        The shared PipelineExecutor wires on_emit automatically for pipelines.
        Standalone cmdlet runs do not, so cmdlets call this explicitly.
        """

        if self._local_ui is None:
            return
        try:
            self._local_ui.on_emit(0, emitted)
        except Exception:
            return

    def ensure_local_ui(
        self,
        *,
        label: str,
        total_items: int,
        items_preview: Optional[Sequence[Any]] = None
    ) -> bool:
        """Start a local PipelineLiveProgress panel if no shared UI exists."""

        try:
            existing = (
                self._ctx.get_live_progress()
                if hasattr(self._ctx,
                           "get_live_progress") else None
            )
        except Exception:
            existing = None

        if existing is not None:
            return False
        if not bool(getattr(sys.stderr, "isatty", lambda: False)()):
            return False

        try:
            from SYS.models import PipelineLiveProgress

            ui = PipelineLiveProgress([str(label or "pipeline")], enabled=True)
            ui.start()
            try:
                if hasattr(self._ctx, "set_live_progress"):
                    self._ctx.set_live_progress(ui)
                    self._local_attached = True
            except Exception:
                self._local_attached = False

            try:
                ui.begin_pipe(
                    0,
                    total_items=max(1,
                                    int(total_items)),
                    items_preview=list(items_preview or [])
                )
            except Exception:
                pass

            self._local_ui = ui
            return True
        except Exception:
            self._local_ui = None
            self._local_attached = False
            return False

    def close_local_ui(self, *, force_complete: bool = True) -> None:
        if self._local_ui is None:
            return
        try:
            try:
                self._local_ui.finish_pipe(0, force_complete=bool(force_complete))
            except Exception:
                pass
            try:
                self._local_ui.stop()
            except Exception:
                pass
        finally:
            self._local_ui = None
            try:
                if self._local_attached and hasattr(self._ctx, "set_live_progress"):
                    self._ctx.set_live_progress(None)
            except Exception:
                pass
            self._local_attached = False

    @contextmanager
    def local_ui_if_needed(
        self,
        *,
        label: str,
        total_items: int,
        items_preview: Optional[Sequence[Any]] = None,
    ) -> Iterator["PipelineProgress"]:
        created = self.ensure_local_ui(
            label=label,
            total_items=total_items,
            items_preview=items_preview
        )
        try:
            yield self
        finally:
            if created:
                self.close_local_ui(force_complete=True)