#!/usr/bin/env python3
"""GUI for freeze-watcher service.

Layout:
  - Status bar (service state + Capture Now)
  - Live metrics row (always visible)
  - Tabs: Live graphs | Snapshots (filterable list + detail) | Summary
"""

import re
import subprocess
import sys
from collections import Counter, deque
from datetime import datetime
from pathlib import Path

from PySide6.QtCore import Qt, QTimer, QPointF, Signal
from PySide6.QtGui import QFont, QColor, QAction, QPainter, QPen, QBrush, QPolygonF
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QListWidget, QListWidgetItem, QTextEdit, QPushButton,
    QSplitter, QGroupBox, QGridLayout, QSystemTrayIcon, QMenu, QStyle,
    QTabWidget, QLineEdit, QSizePolicy, QMessageBox,
)

LOG_DIR = Path("/var/log/freeze-watcher")
TRIGGER = LOG_DIR / ".trigger"
SERVICE = "freeze-watcher.service"
HISTORY_SECONDS = 3600  # 1h of 1-second samples for graph
GRAPH_UPDATE_MS = 1000


# --------------------------------------------------------------------- helpers

def read_psi(name):
    try:
        line = (Path("/proc/pressure") / name).read_text().splitlines()[0]
        m = re.search(r"avg10=([\d.]+)", line)
        return float(m.group(1)) if m else 0.0
    except Exception:
        return 0.0


def read_dstate_count():
    try:
        out = subprocess.run(
            ["ps", "-eo", "state", "--no-headers"],
            capture_output=True, text=True, timeout=2,
        ).stdout
        return sum(1 for ln in out.splitlines() if ln.startswith("D"))
    except Exception:
        return 0


def read_loadavg():
    try:
        return float(Path("/proc/loadavg").read_text().split()[0])
    except Exception:
        return 0.0


def read_procs_blocked():
    try:
        for ln in Path("/proc/stat").read_text().splitlines():
            if ln.startswith("procs_blocked"):
                return int(ln.split()[1])
    except Exception:
        pass
    return 0


def service_active():
    r = subprocess.run(["systemctl", "is-active", SERVICE], capture_output=True, text=True)
    return r.stdout.strip() == "active"


# ----------------------------------------------------------------- widgets

class MetricLabel(QLabel):
    def __init__(self, name):
        super().__init__("--")
        self.name = name
        self.setAlignment(Qt.AlignCenter)
        self.setMinimumHeight(48)
        f = QFont()
        f.setPointSize(18)
        f.setBold(True)
        self.setFont(f)
        self._set_color("#888")

    def _set_color(self, hex_):
        self.setStyleSheet(
            f"background:#1e1e1e; color:{hex_}; border:1px solid #333; "
            f"border-radius:6px; padding:6px;"
        )

    def set_value(self, value, warn, crit, suffix=""):
        self.setText(f"{value}{suffix}")
        if value >= crit:
            self._set_color("#ff5555")
        elif value >= warn:
            self._set_color("#f5a623")
        else:
            self._set_color("#5cdb95")


class LineGraph(QWidget):
    """Simple time-series line graph drawn with QPainter. No extra deps."""

    def __init__(self, label, color="#5cdb95", warn=None, crit=None):
        super().__init__()
        self.label = label
        self.color = QColor(color)
        self.warn = warn
        self.crit = crit
        self.values = deque(maxlen=HISTORY_SECONDS)
        self.markers = set()  # indices to mark as snapshot events
        self.setMinimumHeight(110)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred)

    def push(self, value, marker=False):
        self.values.append(value)
        if marker:
            self.markers.add(len(self.values) - 1)
        # Decay markers as values shift out
        new_markers = set()
        for m in self.markers:
            shifted = m - 1 if len(self.values) == HISTORY_SECONDS else m
            if shifted >= 0:
                new_markers.add(shifted)
        self.markers = new_markers
        self.update()

    def paintEvent(self, _):
        p = QPainter(self)
        p.setRenderHint(QPainter.Antialiasing)
        w, h = self.width(), self.height()

        # Background
        p.fillRect(0, 0, w, h, QColor("#1a1a1a"))

        if not self.values:
            p.setPen(QColor("#666"))
            p.drawText(self.rect(), Qt.AlignCenter, f"{self.label} — no data yet")
            return

        max_v = max(max(self.values), self.crit or 0, 1)
        min_v = 0
        rng = max_v - min_v

        # Title and current value
        p.setPen(QColor("#aaa"))
        p.setFont(QFont("monospace", 9))
        current = self.values[-1]
        p.drawText(8, 16, f"{self.label}  current={current:.1f}  max={max(self.values):.1f}")

        # Threshold lines
        plot_top = 24
        plot_bot = h - 8
        plot_h = plot_bot - plot_top
        for thresh, col in [(self.warn, "#665022"), (self.crit, "#5a2222")]:
            if thresh is None or thresh > max_v:
                continue
            y = plot_bot - (thresh - min_v) / rng * plot_h
            p.setPen(QPen(QColor(col), 1, Qt.DashLine))
            p.drawLine(0, int(y), w, int(y))

        # Line
        n = len(self.values)
        if n < 2:
            return
        step = w / (n - 1) if n > 1 else 0
        poly = QPolygonF()
        for i, v in enumerate(self.values):
            x = i * step
            y = plot_bot - (v - min_v) / rng * plot_h
            poly.append(QPointF(x, y))
        p.setPen(QPen(self.color, 1.5))
        p.drawPolyline(poly)

        # Markers (snapshot captures)
        p.setPen(QPen(QColor("#ff5555"), 1, Qt.DotLine))
        for m in self.markers:
            x = m * step
            p.drawLine(int(x), plot_top, int(x), plot_bot)


# ----------------------------------------------------------------- snapshot parsing

def parse_snapshot(path):
    """Extract reason, diagnosis, dstate procs, wchans, top-cpu procs from a snapshot."""
    try:
        text = path.read_text(errors="replace")
    except Exception:
        return None
    info = {
        "path": str(path),
        "mtime": path.stat().st_mtime,
        "reason": "",
        "diagnoses": [],
        "dstate_procs": [],
        "wchans": [],
        "top_cpu_proc": "",
    }
    lines = text.splitlines()
    i = 0
    while i < len(lines):
        ln = lines[i]
        if ln.startswith("=== FREEZE EVENT TRIGGERED:"):
            info["reason"] = ln.split("TRIGGERED:", 1)[1].strip().rstrip("=").strip()
        elif ln.startswith("## Diagnosis"):
            i += 1
            while i < len(lines) and lines[i].startswith("  •"):
                info["diagnoses"].append(lines[i].lstrip(" •").strip())
                i += 1
            continue
        elif ln.startswith("## D-state processes"):
            i += 1
            while i < len(lines) and not lines[i].startswith("##"):
                parts = lines[i].split()
                if len(parts) >= 6 and parts[0] == "D":
                    info["dstate_procs"].append(parts[4])  # comm
                    info["wchans"].append(parts[3])
                i += 1
            continue
        elif ln.startswith("## Top 15 by CPU"):
            j = i + 1
            while j < len(lines) and not lines[j].startswith("##"):
                parts = lines[j].split()
                if len(parts) >= 6 and parts[1] != "user":
                    info["top_cpu_proc"] = parts[5]
                    break
                j += 1
        i += 1
    return info


# ----------------------------------------------------------------- main window

class FreezeMonitor(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Freeze Watcher")
        self.resize(1200, 800)

        central = QWidget()
        self.setCentralWidget(central)
        root = QVBoxLayout(central)

        # Status row
        status_row = QHBoxLayout()
        self.status_label = QLabel()
        self.status_label.setStyleSheet("padding:6px; font-weight:bold;")
        status_row.addWidget(self.status_label)
        status_row.addStretch()
        self.btn_capture_now = QPushButton("📸 Capture Now")
        self.btn_capture_now.setToolTip("Manually trigger a snapshot capture")
        self.btn_capture_now.clicked.connect(self.capture_now)
        status_row.addWidget(self.btn_capture_now)
        self.btn_start = QPushButton("Start")
        self.btn_stop = QPushButton("Stop")
        self.btn_start.clicked.connect(lambda: self._svc("start"))
        self.btn_stop.clicked.connect(lambda: self._svc("stop"))
        status_row.addWidget(self.btn_start)
        status_row.addWidget(self.btn_stop)
        root.addLayout(status_row)

        # Live metrics row
        metrics_box = QGroupBox("Live metrics (1s refresh)")
        grid = QGridLayout(metrics_box)
        labels = ["PSI IO %", "PSI CPU %", "PSI Mem %", "D-state", "Blocked", "Load (1m)"]
        self.metrics = {}
        for i, name in enumerate(labels):
            grid.addWidget(QLabel(name, alignment=Qt.AlignCenter), 0, i)
            m = MetricLabel(name)
            self.metrics[name] = m
            grid.addWidget(m, 1, i)
        root.addWidget(metrics_box)

        # Tabs
        self.tabs = QTabWidget()
        root.addWidget(self.tabs, stretch=1)

        self._build_graphs_tab()
        self._build_snapshots_tab()
        self._build_summary_tab()

        self._known_events = set()
        self._all_snapshots = []  # cached parses

        # Tray
        self.tray = QSystemTrayIcon(self.style().standardIcon(QStyle.SP_ComputerIcon), self)
        self.tray.setToolTip("Freeze Watcher")
        menu = QMenu()
        show_act = QAction("Show window", self)
        show_act.triggered.connect(self.showNormal)
        capture_act = QAction("Capture now", self)
        capture_act.triggered.connect(self.capture_now)
        quit_act = QAction("Quit", self)
        quit_act.triggered.connect(QApplication.quit)
        menu.addAction(show_act)
        menu.addAction(capture_act)
        menu.addSeparator()
        menu.addAction(quit_act)
        self.tray.setContextMenu(menu)
        self.tray.activated.connect(self._tray_clicked)
        self.tray.show()

        # Tick timer
        self.timer = QTimer()
        self.timer.timeout.connect(self.tick)
        self.timer.start(GRAPH_UPDATE_MS)
        self.load_events()
        self.tick()

    # ---- tabs

    def _build_graphs_tab(self):
        tab = QWidget()
        v = QVBoxLayout(tab)
        v.addWidget(QLabel("<b>Last hour</b> — red dashed lines mark snapshot captures"))
        self.graph_load = LineGraph("Load (1m)", "#5cdb95", warn=8, crit=16)
        self.graph_psi_io = LineGraph("PSI io.some %", "#88ccff", warn=30, crit=60)
        self.graph_psi_cpu = LineGraph("PSI cpu.some %", "#ffcc66", warn=30, crit=60)
        self.graph_dstate = LineGraph("D-state procs", "#ff8888", warn=3, crit=6)
        for g in (self.graph_load, self.graph_psi_io, self.graph_psi_cpu, self.graph_dstate):
            v.addWidget(g)
        self.tabs.addTab(tab, "Live graphs")

    def _build_snapshots_tab(self):
        tab = QWidget()
        v = QVBoxLayout(tab)

        filter_row = QHBoxLayout()
        filter_row.addWidget(QLabel("Filter:"))
        self.filter_edit = QLineEdit()
        self.filter_edit.setPlaceholderText("substring match against reason/diagnosis (e.g. btrfs, dstate, usb)")
        self.filter_edit.textChanged.connect(self._apply_filter)
        filter_row.addWidget(self.filter_edit)
        self.btn_refresh = QPushButton("Refresh")
        self.btn_refresh.clicked.connect(self.load_events)
        self.btn_open = QPushButton("Open log dir")
        self.btn_open.clicked.connect(lambda: subprocess.Popen(["xdg-open", str(LOG_DIR)]))
        self.btn_clear = QPushButton("Delete all")
        self.btn_clear.clicked.connect(self.clear_snapshots)
        filter_row.addWidget(self.btn_refresh)
        filter_row.addWidget(self.btn_open)
        filter_row.addWidget(self.btn_clear)
        v.addLayout(filter_row)

        splitter = QSplitter(Qt.Horizontal)
        self.event_list = QListWidget()
        self.event_list.itemSelectionChanged.connect(self.show_selected)
        splitter.addWidget(self.event_list)
        self.detail = QTextEdit()
        self.detail.setReadOnly(True)
        self.detail.setFont(QFont("monospace"))
        splitter.addWidget(self.detail)
        splitter.setSizes([400, 800])
        v.addWidget(splitter, stretch=1)
        self.tabs.addTab(tab, "Snapshots")

    def _build_summary_tab(self):
        tab = QWidget()
        v = QVBoxLayout(tab)
        top_row = QHBoxLayout()
        top_row.addWidget(QLabel("<b>Cross-snapshot correlation</b> — patterns across all captures"))
        top_row.addStretch()
        btn = QPushButton("Refresh summary")
        btn.clicked.connect(self.refresh_summary)
        top_row.addWidget(btn)
        v.addLayout(top_row)
        self.summary_text = QTextEdit()
        self.summary_text.setReadOnly(True)
        self.summary_text.setFont(QFont("monospace"))
        v.addWidget(self.summary_text)
        self.tabs.addTab(tab, "Summary")

    # ---- actions

    def _tray_clicked(self, reason):
        if reason == QSystemTrayIcon.Trigger:
            if self.isVisible():
                self.hide()
            else:
                self.showNormal()
                self.raise_()

    def closeEvent(self, ev):
        ev.ignore()
        self.hide()

    def _svc(self, action):
        subprocess.run(["pkexec", "systemctl", action, SERVICE])
        self.tick()

    def capture_now(self):
        try:
            TRIGGER.write_text("manual (gui)")
            self.tray.showMessage(
                "Freeze Watcher",
                "Manual capture requested",
                QSystemTrayIcon.Information,
                2000,
            )
        except Exception as e:
            QMessageBox.warning(self, "Capture failed",
                                f"Could not write trigger file at {TRIGGER}:\n{e}\n\n"
                                "Make sure you're in the group owning /var/log/freeze-watcher.")

    # ---- tick

    def tick(self):
        psi_io = read_psi("io")
        psi_cpu = read_psi("cpu")
        psi_mem = read_psi("memory")
        dstate = read_dstate_count()
        blocked = read_procs_blocked()
        load = read_loadavg()

        self.metrics["PSI IO %"].set_value(int(psi_io), 30, 60, "")
        self.metrics["PSI CPU %"].set_value(int(psi_cpu), 30, 60, "")
        self.metrics["PSI Mem %"].set_value(int(psi_mem), 20, 50, "")
        self.metrics["D-state"].set_value(dstate, 3, 6, "")
        self.metrics["Blocked"].set_value(blocked, 3, 6, "")
        self.metrics["Load (1m)"].set_value(int(load), 8, 16, "")

        active = service_active()
        if active:
            self.status_label.setText("● Service: ACTIVE")
            self.status_label.setStyleSheet("padding:6px; font-weight:bold; color:#5cdb95;")
        else:
            self.status_label.setText("● Service: STOPPED")
            self.status_label.setStyleSheet("padding:6px; font-weight:bold; color:#ff5555;")

        self.tray.setToolTip(
            f"Freeze Watcher — {'active' if active else 'stopped'}\n"
            f"PSI io {int(psi_io)}% / cpu {int(psi_cpu)}% — D:{dstate} blk:{blocked} load:{load:.2f}"
        )

        # Detect new snapshot files for graph markers
        current = set(p.name for p in LOG_DIR.glob("snap-*.log")) if LOG_DIR.exists() else set()
        new = current - self._known_events
        is_marker = bool(new)
        self.graph_load.push(load, marker=is_marker)
        self.graph_psi_io.push(psi_io, marker=is_marker)
        self.graph_psi_cpu.push(psi_cpu, marker=is_marker)
        self.graph_dstate.push(dstate, marker=is_marker)

        if current != self._known_events:
            self._known_events = current
            self._all_snapshots = []  # invalidate cache
            self.load_events()
            for n in new:
                self.tray.showMessage(
                    "Freeze Watcher",
                    f"Captured: {n}",
                    QSystemTrayIcon.Warning,
                    5000,
                )

    # ---- snapshots tab

    def load_events(self):
        self.event_list.clear()
        if not LOG_DIR.exists():
            return
        files = sorted(LOG_DIR.glob("snap-*.log"), key=lambda p: p.stat().st_mtime, reverse=True)
        if not self._all_snapshots:
            self._all_snapshots = [parse_snapshot(p) for p in files]
            self._all_snapshots = [s for s in self._all_snapshots if s]
        self._known_events = set(p.name for p in files)
        self._apply_filter()

    def _apply_filter(self):
        needle = self.filter_edit.text().lower().strip()
        self.event_list.clear()
        for info in self._all_snapshots:
            text_blob = (info["reason"] + " " + " ".join(info["diagnoses"])).lower()
            if needle and needle not in text_blob:
                continue
            ts = datetime.fromtimestamp(info["mtime"]).strftime("%Y-%m-%d %H:%M:%S")
            diag_short = info["diagnoses"][0][:80] if info["diagnoses"] else ""
            label = f"{ts}\n  reason: {info['reason']}"
            if diag_short:
                label += f"\n  → {diag_short}"
            item = QListWidgetItem(label)
            item.setData(Qt.UserRole, info["path"])
            self.event_list.addItem(item)

    def show_selected(self):
        items = self.event_list.selectedItems()
        if not items:
            return
        path = items[0].data(Qt.UserRole)
        try:
            self.detail.setPlainText(Path(path).read_text())
        except Exception as e:
            self.detail.setPlainText(f"Error: {e}")

    def clear_snapshots(self):
        if not LOG_DIR.exists():
            return
        if QMessageBox.question(self, "Delete", "Delete all captured snapshots?") != QMessageBox.Yes:
            return
        for p in LOG_DIR.glob("snap-*.log"):
            try:
                p.unlink()
            except Exception:
                try:
                    subprocess.run(["pkexec", "rm", str(p)])
                except Exception:
                    pass
        self._all_snapshots = []
        self.load_events()
        self.detail.clear()

    # ---- summary tab

    def refresh_summary(self):
        if not self._all_snapshots:
            self.load_events()
        snaps = self._all_snapshots
        if not snaps:
            self.summary_text.setPlainText("No snapshots captured yet.")
            return

        wchans = Counter()
        procs = Counter()
        top_cpu = Counter()
        diagnoses = Counter()
        hours = Counter()
        first = None
        last = None
        for s in snaps:
            for w in s["wchans"]:
                wchans[w] += 1
            for p in s["dstate_procs"]:
                procs[p] += 1
            if s["top_cpu_proc"]:
                top_cpu[s["top_cpu_proc"]] += 1
            for d in s["diagnoses"]:
                # First clause only, before " — "
                key = d.split("—")[0].strip()[:60]
                diagnoses[key] += 1
            ts = datetime.fromtimestamp(s["mtime"])
            hours[ts.hour] += 1
            first = min(first, ts) if first else ts
            last = max(last, ts) if last else ts

        def topn(ctr, n=10):
            total = sum(ctr.values()) or 1
            lines = []
            for k, v in ctr.most_common(n):
                pct = 100 * v / total
                bar = "█" * int(pct / 2)
                lines.append(f"  {v:5d} ({pct:5.1f}%)  {bar:<25} {k}")
            return "\n".join(lines) if lines else "  (none)"

        hour_dist = []
        for h in range(24):
            count = hours.get(h, 0)
            bar = "▇" * count
            hour_dist.append(f"  {h:02d}:00  {bar} {count}")

        out = []
        out.append(f"Snapshots analyzed: {len(snaps)}")
        if first and last:
            out.append(f"Time range: {first} → {last}")
        out.append("")
        out.append("─── Top diagnoses ───────────────────────────────────────")
        out.append(topn(diagnoses))
        out.append("")
        out.append("─── Top D-state wchans (where stuck) ────────────────────")
        out.append(topn(wchans))
        out.append("")
        out.append("─── Top processes seen in D-state ───────────────────────")
        out.append(topn(procs))
        out.append("")
        out.append("─── Top CPU consumers at capture time ───────────────────")
        out.append(topn(top_cpu))
        out.append("")
        out.append("─── Captures by hour of day ─────────────────────────────")
        out.append("\n".join(hour_dist))
        self.summary_text.setPlainText("\n".join(out))


def main():
    app = QApplication(sys.argv)
    app.setQuitOnLastWindowClosed(False)
    w = FreezeMonitor()
    w.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    main()
