"""Performance metrics widget for live streaming stats display.

Displays latency, token count, tokens/sec, and memory usage during
and after streaming responses in the TUI chat panel.

Widget lifecycle:
- Hidden when not streaming
- Shown when streaming starts (start())
- Updated on each token (on_token())
- Frozen when streaming ends (stop())
- Cleared on next user submission (hide())
"""

from __future__ import annotations

import time

from textual.widgets import Static


def _get_rss_mb() -> float | None:
    """Read process RSS memory in MB via the resource module.

    Note: ru_maxrss reflects the kernel's last reported RSS, not a
    tracked high-water mark.  On macOS it is in bytes; on Linux, KB.

    Returns:
        RSS in megabytes, or None if the resource module is unavailable.
    """
    try:
        import resource
        import sys

        usage = resource.getrusage(resource.RUSAGE_SELF)
        if sys.platform == "darwin":
            return usage.ru_maxrss / (1024 * 1024)  # bytes -> MB
        else:
            return usage.ru_maxrss / 1024  # KB -> MB
    except (ImportError, AttributeError):
        return None


class PerfMetrics(Static):
    """Single-row status bar showing live streaming performance metrics.

    Display format:
        [dim italic] 1.2s  142 tok  118 tok/s  24.3 RSS [/]

    The widget manages its own visibility state via the _visible flag.
    CSS toggling is handled by the parent (ChatPanel) reading _visible
    to add/remove a CSS class.
    """

    DEFAULT_CSS = """
    PerfMetrics {
        height: 1;
        display: none;
    }

    PerfMetrics.streaming-active {
        display: block;
    }
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self._token_count: int = 0
        self._start_time: float = 0.0
        self._frozen: bool = False
        self._visible: bool = False
        # Cached frozen values for display after stop()
        self._frozen_elapsed: float = 0.0
        self._frozen_rate: float = 0.0

    # -------------------------------------------------------------------------
    # Public lifecycle API
    # -------------------------------------------------------------------------

    def start(self) -> None:
        """Begin tracking metrics. Resets all counters and marks widget visible.

        No-op if already streaming (visible and not frozen) to prevent
        mid-flight counter resets from duplicate start() calls.
        """
        if self._visible and not self._frozen:
            return
        self._token_count = 0
        self._start_time = time.perf_counter()
        self._frozen = False
        self._frozen_elapsed = 0.0
        self._frozen_rate = 0.0
        self._visible = True
        self._sync_css()
        self._refresh_display()

    def on_token(self, token: str) -> None:
        """Record a streaming token arrival.

        Args:
            token: The token text (used only to confirm non-empty; count is
                   incremented regardless of content).
        """
        if self._frozen:
            return
        self._token_count += 1
        self._refresh_display()

    def stop(self) -> None:
        """Freeze current values. Widget stays visible for the user to read."""
        self._frozen = True
        self._frozen_elapsed = self._get_elapsed()
        self._frozen_rate = self._get_tokens_per_sec()
        self._refresh_display()

    def hide(self) -> None:
        """Reset all state and mark widget as hidden."""
        self._token_count = 0
        self._start_time = 0.0
        self._frozen = False
        self._frozen_elapsed = 0.0
        self._frozen_rate = 0.0
        self._visible = False
        self._sync_css()
        self._refresh_display()

    # -------------------------------------------------------------------------
    # Internal calculation helpers
    # -------------------------------------------------------------------------

    def _get_elapsed(self) -> float:
        """Wall-clock seconds since streaming started.

        Returns frozen value if stopped.
        """
        if self._frozen:
            return self._frozen_elapsed
        if self._start_time == 0.0:
            return 0.0
        return time.perf_counter() - self._start_time

    def _get_tokens_per_sec(self) -> float:
        """Current token throughput rate.

        Returns frozen value if stopped.
        """
        if self._frozen:
            return self._frozen_rate
        elapsed = self._get_elapsed()
        if elapsed <= 0.0 or self._token_count == 0:
            return 0.0
        return self._token_count / elapsed

    def _get_memory_mb(self) -> float | None:
        """Current process RSS in megabytes.

        Returns:
            RSS in MB, or None if unavailable on this platform.
        """
        return _get_rss_mb()

    # -------------------------------------------------------------------------
    # Display formatting
    # -------------------------------------------------------------------------

    def _format_display(self) -> str:
        """Build the single-line metrics string.

        Returns:
            Formatted string with latency, tokens, rate, and memory.
        """
        elapsed = self._get_elapsed()
        rate = self._get_tokens_per_sec()
        mem = self._get_memory_mb()

        mem_str = f"{mem:.1f} RSS" if mem is not None else "N/A"

        return (
            f"{elapsed:.1f}s  "
            f"{self._token_count} tok  "
            f"{rate:.0f} tok/s  "
            f"{mem_str}"
        )

    def _refresh_display(self) -> None:
        """Update the widget's rendered content."""
        if not self._visible:
            self.update("")
            return
        display = self._format_display()
        self.update(f"[dim italic]{display}[/]")

    def _sync_css(self) -> None:
        """Synchronize CSS class with visibility state."""
        if self._visible:
            self.add_class("streaming-active")
        else:
            self.remove_class("streaming-active")
