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


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()
        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]}"),
            TextColumn("• {task.fields[rate]}/s"),
            TextColumn("•"),
            TimeElapsedColumn(),
            TextColumn("•"),
            TimeRemainingColumn(),
            console=self.console,
            refresh_per_second=4,
        )
        self.task_id = None
        self.started = 0.0
        self.current_started = 0.0
        self.inflight = 0
        self.heartbeat_task = None

    def start(self, total: int, description: str, completed: int = 0):
        self.started = time.monotonic()
        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,
            waiting="00:00",
            last_result="restored",
        )
        self.heartbeat_task = asyncio.create_task(self._heartbeat())

    def begin_sample(self, label: str):
        if self.task_id is None:
            return
        self.inflight += 1
        self.current_started = time.monotonic()
        self.progress.update(self.task_id, inflight=self.inflight,
                             waiting="00:00", last_result=f"waiting {label}")

    def rollback(self):
        """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)

    def advance(self, success: bool = True):
        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,
            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 and self.inflight:
                waiting = int(time.monotonic() - self.current_started)
                self.progress.update(self.task_id, waiting=f"{waiting // 60:02d}:{waiting % 60:02d}")
            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.task_id is not None:
            self.progress.print(message)
        else:
            self.console.print(message)

    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