"""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,
)


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}"),
            TextColumn("• in-flight {task.fields[inflight]} ({task.fields[cur]})"),
            TextColumn("• {task.fields[rate]}/s"),
            TextColumn("• {task.fields[elapsed]}"),
            TextColumn("• eta {task.fields[eta]}"),
            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.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),
                inflight=0, rate='0.00', success=done, failed=0,
                waiting='00:00', last_result='')
        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()
        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),
            success=completed,
            failed=0,
            rate="0.00",
            inflight=0,
            cur="0s",
            elapsed="0s",
            eta="-",
            waiting="00:00",
            last_result="restored",
        )
        self.heartbeat_task = asyncio.create_task(self._heartbeat())

    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
        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),
                success=completed, failed=0, rate='0.00', inflight=0,
                cur='0s', elapsed='0s', eta='-',
                waiting='00:00', last_result='restored')
            self.heartbeat_task = asyncio.create_task(self._heartbeat())
        else:
            self.progress.update(self.task_id, description=desc,
                                 total=total, completed=min(completed, total),
                                 success=completed, failed=0, rate='0.00',
                                 inflight=0, cur='0s', elapsed='0s', eta='-',
                                 last_result='restored')

    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',
                             waiting="00:00", last_result=f"waiting {label}")

    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
        ok = task.fields["success"] + (1 if success else 0)
        failed = task.fields["failed"] + (0 if success else 1)
        self.inflight = max(0, self.inflight - 1)
        elapsed = max(time.monotonic() - self.started, 1e-6)
        self.progress.update(
            self.task_id,
            advance=1,
            success=ok,
            failed=failed,
            rate=f"{completed / elapsed:.2f}",
            inflight=self.inflight, cur='0s',
            elapsed=_fmt(elapsed),
            eta=_fmt((task.total - completed) * elapsed / completed)
                 if completed and task.total and task.total > completed else '-',
            waiting="00:00",
            last_result="success" if success else "failed",
        )

    async def _heartbeat(self):
        while self.task_id is not None:
            if 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')
                self.progress.update(self.task_id, **upd)
            await asyncio.sleep(1)

    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