#!/usr/bin/env python3
"""A full-screen task runner for long setup steps that stay in the TUI.

Long backend-setup steps (git clone, audiocpp_server build, model downloads,
pip installs, whisper transcription) used to run under ``tui.suspend``, which
dumped the user into plain console output. This widget keeps them inside the
hub's curses session: a worker thread runs an ordered list of ``TaskStep``s
while the main thread redraws a DOS-style frame showing each step's state
(pending / running with a spinner and elapsed clock / [OK] / [FAIL]), an
optional progress bar for the current step, and a dim scrolling log tail of
the step's output.

Steps stream their output by calling ``emit(line)`` (or simply printing to
stdout/stderr, which the view captures). The view turns output into progress
three ways, best-effort:

  * ``AUDIOCPP_PROGRESS downloaded=N total=M``  (audio.cpp model downloads,
    hidden from the log) — an exact bytes bar;
  * ``NN%`` (git ``Receiving objects: 45%``, cmake/make ``[ 45%]``, tqdm) —
    a percent bar;
  * ``[done/total]`` (ninja build output) — a count bar.

A ``threading.Event`` passed to every step is set when the user confirms
cancel (Esc/q); subprocess runners kill their child process group, and
in-process steps are expected to check it between units of work. When all
steps finish (or are cancelled) the view shows a summary and waits for a key
press, so a failure is never scrolled away. ``run_steps`` returns the first
non-zero step exit code (0 when every step succeeded).
"""

import contextlib
import re
import threading
import time
from dataclasses import dataclass
from queue import Empty, Queue
from typing import Callable, List, Optional, Tuple

from ui import tui

# Redraw cadence for the timed getch (milliseconds).
_DRAW_TIMEOUT_MS = 250

# How many recent output lines the log tail keeps.
_LOG_TAIL = 10

# Terminal state: the run is over and the screen waits for a key.
_TERMINAL = ("done", "error", "cancelled")

# Progress-line matchers, in order of precedence.
_PROGRESS_BYTES = re.compile(r"AUDIOCPP_PROGRESS downloaded=(\d+) total=(\d+)")
_PROGRESS_PERCENT = re.compile(r"(\d{1,3})%")
_PROGRESS_COUNT = re.compile(r"\[(\d+)/(\d+)\]")

# A spinner frame set for the running step marker.
_SPINNER = ("|", "/", "-", "\\")


@dataclass
class TaskStep:
    """One step of a task view run.

    WORK is ``work(emit, cancel) -> int``: it streams output lines through
    EMIT and returns its exit code (0 = success). CANCEL is a
    ``threading.Event`` the view sets when the user confirms cancel; WORK
    should stop promptly and may return any code (the view reports the run
    as "cancelled" regardless).
    """
    title: str
    work: Callable[[Callable[[str], None], threading.Event], int]


def run_steps(scr, title: str, steps: List[TaskStep]) -> int:
    """Run STEPS in order inside the curses screen; return the first bad rc.

    Returns 0 when every step succeeded, otherwise the first non-zero exit
    code (a cancelled run returns a non-zero code too).
    """
    view = TaskView(scr, title, steps)
    return view.run()


def run_steps_inline(steps: List[TaskStep], emit=None, cancel=None) -> int:
    """Run STEPS in order without the curses view; return the first bad rc.

    The console/CLI counterpart of ``run_steps``: each step's work is called
    directly (EMIT None keeps the current plain-console subprocess behavior),
    and every step runs even when an earlier one failed — matching how the
    wizards warn-and-continue today.
    """
    first = 0
    for step in steps:
        rc = step.work(emit, cancel)
        if rc and not first:
            first = rc
    return first


class TaskView:
    """Draws and drives one list of setup steps; see the module docstring."""

    def __init__(self, scr, title: str, steps: List[TaskStep],
                 clock: Callable[[], float] = time.time):
        import curses
        self.curses = curses
        self.scr = scr
        self.title = title
        self.steps = steps
        self.theme = tui._ensure_theme(curses)
        self._clock = clock
        # -- state -----------------------------------------------------
        self.phase = "running"   # running | done | error | cancelled
        self.current: Optional[int] = None   # index of the running step
        self.results: List[Optional[int]] = [None] * len(steps)
        self.cancelled_step: Optional[int] = None
        self.log_tail: List[str] = []
        self._progress: Optional[Tuple[float, float]] = None  # (done, total)
        self._progress_kind = ""   # "bytes" | "percent" | "count" | ""
        self.step_started: List[Optional[float]] = [None] * len(steps)
        self.finished_at: Optional[float] = None
        self.cancelled = False
        self.cancelling = False
        # -- threads ---------------------------------------------------
        self._queue: Queue = Queue()
        self._cancel = threading.Event()
        self._worker = threading.Thread(target=self._worker_main, daemon=True)

    # ------------------------------------------------------------------
    # Worker
    # ------------------------------------------------------------------

    def _worker_main(self) -> None:
        first_failure = 0
        for index, step in enumerate(self.steps):
            if self._cancel.is_set():
                break
            self._queue.put({"kind": "step_start", "index": index,
                             "title": step.title})
            try:
                with contextlib.redirect_stdout(_LineWriter(self._emit)), \
                        contextlib.redirect_stderr(_LineWriter(self._emit)):
                    rc = step.work(self._emit, self._cancel)
            except Exception as exc:  # noqa: BLE001 - reported to the view
                self._queue.put({"kind": "line",
                                 "text": f"[ERROR] {exc}"})
                rc = 1
            if self._cancel.is_set():
                self._queue.put({"kind": "step_cancelled", "index": index})
                break
            self._queue.put({"kind": "step_done", "index": index, "rc": rc})
            if rc != 0:
                first_failure = first_failure or rc
                # Keep going where the console path would only warn; the
                # failing step stays marked [FAIL].
        if self._cancel.is_set():
            self._queue.put({"kind": "finish", "phase": "cancelled",
                             "rc": first_failure or 1})
        elif first_failure:
            self._queue.put({"kind": "finish", "phase": "error",
                             "rc": first_failure})
        else:
            self._queue.put({"kind": "finish", "phase": "done", "rc": 0})

    def _emit(self, line: str) -> None:
        """Forward one output line to the view queue (progress-aware)."""
        self._queue.put({"kind": "line", "text": line})

    # ------------------------------------------------------------------
    # Event handling
    # ------------------------------------------------------------------

    def handle_event(self, event: dict) -> None:
        kind = event.get("kind")
        if kind == "step_start":
            self.current = event["index"]
            self.step_started[self.current] = self._now()
            self._progress = None
            self._progress_kind = ""
        elif kind == "line":
            text = event.get("text") or ""
            self._ingest_line(text)
        elif kind == "step_done":
            index = event["index"]
            self.results[index] = event.get("rc") or 0
            self.current = None
            self._progress = None
            self._progress_kind = ""
        elif kind == "step_cancelled":
            self.cancelled_step = event["index"]
            self.current = None
            self._progress = None
            self._progress_kind = ""
        elif kind == "finish":
            self.phase = event.get("phase") or "done"
            self.cancelled = self.phase == "cancelled"
            self.finished_at = self._now()
            self.current = None

    def _ingest_line(self, text: str) -> None:
        """Fold one output line into the log tail and progress bar."""
        line = text.rstrip("\r\n")
        if not line:
            return
        match = _PROGRESS_BYTES.search(line)
        if match:
            total = int(match.group(2))
            done = int(match.group(1))
            self._progress = (done, total)
            self._progress_kind = "bytes"
            return  # machine-readable progress is not part of the log
        match = _PROGRESS_PERCENT.search(line)
        if match:
            percent = int(match.group(1))
            if percent <= 100:
                self._progress = (percent, 100)
                self._progress_kind = "percent"
                # Fall through: keep the line in the log (the tail already
                # collapses rapid \r updates to the last full line).
        else:
            match = _PROGRESS_COUNT.search(line)
            if match:
                done = int(match.group(1))
                total = int(match.group(2))
                if total > 0 and done <= total:
                    self._progress = (done, total)
                    self._progress_kind = "count"
        self.log_tail.append(line)
        if len(self.log_tail) > _LOG_TAIL:
            del self.log_tail[: len(self.log_tail) - _LOG_TAIL]

    def _finish(self, phase: str) -> None:
        self.phase = phase
        if self.finished_at is None:
            self.finished_at = self._now()

    def _now(self) -> float:
        return self._clock()

    # ------------------------------------------------------------------
    # Main loop
    # ------------------------------------------------------------------

    def run(self) -> int:
        scr = self.scr
        try:
            scr.timeout(_DRAW_TIMEOUT_MS)
        except Exception:
            pass
        self._worker.start()
        first_failure = 0
        try:
            while True:
                self._drain()
                self.render()
                key = self._get_key()
                if key is None:
                    continue
                if self.phase in _TERMINAL:
                    return self._result_rc()
                if key in (27, ord("q"), 3) and not self.cancelling:
                    if self._prompt_cancel():
                        self._drain()
                        return self._result_rc()
        finally:
            self._cancel.set()

    def _result_rc(self) -> int:
        """The exit code for the whole run (cancelled counts as failure)."""
        if self.cancelled:
            return 1
        return next((rc for rc in self.results if rc), 0)

    def _get_key(self) -> Optional[int]:
        try:
            key = self.scr.getch()
        except KeyboardInterrupt:
            return 3
        if key == -1:
            return None
        return key

    def _drain(self) -> None:
        while True:
            try:
                event = self._queue.get_nowait()
            except Empty:
                return
            self.handle_event(event)

    def _prompt_cancel(self) -> bool:
        """Esc/q: confirm cancel, then wait for the worker to wind down."""
        self._blocking()
        try:
            answer = tui.confirm(self.scr, "Cancel this step?", default=False,
                                 cancel_value=False)
        finally:
            self._nonblocking()
        if not answer:
            return False
        self.cancelling = True
        self._cancel.set()
        self._worker.join(timeout=60)
        return True

    def _blocking(self) -> None:
        try:
            self.scr.timeout(-1)
        except Exception:
            pass

    def _nonblocking(self) -> None:
        try:
            self.scr.timeout(_DRAW_TIMEOUT_MS)
        except Exception:
            pass

    # ------------------------------------------------------------------
    # Drawing
    # ------------------------------------------------------------------

    def render(self) -> None:
        curses, theme = self.curses, self.theme
        scr = self.scr
        scr.erase()
        height, width = scr.getmaxyx()
        if height < 12 or width < 40:
            _text(scr, theme, height // 2, 2, "Terminal too small",
                  curses.A_BOLD)
            scr.refresh()
            return

        _box(scr, curses, theme, height, width)
        _text(scr, theme, 0, 2, _fit(f" {self.title} ", width - 4),
              theme["title"])

        inner_x = 2
        y = 2
        # -- step list -------------------------------------------------
        for index, step in enumerate(self.steps):
            mark, kind = self._step_mark(index)
            label = _fit(f" {step.title} ", max(8, width - inner_x - 14))
            _text(scr, theme, y, inner_x, mark, theme.get(kind, theme["body"]))
            _text(scr, theme, y, inner_x + 5, label, theme["body"])
            if index == self.current and self.phase not in _TERMINAL:
                started = self.step_started[index] or self._now()
                _text(scr, theme, y, inner_x + 5 + len(label) + 1,
                      f" {_format_elapsed(self._now() - started)}",
                      theme["dim"])
            y += 1

        y += 1
        _sep(scr, curses, theme, y, width)
        y += 1

        # -- progress bar ----------------------------------------------
        if self._progress is not None and self.phase not in _TERMINAL:
            done, total = self._progress
            bar_x = inner_x + 10
            bar_room = max(10, width - bar_x - 16)
            filled = 0
            if total:
                filled = round(bar_room * min(done, total) / total)
                filled = max(0, min(bar_room, filled))
            _text(scr, theme, y, inner_x, "Progress".ljust(9), theme["dim"])
            try:
                scr.addstr(y, bar_x, " " * filled, theme["bar"])
            except Exception:
                pass
            _text(scr, theme, y, bar_x + bar_room + 1,
                  _progress_label(self._progress, self._progress_kind),
                  theme["accent"])
            y += 1

        # -- log tail --------------------------------------------------
        for line in self.log_tail[-_LOG_TAIL:]:
            _text(scr, theme, y, inner_x, _fit(line, width - inner_x - 2),
                  theme["dim"])
            y += 1
            if y >= height - 3:
                break

        # -- footer ----------------------------------------------------
        if self.phase == "done":
            footer = "completed — press any key to return"
            kind = "ok"
        elif self.phase == "cancelled":
            footer = "cancelled — press any key to return"
            kind = "warn"
        elif self.phase == "error":
            footer = "finished with errors — press any key to return"
            kind = "err"
        elif self.cancelling:
            footer = "cancelling..."
            kind = "warn"
        else:
            footer = "Esc or q: cancel"
            kind = "dim"
        _text(scr, theme, height - 2, 2, _fit(footer, width - 4),
              theme[kind])
        scr.refresh()

    def _step_mark(self, index: int) -> Tuple[str, str]:
        """The (mark, kind) for step INDEX."""
        if self.phase in _TERMINAL:
            if index == self.cancelled_step:
                return "[x]", "warn"
            if self.results[index] == 0:
                return "[OK]", "ok"
            if self.results[index] is not None:
                return "[FAIL]", "err"
            return "[  ]", "dim"
        if index == self.current:
            frame = _SPINNER[int(self._now() * 4) % len(_SPINNER)]
            return f"[{frame} ]", "warn"
        if self.results[index] == 0:
            return "[OK]", "ok"
        if self.results[index] is not None:
            return "[FAIL]", "err"
        return "[  ]", "dim"


# ---------------------------------------------------------------------------
# Small helpers (module-level for testability)
# ---------------------------------------------------------------------------

class _LineWriter:
    """A file-like object that forwards writes to a per-line callback.

    Handles carriage-return progress updates (git/tqdm) by treating ``\r``
    as a line terminator too, so the last full line always reflects the
    latest progress.
    """

    def __init__(self, emit: Callable[[str], None]):
        self._emit = emit
        self._buffer = ""

    def write(self, text: str) -> int:
        if not text:
            return 0
        self._buffer += text
        while True:
            cut = _find_line_end(self._buffer)
            if cut < 0:
                break
            line, self._buffer = self._buffer[:cut], self._buffer[cut + 1:]
            if line:
                self._emit(line)
        return len(text)

    def flush(self) -> None:
        if self._buffer:
            self._emit(self._buffer)
            self._buffer = ""

    def isatty(self) -> bool:
        return False


def _find_line_end(text: str) -> int:
    """Index of the earliest ``\n`` or ``\r`` in TEXT, else -1."""
    newline = text.find("\n")
    carriage = text.find("\r")
    if newline < 0:
        return carriage
    if carriage < 0:
        return newline
    return min(newline, carriage)


def _text(scr, theme, y, x, text, attr) -> None:
    try:
        scr.addstr(y, x, text, attr)
    except Exception:
        pass


def _box(scr, curses, theme, height, width) -> None:
    border = theme["border"]
    try:
        scr.addch(0, 0, curses.ACS_ULCORNER, border)
        scr.addch(0, width - 1, curses.ACS_URCORNER, border)
        scr.addch(height - 1, 0, curses.ACS_LLCORNER, border)
        scr.addch(height - 1, width - 1, curses.ACS_LRCORNER, border)
        scr.hline(0, 1, curses.ACS_HLINE, width - 2, border)
        scr.hline(height - 1, 1, curses.ACS_HLINE, width - 2, border)
        for y in range(1, height - 1):
            scr.addch(y, 0, curses.ACS_VLINE, border)
            scr.addch(y, width - 1, curses.ACS_VLINE, border)
    except Exception:
        pass


def _sep(scr, curses, theme, y, width) -> None:
    try:
        scr.addch(y, 0, curses.ACS_LTEE, theme["border"])
        scr.addch(y, width - 1, curses.ACS_RTEE, theme["border"])
        scr.hline(y, 1, curses.ACS_HLINE, width - 2, theme["dim"])
    except Exception:
        pass


def _fit(text: str, width: int) -> str:
    if width < 1:
        return ""
    if len(text) <= width:
        return text
    return text[: max(0, width - 1)] + "~"


def _format_elapsed(seconds: float) -> str:
    seconds = max(0, int(seconds))
    hours, remainder = divmod(seconds, 3600)
    minutes, secs = divmod(remainder, 60)
    if hours:
        return f"{hours}:{minutes:02d}:{secs:02d}"
    return f"{minutes}:{secs:02d}"


def _progress_label(progress: Tuple[float, float], kind: str) -> str:
    done, total = progress
    if kind == "bytes":
        return f"{_fmt_bytes(done)} / {_fmt_bytes(total)}"
    if kind == "count":
        return f"{int(done)}/{int(total)}"
    return f"{int(done)}%"


def _fmt_bytes(size: float) -> str:
    value = float(size)
    for unit in ("B", "KB", "MB", "GB"):
        if value < 1024 or unit == "GB":
            if unit == "B":
                return f"{int(value)}{unit}"
            return f"{value:.1f}{unit}"
        value /= 1024
    return f"{value:.1f}GB"