"""Rich terminal progress reporter for per-sample model generation."""

import asyncio
import time

from rich.console import Console
from rich.progress import (
    BarColumn,
    Progress,
    SpinnerColumn,
    TaskProgressColumn,
    TextColumn,
)



def _fmt(sec):
    sec = int(sec)
    return f'{sec // 60}m{sec % 60:02d}s' if sec >= 60 else f'{sec}s'


class RichTerminalProgress:
    def __init__(self, console=None):
        # accept an EXTERNAL console: CLI phase messages and the live bar must
        # share one console, or the two writers interleave and repaint wrongly
        self.console = console or Console()
        # NON-TERMINAL (pipes, file redirects): rich's live refresh thread
        # misbehaves and can stall the whole run. Disable everything live;
        # log() degrades to a plain console print. Checked HERE so every
        # creation path is safe regardless of caller logic.
        # double check: FORCE_COLOR/FORCE_TERMINAL env can make rich claim
        # terminal-ness while stdout is actually a pipe/file -> live thread
        # stalls. Require the REAL underlying stream to be a tty.
        import sys as _sys

        _tty = getattr(self.console.file, 'isatty', None)
        self.disabled = not (self.console.is_terminal
                             and callable(_tty) and _tty())
        self.overall_id = None
        self.overall_total = 0
        self.progress = Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(complete_style="green", finished_style="bold green"),
            TaskProgressColumn(),
            TextColumn("• Completed {task.completed}/{task.total} [dim]{task.fields[new]}[/dim]"),
            TextColumn("• in-flight [yellow]{task.fields[inflight]}[/yellow] ([yellow]{task.fields[cur]}[/yellow])"),
            TextColumn("• [red]retries {task.fields[retries]}[/red]"),
            TextColumn("• [dim]{task.fields[rate]}/s[/dim]"),
            TextColumn("• [green]{task.fields[elapsed]}[/green]"),
            TextColumn("• [cyan]eta {task.fields[eta]}[/cyan]"),
            console=self.console,
            refresh_per_second=4,
        )
        self.task_id = None
        self.bench_name = ''
        self.bench_tag = ''   # e.g. '[1/6]': persistent benchmark counter
        self._last_phase = ''
        self.started = 0.0
        self.current_started = 0.0
        self.inflight = 0
        self.restored = 0     # checkpoint head start (drives the '+N new' marker)
        self.heartbeat_task = None

    def set_overall(self, total: int, done: int, label: str = 'benches'):
        if self.disabled:
            return
        """Multi-benchmark runs: a bar above the sample bar showing which
        benchmark we are on (also covers loading/scoring phases)."""
        self.overall_total = total
        if self.overall_id is None:
            self.progress.start()
            # the shared column set reads fields from EVERY task: give the
            # overall task the same fields or rendering raises KeyError
            self.overall_id = self.progress.add_task(
                f'[cyan]{label}[/cyan]', total=total, completed=min(done, total),
                new='', inflight=0, cur='0s', retries=0, rate='0.00',
                elapsed='0s', eta='-')
        else:
            self.progress.update(self.overall_id, completed=min(done, total))

    def advance_overall(self):
        if self.disabled:
            return
        if self.overall_id is not None:
            self.progress.advance(self.overall_id)

    def start(self, total: int, description: str, completed: int = 0):
        if self.disabled:
            return
        self.started = time.monotonic()
        self.restored = max(completed, 0)
        if self.task_id is not None:
            return  # one live reporter at a time; reuse across benchmarks
        self.progress.start()
        self.task_id = self.progress.add_task(
            f"[green]{description}",
            total=total,
            completed=min(completed, total),
            new=self._new_txt(completed),
            rate="0.00",
            inflight=0,
            cur="0s",
            retries=0,
            elapsed="0s",
            eta="-",
        )
        self.heartbeat_task = asyncio.create_task(self._heartbeat())

    def _new_txt(self, absolute_done: int) -> str:
        """'(+N new)' marker: samples completed by THIS run, i.e. absolute
        progress minus the checkpoint-restored head start. Empty when the
        run started fresh (nothing was restored, nothing to distinguish)."""
        if not getattr(self, 'restored', 0):
            return ''
        return f'(+{max(absolute_done - self.restored, 0)} new)'

    def reset_samples(self, total: int, description: str, completed: int = 0):
        if self.disabled:
            return
        """Start (or re-target) the per-sample task for the next benchmark."""
        self.started = time.monotonic()
        self.inflight = 0
        self.bench_name = description
        self.restored = max(completed, 0)  # checkpoint head start this bench
        desc = f'[green]{self.bench_tag}{description} · generating[/green]'
        if self.task_id is None:
            self.progress.start()
            self.task_id = self.progress.add_task(
                desc, total=total, completed=min(completed, total),
                new=self._new_txt(completed), rate='0.00', inflight=0,
                cur='0s', elapsed='0s', eta='-', retries=0)
        else:
            self.progress.update(self.task_id, description=desc,
                                 total=total, completed=min(completed, total),
                                 new=self._new_txt(completed), rate='0.00',
                                 inflight=0, cur='0s', elapsed='0s', eta='-',
                                 retries=0)
        # ALWAYS recreate the heartbeat: the previous one may have died during
        # pause/resume cycles between benchmarks (stale reference -> silent
        # death -> frozen clock while the spinner still animates)
        if self.heartbeat_task is not None:
            self.heartbeat_task.cancel()
        self.heartbeat_task = asyncio.create_task(self._heartbeat())

    def set_bench_tag(self, tag: str):
        if self.disabled:
            return
        """Persistent counter shown on the sample bar, e.g. '[1/6]'."""
        self.bench_tag = tag + ' ' if tag else ''
        if self._last_phase:
            self.set_phase(self._last_phase)

    def set_phase(self, phase: str):
        if self.disabled:
            return
        self._last_phase = phase
        """Retag the sample bar with what is happening (generating/scoring/
        writing) -- the bar alone does not say which stage we are in."""
        if self.task_id is not None:
            self.progress.update(
                self.task_id,
                description=f'[green]{self.bench_tag}{self.bench_name} · {phase}[/green]')

    def begin_sample(self, label: str):
        if self.disabled:
            return
        if self.task_id is None:
            return
        self.inflight += 1
        self.current_started = time.monotonic()
        self.progress.update(self.task_id, inflight=self.inflight, cur='0s')

    def set_retries(self, n: int):
        """Show the retry count on the bar (from the adapter's attempt)."""
        if self.task_id is not None:
            self.progress.update(self.task_id, retries=n)

    def rollback(self):
        if self.disabled:
            return
        """Pair a begin_sample that will NOT reach advance (retry path):
        just decrement the in-flight count, no success/fail bookkeeping."""
        self.inflight = max(0, self.inflight - 1)
        if self.task_id is not None:
            self.progress.update(self.task_id, inflight=self.inflight, cur='0s')

    def advance(self, success: bool = True):
        if self.disabled:
            return
        if self.task_id is None:
            return
        task = self.progress.tasks[self.task_id]
        completed = task.completed + 1
        self.inflight = max(0, self.inflight - 1)
        elapsed = max(time.monotonic() - self.started, 1e-6)
        # rate/eta over THIS RUN's fresh samples only: counting the restored
        # head start would print 141/6s = 23/s when 1 sample was generated
        fresh = max(completed - getattr(self, 'restored', 0), 0)
        self.progress.update(
            self.task_id,
            advance=1,
            new=self._new_txt(completed),
            rate=f"{fresh / elapsed:.2f}",
            inflight=self.inflight, cur='0s',
            elapsed=_fmt(elapsed),
            eta=_fmt((task.total - completed) * elapsed / fresh)
                 if fresh and task.total and task.total > completed else '-',
        )

    async def _heartbeat(self):
        """One tick per second: refresh elapsed + current-sample timer.
        Recreated on every reset_samples; must never raise or the clock
        freezes silently."""
        try:
            while self.task_id is not None:
                e = time.monotonic() - self.started
                upd = {'elapsed': f'{int(e) // 60}m{int(e) % 60:02d}s' if e >= 60
                       else f'{int(e)}s'}
                if self.inflight:
                    secs = int(time.monotonic() - self.current_started)
                    upd['cur'] = (f'{secs // 60}m{secs % 60:02d}s'
                                  if secs >= 60 else f'{secs}s')
                try:
                    self.progress.update(self.task_id, **upd)
                except Exception:
                    pass  # task may have been removed mid-tick
                await asyncio.sleep(1)
        except asyncio.CancelledError:
            pass

    def log(self, message: str):
        """Print a status line ABOVE the live bar (safe during live display).

        Plain console.print from another writer while Progress is live causes
        interleaved/repainted output; Progress.print routes through the live
        region correctly.
        """
        if self.disabled or self.task_id is None:
            self.console.print(message, highlight=False)
        else:
            self.progress.print(message, highlight=False)

    def pause(self):
        """Temporarily stop the live display (e.g. while hub downloads print
        their own tqdm); task state is kept and resume() restores the bars."""
        if self.task_id is not None or self.overall_id is not None:
            self.progress.stop()

    def resume(self):
        if self.task_id is not None or self.overall_id is not None:
            self.progress.start()

    def close(self):
        if self.task_id is not None:
            if self.heartbeat_task is not None:
                self.heartbeat_task.cancel()
                self.heartbeat_task = None
            self.progress.stop()
            self.task_id = None