#!/usr/bin/env python3
import argparse
import json
import pathlib
import subprocess
import sys
from datetime import datetime, timezone

from PyQt5 import QtCore, QtGui, QtWidgets

REPORTS_DIR_DEFAULT = "/var/lib/saikyo-av/reports"


def utc_now() -> str:
    return datetime.now(tz=timezone.utc).isoformat(timespec="seconds")


def load_reports(reports_dir: pathlib.Path):
    items = []
    if not reports_dir.exists():
        return items
    for p in sorted(reports_dir.glob("*.json"), key=lambda x: x.stat().st_mtime, reverse=True):
        try:
            data = json.loads(p.read_text(encoding="utf-8"))
        except Exception:
            data = {}
        items.append(
            {
                "id": p.stem,
                "path": str(p),
                "summary": data.get("summary") or "(no summary)",
                "severity": data.get("severity") or "unknown",
                "created_utc": data.get("created_utc") or "",
                "raw": data,
            }
        )
    return items


class ReportsModel(QtCore.QAbstractTableModel):
    def __init__(self):
        super().__init__()
        self.items = []

    def rowCount(self, parent=QtCore.QModelIndex()):
        return len(self.items)

    def columnCount(self, parent=QtCore.QModelIndex()):
        return 3

    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
        if role != QtCore.Qt.DisplayRole:
            return None
        if orientation == QtCore.Qt.Horizontal:
            return ["Summary", "Severity", "Created (UTC)"][section]
        return str(section + 1)

    def data(self, index, role=QtCore.Qt.DisplayRole):
        if not index.isValid():
            return None
        item = self.items[index.row()]
        col = index.column()

        if role == QtCore.Qt.DisplayRole:
            if col == 0:
                return item.get("summary")
            if col == 1:
                return item.get("severity")
            if col == 2:
                return item.get("created_utc")

        if role == QtCore.Qt.ForegroundRole and col == 1:
            sev = (item.get("severity") or "").lower()
            if sev in {"ok", "info", "low"}:
                return QtGui.QBrush(QtGui.QColor("#2dd4bf"))
            if sev in {"medium", "warn", "warning"}:
                return QtGui.QBrush(QtGui.QColor("#fbbf24"))
            if sev in {"high", "critical", "bad"}:
                return QtGui.QBrush(QtGui.QColor("#fb7185"))
            return QtGui.QBrush(QtGui.QColor("#9bb0d1"))

        return None

    def set_items(self, items):
        self.beginResetModel()
        self.items = items
        self.endResetModel()


class MainWindow(QtWidgets.QMainWindow):
    def __init__(self, reports_dir: pathlib.Path, start_in_tray: bool):
        super().__init__()
        self.reports_dir = reports_dir
        self.start_in_tray = start_in_tray

        self.setWindowTitle("Saikyo Antivirus")
        self.setWindowIcon(QtGui.QIcon.fromTheme("saikyo-av", QtGui.QIcon.fromTheme("security-high")))
        self.setMinimumSize(980, 620)

        self.model = ReportsModel()
        self.table = QtWidgets.QTableView()
        self.table.setModel(self.model)
        self.table.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
        self.table.setSelectionMode(QtWidgets.QAbstractItemView.SingleSelection)
        self.table.horizontalHeader().setStretchLastSection(True)
        self.table.horizontalHeader().setSectionResizeMode(0, QtWidgets.QHeaderView.Stretch)
        self.table.horizontalHeader().setSectionResizeMode(1, QtWidgets.QHeaderView.ResizeToContents)
        self.table.horizontalHeader().setSectionResizeMode(2, QtWidgets.QHeaderView.ResizeToContents)

        self.details = QtWidgets.QPlainTextEdit()
        self.details.setReadOnly(True)
        self.details.setFont(QtGui.QFontDatabase.systemFont(QtGui.QFontDatabase.FixedFont))

        self.det_list = QtWidgets.QListWidget()
        self.det_list.setSelectionMode(QtWidgets.QAbstractItemView.SingleSelection)

        self.btn_quarantine = QtWidgets.QPushButton("Quarantine")
        self.btn_delete = QtWidgets.QPushButton("Delete")
        for b in [self.btn_quarantine, self.btn_delete]:
            b.setCursor(QtCore.Qt.PointingHandCursor)
            b.setMinimumHeight(32)
        self.btn_delete.setStyleSheet("background:#2b1420;color:#ffd4df;border:1px solid #5b2031;")
        self.btn_quarantine.setEnabled(False)
        self.btn_delete.setEnabled(False)

        self.scan_status = QtWidgets.QLabel("")
        self.scan_status.setStyleSheet("color:#9bb0d1")
        self.scan_progress = QtWidgets.QProgressBar()
        self.scan_progress.setRange(0, 0)
        self.scan_progress.setVisible(False)
        self.btn_stop_scan = QtWidgets.QPushButton("Stop scan")
        self.btn_stop_scan.setCursor(QtCore.Qt.PointingHandCursor)
        self.btn_stop_scan.setMinimumHeight(32)
        self.btn_stop_scan.setEnabled(False)

        self._scan_proc: QtCore.QProcess | None = None
        self._scan_timer = QtCore.QElapsedTimer()
        self._scan_tick = QtCore.QTimer(self)
        self._scan_tick.setInterval(500)
        self._scan_tick.timeout.connect(self._update_scan_elapsed)

        self.status = QtWidgets.QLabel("local")
        self.status.setStyleSheet("color:#9bb0d1")

        btn_evidence = QtWidgets.QPushButton("Evidence")
        btn_audit = QtWidgets.QPushButton("Audit")
        btn_collect = QtWidgets.QPushButton("Collect")
        btn_scan = QtWidgets.QPushButton("SCAN")
        btn_enable = QtWidgets.QPushButton("Enable")
        btn_disable = QtWidgets.QPushButton("Disable…")
        btn_refresh = QtWidgets.QPushButton("Refresh")

        for b in [btn_evidence, btn_audit, btn_collect, btn_scan, btn_enable, btn_disable, btn_refresh]:
            b.setCursor(QtCore.Qt.PointingHandCursor)
            b.setMinimumHeight(36)

        btn_scan.setStyleSheet("background:rgba(45,212,191,0.18);border:1px solid rgba(45,212,191,0.35);")

        btn_disable.setStyleSheet("background:#2b1420;color:#ffd4df;border:1px solid #5b2031;")

        left = QtWidgets.QFrame()
        left.setFrameShape(QtWidgets.QFrame.StyledPanel)
        ll = QtWidgets.QVBoxLayout(left)
        ll.setContentsMargins(12, 12, 12, 12)
        ll.setSpacing(10)

        title = QtWidgets.QLabel("Admin checks")
        title.setStyleSheet("font-weight:800;font-size:16px")
        hint = QtWidgets.QLabel("Кнопки запускают проверки с подтверждением (pkexec).")
        hint.setWordWrap(True)
        hint.setStyleSheet("color:#9bb0d1;font-size:12px")

        ll.addWidget(title)
        ll.addWidget(hint)
        ll.addSpacing(4)
        ll.addWidget(btn_evidence)
        ll.addWidget(btn_audit)
        ll.addWidget(btn_collect)
        ll.addSpacing(8)
        ll.addWidget(btn_scan)
        ll.addSpacing(8)
        ll.addWidget(btn_enable)
        ll.addWidget(btn_disable)
        ll.addStretch(1)
        ll.addWidget(btn_refresh)
        ll.addWidget(self.status)

        splitter = QtWidgets.QSplitter()
        splitter.setOrientation(QtCore.Qt.Vertical)
        top = QtWidgets.QWidget()
        tl = QtWidgets.QVBoxLayout(top)
        tl.setContentsMargins(0, 0, 0, 0)
        tl.addWidget(self.table)
        splitter.addWidget(top)
        details_box = QtWidgets.QWidget()
        dl = QtWidgets.QVBoxLayout(details_box)
        dl.setContentsMargins(0, 0, 0, 0)
        dl.setSpacing(8)

        det_head = QtWidgets.QHBoxLayout()
        det_title = QtWidgets.QLabel("Detections")
        det_title.setStyleSheet("font-weight:800;")
        det_head.addWidget(det_title)
        det_head.addStretch(1)
        det_head.addWidget(self.scan_status)
        det_head.addWidget(self.scan_progress)
        det_head.addWidget(self.btn_stop_scan)
        det_head.addWidget(self.btn_quarantine)
        det_head.addWidget(self.btn_delete)
        dl.addLayout(det_head)
        dl.addWidget(self.det_list)

        raw_title = QtWidgets.QLabel("Report JSON")
        raw_title.setStyleSheet("font-weight:800;")
        dl.addWidget(raw_title)
        dl.addWidget(self.details)

        splitter.addWidget(details_box)
        splitter.setStretchFactor(0, 2)
        splitter.setStretchFactor(1, 1)

        central = QtWidgets.QWidget()
        root = QtWidgets.QHBoxLayout(central)
        root.setContentsMargins(16, 16, 16, 16)
        root.setSpacing(12)
        left.setFixedWidth(300)
        root.addWidget(left)
        root.addWidget(splitter)
        self.setCentralWidget(central)

        self.setStyleSheet(
            "QMainWindow{background:#0b1220;color:#e7eefc;}"
            "QFrame{background:#0f1a2e;border:1px solid rgba(255,255,255,0.08);border-radius:14px;}"
            "QTableView{background:#0f1a2e;border:1px solid rgba(255,255,255,0.08);border-radius:14px;}"
            "QHeaderView::section{background:#0c1628;color:#9bb0d1;border:none;padding:8px;font-weight:700;}"
            "QPushButton{background:rgba(96,165,250,0.15);border:1px solid rgba(96,165,250,0.35);border-radius:10px;padding:8px 10px;font-weight:700;}"
            "QPushButton:hover{background:rgba(96,165,250,0.22);}"
            "QPlainTextEdit{background:#0c1628;border:1px solid rgba(255,255,255,0.08);border-radius:14px;}"
        )

        self.table.selectionModel().selectionChanged.connect(self._on_select)

        btn_refresh.clicked.connect(self.refresh)
        btn_evidence.clicked.connect(lambda: self._run_admin("run-evidence"))
        btn_audit.clicked.connect(lambda: self._run_admin("run-audit"))
        btn_collect.clicked.connect(lambda: self._run_admin("collect-artifacts"))
        btn_scan.clicked.connect(self._scan)
        btn_enable.clicked.connect(lambda: self._run_admin("enable-protection"))
        btn_disable.clicked.connect(self._disable)

        self.btn_quarantine.clicked.connect(self._quarantine_selected)
        self.btn_delete.clicked.connect(self._delete_selected)
        self.det_list.itemSelectionChanged.connect(self._on_det_select)
        self.btn_stop_scan.clicked.connect(self._stop_scan)

        self._setup_tray()
        self.refresh()
        if self.start_in_tray:
            QtCore.QTimer.singleShot(50, self.hide)

    def _setup_tray(self):
        self.tray = QtWidgets.QSystemTrayIcon(self.windowIcon(), self)
        menu = QtWidgets.QMenu()
        act_open = menu.addAction("Open")
        act_refresh = menu.addAction("Refresh")
        menu.addSeparator()
        act_quit = menu.addAction("Quit")
        act_open.triggered.connect(self._show)
        act_refresh.triggered.connect(self.refresh)
        act_quit.triggered.connect(QtWidgets.QApplication.quit)
        self.tray.setContextMenu(menu)
        self.tray.setToolTip("Saikyo Antivirus")
        self.tray.activated.connect(lambda r: self._show() if r == QtWidgets.QSystemTrayIcon.Trigger else None)
        self.tray.show()

    def _show(self):
        self.showNormal()
        self.raise_()
        self.activateWindow()

    def closeEvent(self, event: QtGui.QCloseEvent):
        event.ignore()
        self.hide()
        self.tray.showMessage(
            "Saikyo Antivirus",
            "Свернуто в трей. Откройте из значка.",
            QtWidgets.QSystemTrayIcon.Information,
            2000,
        )

    def refresh(self):
        self.reports_dir.mkdir(parents=True, exist_ok=True)
        items = load_reports(self.reports_dir)
        self.model.set_items(items)
        # systemctl status doesn't require root for is-enabled/is-active.
        enabled = "unknown"
        active = "unknown"
        try:
            enabled = subprocess.run(
                ["systemctl", "is-enabled", "saikyo-avd.timer"],
                capture_output=True,
                text=True,
            ).stdout.strip() or enabled
        except Exception:
            pass
        try:
            active = subprocess.run(
                ["systemctl", "is-active", "saikyo-avd.timer"],
                capture_output=True,
                text=True,
            ).stdout.strip() or active
        except Exception:
            pass

        self.status.setText(f"protection: {enabled}/{active} | reports: {len(items)} | {utc_now()}")
        if items:
            self.table.selectRow(0)
        else:
            self.details.setPlainText("No reports yet. Use the buttons to generate reports.")

    def _selected(self):
        rows = self.table.selectionModel().selectedRows()
        if not rows:
            return None
        i = rows[0].row()
        if i < 0 or i >= len(self.model.items):
            return None
        return self.model.items[i]

    def _on_select(self, *_):
        it = self._selected()
        if not it:
            return
        raw = it.get("raw") or {}
        self.details.setPlainText(json.dumps(raw, ensure_ascii=False, indent=2))

        self.det_list.clear()
        self.btn_quarantine.setEnabled(False)
        self.btn_delete.setEnabled(False)

        dets = []
        try:
            dets = (raw.get("details") or {}).get("detections") or []
        except Exception:
            dets = []

        for d in dets:
            p = d.get("path")
            n = d.get("name")
            if not p:
                continue
            item = QtWidgets.QListWidgetItem(f"{p}  [{n or 'unknown'}]")
            item.setData(QtCore.Qt.UserRole, {"path": p, "name": n})
            self.det_list.addItem(item)

    def _on_det_select(self):
        sel = self.det_list.selectedItems()
        ok = bool(sel)
        self.btn_quarantine.setEnabled(ok)
        self.btn_delete.setEnabled(ok)

    def _scan(self):
        if self._scan_proc is not None:
            QtWidgets.QMessageBox.information(self, "Scan running", "Scan is already running.")
            return

        r = QtWidgets.QMessageBox.question(
            self,
            "Full system scan",
            "Run full ClamAV scan of '/'? This may take a long time.\n\nProceed?",
        )
        if r != QtWidgets.QMessageBox.Yes:
            return

        self.scan_progress.setVisible(True)
        self.btn_stop_scan.setEnabled(True)
        self.scan_status.setText("Scanning… 0s")
        self._scan_timer.start()
        self._scan_tick.start()

        proc = QtCore.QProcess(self)
        proc.setProgram("pkexec")
        proc.setArguments(["/usr/sbin/saikyo-av-admin", "scan"])
        proc.setProcessChannelMode(QtCore.QProcess.MergedChannels)

        buf: list[str] = []

        def on_ready():
            data = bytes(proc.readAllStandardOutput()).decode("utf-8", errors="replace")
            if data:
                buf.append(data)

        def on_finished(exit_code, _status):
            self._scan_tick.stop()
            self.scan_progress.setVisible(False)
            self.btn_stop_scan.setEnabled(False)

            out = "".join(buf).strip()
            self._scan_proc = None
            self.scan_status.setText("")

            # Refresh reports list after scan.
            self.refresh()

            if exit_code == 0:
                if out:
                    QtWidgets.QMessageBox.information(self, "Scan complete", out)
            else:
                msg = out or "Scan failed"
                QtWidgets.QMessageBox.warning(self, "Scan failed", msg)

        proc.readyReadStandardOutput.connect(on_ready)
        proc.finished.connect(on_finished)

        self._scan_proc = proc
        proc.start()

    def _selected_detection_path(self) -> str | None:
        sel = self.det_list.selectedItems()
        if not sel:
            return None
        data = sel[0].data(QtCore.Qt.UserRole) or {}
        p = data.get("path")
        if not p:
            return None
        return str(p)

    def _quarantine_selected(self):
        p = self._selected_detection_path()
        if not p:
            return
        r = QtWidgets.QMessageBox.question(
            self,
            "Quarantine file",
            f"Move file to quarantine?\n\n{p}",
        )
        if r != QtWidgets.QMessageBox.Yes:
            return
        self._run_admin_with_arg("quarantine", p)

    def _delete_selected(self):
        p = self._selected_detection_path()
        if not p:
            return
        r1 = QtWidgets.QMessageBox.warning(
            self,
            "Delete infected file",
            f"This will permanently delete the file:\n\n{p}\n\nContinue?",
            QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No,
        )
        if r1 != QtWidgets.QMessageBox.Yes:
            return
        r2 = QtWidgets.QMessageBox.warning(
            self,
            "Confirm delete",
            "Are you absolutely sure? This cannot be undone.",
            QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No,
        )
        if r2 != QtWidgets.QMessageBox.Yes:
            return
        self._run_admin_with_arg("delete", p)

    def _disable(self):
        r = QtWidgets.QMessageBox.question(
            self,
            "Disable protection",
            "Disable protection requires explicit operator consent. Action will be logged.",
        )
        if r != QtWidgets.QMessageBox.Yes:
            return
        self._run_admin("disable-protection")

    def _run_admin(self, action: str):
        cmd = ["pkexec", "/usr/sbin/saikyo-av-admin", action]
        proc = subprocess.run(cmd, capture_output=True, text=True)
        if proc.returncode == 0:
            self.refresh()
            out = (proc.stdout or "").strip()
            if out:
                QtWidgets.QMessageBox.information(self, "Done", out)
        else:
            msg = (proc.stderr or proc.stdout or "failed").strip()
            QtWidgets.QMessageBox.warning(self, "Failed", msg)

    def _update_scan_elapsed(self):
        if self._scan_proc is None:
            return
        ms = self._scan_timer.elapsed()
        self.scan_status.setText(f"Scanning… {int(ms/1000)}s")

    def _stop_scan(self):
        if self._scan_proc is None:
            return
        r = QtWidgets.QMessageBox.question(
            self,
            "Stop scan",
            "Stop the running scan?",
        )
        if r != QtWidgets.QMessageBox.Yes:
            return
        try:
            self._scan_proc.kill()
        except Exception:
            pass

    def _run_admin_with_arg(self, action: str, arg: str):
        cmd = ["pkexec", "/usr/sbin/saikyo-av-admin", action, arg]
        proc = subprocess.run(cmd, capture_output=True, text=True)
        if proc.returncode == 0:
            self.refresh()
            out = (proc.stdout or "").strip()
            if out:
                QtWidgets.QMessageBox.information(self, "Done", out)
        else:
            msg = (proc.stderr or proc.stdout or "failed").strip()
            QtWidgets.QMessageBox.warning(self, "Failed", msg)


def main() -> int:
    ap = argparse.ArgumentParser(prog="saikyo-av-gui")
    ap.add_argument("--tray", action="store_true")
    ap.add_argument("--reports-dir", default=REPORTS_DIR_DEFAULT)
    args = ap.parse_args()

    app = QtWidgets.QApplication(sys.argv)
    w = MainWindow(pathlib.Path(args.reports_dir), start_in_tray=args.tray)
    w.show()
    return app.exec_()


if __name__ == "__main__":
    raise SystemExit(main())
