from __future__ import annotations

import time
from collections.abc import Sequence
from typing import ClassVar

from textual.app import App, ComposeResult
from textual.containers import Vertical
from textual.widgets import DataTable, Footer, Header, Label

from hytop.core.format import fmt_elapsed, fmt_window
from hytop.core.sorting import next_sort_field_index, sort_with_missing_last
from hytop.gpu.metrics import RenderColumn, render_columns_for_show_flags
from hytop.gpu.models import MonitorState
from hytop.gpu.service import (
    apply_node_results,
    availability_ready,
    drain_pending_nodes,
)
from hytop.gpu.sort import sort_gpu_keys_grouped


def _format_metric(metric: str, value: object) -> str:
    """Format a GPU metric value for display."""
    if value is None:
        return "-"
    if metric == "temp_c":
        return f"{float(value):.1f}C"
    if metric == "avg_pwr_w":
        return f"{float(value):.1f}W"
    if metric in {"vram_pct", "gpu_pct"}:
        return f"{float(value):.2f}%"
    if metric == "sclk_mhz":
        return f"{float(value):.0f}MHz"
    return str(value)


class GpuMonitorApp(App[int]):
    """Textual TUI application for real-time GPU monitoring via hy-smi.

    This replaces the previous Rich Live + Table rendering loop. The GPU data
    is collected by background threads managed by the caller; this App simply
    polls the shared MonitorState on a timer and updates the DataTable in place.

    The app returns an integer exit code via ``App.exit(return_code)`` so that
    the calling ``run_monitor`` function can propagate the correct process exit
    code to the shell.
    """

    CSS = """
    Screen {
        layout: vertical;
    }

    #gpu-table {
        height: 1fr;
        border: none;
    }

    #error-label {
        height: auto;
        padding: 0 1;
        color: $warning;
    }

    #error-label.hidden {
        display: none;
    }
    """

    BINDINGS: ClassVar[list] = [
        ("q", "quit_with_code", "Quit"),
        ("ctrl+c", "quit_with_code", "Quit"),
        ("s", "next_sort_field", "Next Sort Field"),
        ("S", "toggle_sort_order", "Toggle Sort Order"),
        ("g", "group_sort", "Grouped Sort"),
    ]
    SORT_FIELDS: ClassVar[list[tuple[str, str]]] = [
        ("host", "Host"),
        ("gpu", "GPU"),
        ("temp_c", "Temp"),
        ("avg_pwr_w", "Power"),
        ("vram_pct", "VRAM"),
        ("gpu_pct", "GPU%"),
        ("sclk_mhz", "SCLK"),
    ]

    def __init__(
        self,
        hosts: list[str],
        show_flags: Sequence[str],
        window: float,
        interval: float,
        wait_idle: bool,
        wait_idle_duration: float,
        timeout: float | None,
        state: MonitorState,
        device_filter: set[int] | None,
        started: float,
        **kwargs: object,
    ) -> None:
        super().__init__(**kwargs)
        self._hosts = hosts
        self._show_flags = list(show_flags)
        self._window = window
        self._interval = interval
        self._wait_idle = wait_idle
        self._wait_idle_duration = wait_idle_duration
        self._timeout = timeout
        self._state = state
        self._device_filter = device_filter
        self._started = started
        self._columns: list[RenderColumn] = render_columns_for_show_flags(show_flags)
        # Map (host, gpu) -> DataTable row key for efficient cell updates.
        self._row_keys: dict[tuple[str, int], str] = {}
        # Last displayed key order; when unchanged we use update_cell to preserve scroll/cursor.
        self._ordered_keys: list[tuple[str, int]] = []
        # Column keys for metric cells (host, gpu, then metric columns).
        self._cell_column_keys: list[str] = ["host", "gpu"]
        for col in self._columns:
            self._cell_column_keys.append(col.metric)
            if col.avg_label is not None:
                self._cell_column_keys.append(f"{col.metric}_avg")
        self._sort_mode: str = "grouped"
        self._sort_desc: bool = False
        self._sort_field_index: int = 0

    # ------------------------------------------------------------------
    # Composition
    # ------------------------------------------------------------------

    def compose(self) -> ComposeResult:
        yield Header()
        with Vertical():
            yield DataTable(id="gpu-table", cursor_type="row")
            yield Label("", id="error-label")
        yield Footer()

    def on_mount(self) -> None:
        """Initialise columns and start the polling timer."""
        table: DataTable = self.query_one("#gpu-table", DataTable)
        table.add_column("Host", key="host")
        table.add_column("GPU", key="gpu")
        for col in self._columns:
            table.add_column(col.label, key=col.metric)
            if col.avg_label is not None:
                avg_key = f"{col.metric}_avg"
                table.add_column(
                    f"{col.avg_label}@{fmt_window(self._window)}",
                    key=avg_key,
                )
        render_interval = min(self._interval, 0.5)
        self.set_interval(render_interval, self._tick)

    # ------------------------------------------------------------------
    # Actions
    # ------------------------------------------------------------------

    def action_quit_with_code(self) -> None:
        """Quit the monitor with a user-interrupt exit code."""
        self.exit(130)

    def action_next_sort_field(self) -> None:
        """Cycle sortable fields and switch to metric sorting mode."""
        if self._sort_mode == "grouped":
            self._sort_mode = "metric"
        self._sort_field_index = next_sort_field_index(
            self._sort_field_index, len(self.SORT_FIELDS)
        )

    def action_toggle_sort_order(self) -> None:
        """Toggle ascending / descending metric sort order."""
        self._sort_desc = not self._sort_desc

    def action_group_sort(self) -> None:
        """Reset to grouped default ordering."""
        self._sort_mode = "grouped"
        self._sort_desc = False

    def _sort_status_text(self) -> str:
        if self._sort_mode == "grouped":
            return "sort=grouped(host->gpu)"
        _, label = self.SORT_FIELDS[self._sort_field_index]
        order = "desc" if self._sort_desc else "asc"
        return f"sort={label.lower()} {order}"

    # ------------------------------------------------------------------
    # Internal timer callback
    # ------------------------------------------------------------------

    def _tick(self) -> None:
        """Drain pending collector results and refresh the DataTable."""
        state = self._state

        # Process any new data from collector threads.
        apply_node_results(
            nodes=drain_pending_nodes(hosts=self._hosts, state=state),
            device_filter=self._device_filter,
            state=state,
        )

        # Update the effective monitored key set when no filter is applied.
        if self._device_filter is None:
            state.monitored_keys = state.discovered_keys.copy()

        now = time.monotonic()
        elapsed = now - self._started

        # ---- Update header subtitle with current runtime ----
        self.title = (
            f"hytop gpu  |  interval={self._interval:.2f}s"
            f"  |  window={fmt_window(self._window)}"
            f"  |  elapsed={fmt_elapsed(elapsed)}"
            f"  |  {self._sort_status_text()}"
        )

        # ---- Rebuild table rows ----
        key_list = sort_gpu_keys_grouped(state.monitored_keys, self._hosts)
        if self._sort_mode == "metric":
            field, _ = self.SORT_FIELDS[self._sort_field_index]

            def _metric_value(key: tuple[str, int]) -> float | int | None:
                if field == "host":
                    return key[0]
                if field == "gpu":
                    return key[1]
                history = state.histories.get(key)
                if history is None:
                    return None
                latest = history.latest()
                if latest is None:
                    return None
                if field in {"temp_c", "avg_pwr_w", "vram_pct", "gpu_pct", "sclk_mhz"}:
                    value = getattr(latest, field, None)
                    return float(value) if value is not None else None
                return None

            key_list = sort_with_missing_last(key_list, _metric_value, self._sort_desc)
        table: DataTable = self.query_one("#gpu-table", DataTable)

        # Build keys_to_display: keys we actually have data for.
        keys_to_display: list[tuple[str, int]] = []
        for key in key_list:
            history = state.histories.get(key)
            if history is None:
                continue
            latest = history.latest()
            if latest is None:
                continue
            keys_to_display.append(key)

        def _cell_values(key: tuple[str, int]) -> list[str]:
            host, gpu = key
            history = state.histories.get(key)
            latest = history.latest() if history else None
            if history is None or latest is None:
                return []
            stale = (now - latest.ts) > self._window
            if stale:
                return [host, str(gpu)] + ["-"] * (len(self._cell_column_keys) - 2)
            values: list[str] = []
            for col in self._columns:
                metric_value = getattr(latest, col.metric, None)
                values.append(_format_metric(col.metric, metric_value))
                if col.avg_label is not None:
                    if metric_value is None:
                        values.append("-")
                    else:
                        values.append(
                            _format_metric(
                                col.metric,
                                history.avg(col.metric, self._window, now),
                            )
                        )
            return [host, str(gpu), *values]

        if keys_to_display == self._ordered_keys:
            # Same order: incremental update to preserve scroll position and cursor.
            for key in keys_to_display:
                row_key = self._row_keys.get(key)
                if row_key is None:
                    continue
                values = _cell_values(key)
                if len(values) != len(self._cell_column_keys):
                    continue
                for col_key, val in zip(self._cell_column_keys, values, strict=True):
                    table.update_cell(row_key, col_key, val)
        else:
            # Order or set changed: full rebuild.
            table.clear(columns=False)
            self._row_keys.clear()
            self._ordered_keys = keys_to_display.copy()
            for key in keys_to_display:
                host, gpu = key
                values = _cell_values(key)
                if not values:
                    continue
                row_key = f"{host}:{gpu}"
                table.add_row(*values, key=row_key)
                self._row_keys[key] = row_key

        # ---- Update error label ----
        error_label: Label = self.query_one("#error-label", Label)
        if state.errors:
            parts = [f"{h}: {err}" for h, err in state.errors.items()]
            error_label.update("  ".join(parts))
            error_label.remove_class("hidden")
        else:
            error_label.update("")
            error_label.add_class("hidden")

        # ---- Check wait-idle exit condition ----
        if self._wait_idle:
            warmup_done = elapsed >= self._wait_idle_duration
            if warmup_done and availability_ready(
                window=self._window,
                histories=state.histories,
                monitored_keys=state.monitored_keys,
                hosts=self._hosts,
                errors=state.errors,
            ):
                self.exit(0)

        # ---- Check global timeout ----
        if self._timeout is not None and elapsed >= self._timeout:
            self.exit(124)
