import json
import subprocess
import time
from collections import deque
from typing import Any

import docker
import typer
from rich.text import Text
from textual import work
from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical
from textual.message import Message
from textual.screen import Screen
from textual.widget import Widget
from textual.widgets import DataTable, Footer, Header, RichLog, Sparkline, Static

app = typer.Typer()

_COLUMNS: list[tuple[str, str]] = [
    ("Container", "name"),
    ("CPU %", "cpu"),
    ("Mem %", "mem_pct"),
    ("Mem Usage", "mem_usage"),
    ("Net I/O", "net"),
    ("Block I/O", "block"),
    ("PIDs", "pids"),
]

# Metrics shown in DetailScreen's stats panel, in order.
# (label, docker_key, chartable)
_METRICS: list[tuple[str, str, bool]] = [
    ("CPU %", "CPUPerc", True),
    ("Mem %", "MemPerc", True),
    ("Mem Usage", "MemUsage", False),
    ("Net I/O", "NetIO", False),
    ("Net RX/s", "NetRXs", False),
    ("Net TX/s", "NetTXs", False),
    ("Block I/O", "BlockIO", False),
    ("PIDs", "PIDs", False),
]

_CHARTABLE = {"CPUPerc", "MemPerc"}

_NET_SPARKLINE_METRICS: list[tuple[str, str, str]] = [
    ("NetRXs", "RX/s", "#ffb74d"),
    ("NetTXs", "TX/s", "#f06292"),
]


# ---------------------------------------------------------------------------
# Docker helpers
# ---------------------------------------------------------------------------


def list_compose_stacks() -> list[str]:
    result = subprocess.run(
        ["docker", "compose", "ls", "--format", "json"],
        capture_output=True,
        text=True,
        check=True,
    )
    stacks: list[dict[str, Any]] = json.loads(result.stdout)
    return [s["Name"] for s in stacks if "running" in s.get("Status", "")]


def select_stack() -> str:
    stacks = list_compose_stacks()

    if not stacks:
        typer.echo("No running Docker Compose stacks found.", err=True)
        raise typer.Exit(1)

    if len(stacks) == 1:
        typer.echo(f"Using stack: {stacks[0]}")
        return stacks[0]

    typer.echo("Multiple stacks found:")
    for i, name in enumerate(stacks, start=1):
        typer.echo(f"  {i}. {name}")

    choice = int(typer.prompt("Select stack number", type=int))
    if not 1 <= choice <= len(stacks):
        typer.echo(f"Invalid choice: {choice}", err=True)
        raise typer.Exit(1)

    return stacks[choice - 1]


def validate_broker(stack: str) -> None:
    container = f"{stack}-topology-broker.com-1"
    result = subprocess.run(
        ["docker", "ps", "--filter", f"name=^/{container}$", "--format", "{{.Names}}"],
        capture_output=True,
        text=True,
        check=True,
    )
    if container not in result.stdout.splitlines():
        typer.echo(f"Required container '{container}' is not running.", err=True)
        raise typer.Exit(1)


def fetch_stats(stack: str) -> list[dict[str, Any]]:
    ps = subprocess.run(
        ["docker", "ps", "--filter", f"name={stack}-", "--format", "{{.Names}}"],
        capture_output=True,
        text=True,
        check=True,
    )
    containers = [n for n in ps.stdout.splitlines() if n.strip()]
    if not containers:
        return []

    result = subprocess.run(
        ["docker", "stats", "--no-stream", "--format", "{{ json . }}", *containers],
        capture_output=True,
        text=True,
        check=True,
    )
    return [json.loads(line) for line in result.stdout.splitlines() if line.strip()]


def fetch_container_stats(container: str) -> dict[str, Any]:
    result = subprocess.run(
        ["docker", "stats", "--no-stream", "--format", "{{ json . }}", container],
        capture_output=True,
        text=True,
        check=True,
    )
    lines = [json.loads(line) for line in result.stdout.splitlines() if line.strip()]
    return lines[0] if lines else {}


def fetch_raw_net_bytes(container: str) -> tuple[int, int]:
    client = docker.from_env()
    c = client.containers.get(container)
    stats = c.stats(stream=False)
    nets = stats.get("networks", {})
    rx = sum(v["rx_bytes"] for v in nets.values())
    tx = sum(v["tx_bytes"] for v in nets.values())
    return rx, tx


def _human_net(n: float) -> str:
    for unit in ["B/s", "KB/s", "MB/s", "GB/s"]:
        if n < 1024:
            return f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} TB/s"


# ---------------------------------------------------------------------------
# Widgets
# ---------------------------------------------------------------------------


class MetricsPanel(Widget):
    """Navigable list of metrics. Posts MetricSelected when the cursor moves."""

    class MetricSelected(Message):
        def __init__(self, key: str, label: str, chartable: bool) -> None:
            super().__init__()
            self.key = key
            self.label = label
            self.chartable = chartable

    DEFAULT_CSS = """
    MetricsPanel {
        border: solid $accent;
        padding: 0 1;
    }
    """

    BINDINGS = [
        Binding("up", "cursor_up", "Up", show=False),
        Binding("down", "cursor_down", "Down", show=False),
    ]

    def __init__(self) -> None:
        super().__init__()
        self._values: dict[str, str] = {}
        self._cursor: int = 0

    can_focus = True

    def render(self) -> Text:
        lines = Text()
        for i, (label, key, chartable) in enumerate(_METRICS):
            value = self._values.get(key, "…")
            marker = " *" if chartable else "  "
            line = f"{marker} {label:<12} {value}"
            if i == self._cursor:
                lines.append(line, style="bold reverse")
            else:
                lines.append(line)
            if i < len(_METRICS) - 1:
                lines.append("\n")
        return lines

    def update_values(self, values: dict[str, str]) -> None:
        self._values = values
        self.refresh()

    def action_cursor_up(self) -> None:
        self._cursor = max(0, self._cursor - 1)
        self.refresh()
        self._emit_selected()

    def action_cursor_down(self) -> None:
        self._cursor = min(len(_METRICS) - 1, self._cursor + 1)
        self.refresh()
        self._emit_selected()

    def _emit_selected(self) -> None:
        label, key, chartable = _METRICS[self._cursor]
        self.post_message(MetricsPanel.MetricSelected(key, label, chartable))

    def current_metric(self) -> tuple[str, str, bool]:
        return _METRICS[self._cursor]


_GAUGE_METRICS: list[tuple[str, str, str]] = [
    ("CPUPerc", "CPU %", "#4dd0e1"),
    ("MemPerc", "Mem %", "#81c784"),
]


class LinearGauge(Widget):
    """Htop-style horizontal bar gauge for a single 0–100 % metric."""

    DEFAULT_CSS = """
    LinearGauge {
        height: 1;
        padding: 0 1;
    }
    """

    def __init__(self, label: str, color: str, id: str | None = None) -> None:
        super().__init__(id=id)
        self._label = label
        self._color = color
        self._value: float = 0.0
        self._display: str = "0.0%"

    def set_values(self, value: float, display: str) -> None:
        self._value = value
        self._display = display
        self.refresh()

    def render(self) -> Text:
        label_part = f"{self._label:<6} ["
        suffix = f"] {self._display}"
        bar_width = max(1, self.size.width - len(label_part) - len(suffix))
        filled = round(bar_width * min(self._value, 100.0) / 100.0)
        empty = bar_width - filled

        t = Text(no_wrap=True, overflow="ellipsis")
        t.append(label_part)
        t.append("█" * filled, style=self._color)
        t.append("░" * empty, style="dim")
        t.append(suffix)
        return t


class ChartPanel(Widget):
    """Always displays CPU % and Mem % gauges with sparkline history stacked."""

    DEFAULT_CSS = """
    ChartPanel {
        border: solid $accent;
        height: 100%;
        width: 1fr;
        layout: vertical;
    }
    ChartPanel .gauge-row {
        height: auto;
        layout: vertical;
        margin: 0;
    }
    ChartPanel .net-row {
        height: auto;
        layout: vertical;
        margin: 0;
    }
    ChartPanel Sparkline {
        height: 2;
    }
    ChartPanel #chart-footer {
        height: 1;
        padding: 0 1;
        color: $text-muted;
    }
    #spark-CPUPerc {
        color: #4dd0e1;
    }
    #spark-MemPerc {
        color: #81c784;
    }
    #spark-NetRXs {
        color: #ffb74d;
    }
    #spark-NetTXs {
        color: #f06292;
    }
    ChartPanel .net-label {
        height: 1;
        padding: 0 1;
        margin: 0;
    }
    """

    def compose(self) -> ComposeResult:
        for key, label, color in _GAUGE_METRICS:
            with Vertical(classes="gauge-row"):
                yield LinearGauge(label=label, color=color, id=f"gauge-{key}")
                yield Sparkline([], id=f"spark-{key}")
        for key, label, _color in _NET_SPARKLINE_METRICS:
            with Vertical(classes="net-row"):
                yield Static(f"[dim]{label}[/dim]", classes="net-label", id=f"netlabel-{key}")
                yield Sparkline([], id=f"spark-{key}")
        yield Static("[dim]last 3 minutes[/dim]", id="chart-footer")

    def refresh_charts(self, histories: dict[str, deque[float]], mem_usage: str = "") -> None:
        for key, label, _color in _GAUGE_METRICS:
            data = list(histories.get(key, []))
            maxlen = histories[key].maxlen or 1
            padded = [0.0] * (maxlen - len(data)) + data

            gauge = self.query_one(f"#gauge-{key}", LinearGauge)
            sparkline = self.query_one(f"#spark-{key}", Sparkline)

            sparkline.data = padded

            if data:
                current = data[-1]
                if key == "CPUPerc":
                    display = f"{current:.1f}%"
                else:
                    display = mem_usage if mem_usage else f"{current:.1f}%"
                gauge.set_values(current, display)
            else:
                gauge.set_values(0.0, "n/a")

        for key, label, _color in _NET_SPARKLINE_METRICS:
            data = list(histories.get(key, []))
            if key in histories:
                maxlen = histories[key].maxlen or 1
                padded = [0.0] * (maxlen - len(data)) + data
            else:
                padded = []
            sparkline = self.query_one(f"#spark-{key}", Sparkline)
            sparkline.data = padded
            rate_str = _human_net(data[-1]) if data else "n/a"
            self.query_one(f"#netlabel-{key}", Static).update(f"[dim]{label}[/dim] {rate_str}")


# ---------------------------------------------------------------------------
# TUI screens
# ---------------------------------------------------------------------------


class TableScreen(Screen[None]):
    BINDINGS = [
        Binding("q", "app.quit", "Quit"),
    ]

    def __init__(self, stack: str, interval: float) -> None:
        super().__init__()
        self.stack = stack
        self.interval = interval

    def compose(self) -> ComposeResult:
        yield Header(show_clock=True)
        yield DataTable()
        yield Footer()

    def on_mount(self) -> None:
        table = self.query_one(DataTable)
        table.cursor_type = "row"
        for label, key in _COLUMNS:
            table.add_column(label, key=key)
        self._do_refresh()
        self.set_interval(self.interval, self._do_refresh)

    def _do_refresh(self) -> None:
        self._refresh_stats()

    @work(thread=True)
    def _refresh_stats(self) -> None:
        stats = fetch_stats(self.stack)
        self.app.call_from_thread(self._update_table, stats)

    def _update_table(self, stats: list[dict[str, Any]]) -> None:
        prefix = f"{self.stack}-"
        table = self.query_one(DataTable)
        cursor_row = table.cursor_row

        sorted_stats = sorted(
            stats,
            key=lambda x: float(x.get("CPUPerc", "0%").rstrip("%")),
            reverse=True,
        )

        table.clear()
        for s in sorted_stats:
            full_name = s.get("Name", "")
            display = full_name[len(prefix) :] if full_name.startswith(prefix) else full_name
            table.add_row(
                display,
                s.get("CPUPerc", ""),
                s.get("MemPerc", ""),
                s.get("MemUsage", ""),
                s.get("NetIO", ""),
                s.get("BlockIO", ""),
                s.get("PIDs", ""),
                key=full_name,
            )

        if table.row_count > 0:
            table.move_cursor(row=min(cursor_row, table.row_count - 1))

    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        container = str(event.row_key.value)
        self.app.push_screen(DetailScreen(container, self.stack, self.interval))


class DetailScreen(Screen[None]):
    DEFAULT_CSS = """
    DetailScreen {
        layout: vertical;
    }
    #top-pane {
        height: 16;
        layout: horizontal;
    }
    MetricsPanel {
        width: 1fr;
    }
    ChartPanel {
        width: 1fr;
    }
    #logs-panel {
        height: 1fr;
        border: solid $accent;
    }
    """

    BINDINGS = [
        Binding("q", "pop_screen", "Back"),
        Binding("escape", "pop_screen", "Back"),
    ]

    def __init__(self, container: str, stack: str, interval: float) -> None:
        super().__init__()
        self.container = container
        self.stack = stack
        self.interval = interval
        maxlen = max(1, round(180 / interval))
        self._history: dict[str, deque[float]] = {
            "CPUPerc": deque(maxlen=maxlen),
            "MemPerc": deque(maxlen=maxlen),
            "NetRXs": deque(maxlen=maxlen),
            "NetTXs": deque(maxlen=maxlen),
        }
        self._prev_rx: int | None = None
        self._prev_tx: int | None = None
        self._prev_net_time: float | None = None

    def compose(self) -> ComposeResult:
        yield Header(show_clock=True)
        yield Horizontal(
            MetricsPanel(),
            ChartPanel(id="chart-panel"),
            id="top-pane",
        )
        yield RichLog(id="logs-panel", highlight=True, markup=True)
        yield Footer()

    def on_mount(self) -> None:
        self.sub_title = self.container
        self.query_one(MetricsPanel).focus()
        self._fetch_and_update_stats()
        self.set_interval(self.interval, self._fetch_and_update_stats)
        self._stream_logs()

    @work(thread=True)
    def _fetch_and_update_stats(self) -> None:
        s = fetch_container_stats(self.container)
        now = time.monotonic()
        try:
            rx, tx = fetch_raw_net_bytes(self.container)
        except Exception:
            rx, tx = 0, 0
        self.app.call_from_thread(self._update_metrics, s, rx, tx, now)

    def _update_metrics(self, s: dict[str, Any], rx: int, tx: int, now: float) -> None:
        if self._prev_rx is not None and self._prev_net_time is not None:
            dt = now - self._prev_net_time
            if dt > 0:
                rx_rate = (rx - self._prev_rx) / dt
                tx_rate = (tx - self._prev_tx) / dt  # type: ignore[operator]
                self._history["NetRXs"].append(max(0.0, rx_rate))
                self._history["NetTXs"].append(max(0.0, tx_rate))

        self._prev_rx = rx
        self._prev_tx = tx
        self._prev_net_time = now

        net_rx_str = _human_net(self._history["NetRXs"][-1]) if self._history["NetRXs"] else "…"
        net_tx_str = _human_net(self._history["NetTXs"][-1]) if self._history["NetTXs"] else "…"

        values = {key: s.get(key, "n/a") for _, key, _ in _METRICS}
        values["NetRXs"] = net_rx_str
        values["NetTXs"] = net_tx_str
        self.query_one(MetricsPanel).update_values(values)

        for key in _CHARTABLE:
            raw = s.get(key, "0%")
            try:
                self._history[key].append(float(str(raw).rstrip("%")))
            except ValueError:
                pass

        mem_usage = str(s.get("MemUsage", ""))
        self.query_one("#chart-panel", ChartPanel).refresh_charts(self._history, mem_usage)

    @work(thread=True, exclusive=True)
    def _stream_logs(self) -> None:
        log = self.query_one("#logs-panel", RichLog)
        with subprocess.Popen(
            ["docker", "logs", "--follow", "--tail", "100", self.container],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        ) as proc:
            if proc.stdout is None:
                return
            for line in proc.stdout:
                self.app.call_from_thread(log.write, line.rstrip())

    def action_pop_screen(self) -> None:
        self.app.pop_screen()


# ---------------------------------------------------------------------------
# Textual app
# ---------------------------------------------------------------------------


class StatsApp(App[None]):
    def __init__(self, stack: str, interval: float) -> None:
        super().__init__()
        self.stack = stack
        self.interval = interval

    def on_mount(self) -> None:
        self.push_screen(TableScreen(self.stack, self.interval))


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------


@app.command()
def watch(interval: float = 2.0) -> None:
    stack = select_stack()
    validate_broker(stack)
    StatsApp(stack, interval).run()


if __name__ == "__main__":
    app()
