#!/usr/bin/env python3
"""Headless GUI test: drives SrasViewerWindow through the real Qt widgets,
signals and worker threads under the offscreen platform plugin.

Covers the interactions a manual smoke test would: load, switch angles and
channels, background DC precompute, lazy FFT compute, threshold and bg-sub
changes, angle alignment, manual angle alignment, aligned view, ROI
draw/move, and CSV export.

Usage:  QT_QPA_PLATFORM=offscreen python tools/test_gui.py [file.sras]
"""

import json
import os
import sys
import tempfile
from pathlib import Path
from unittest.mock import patch

os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")

import numpy as np                                          # noqa: E402
from PyQt6.QtCore import QEventLoop, Qt, QTimer              # noqa: E402
from PyQt6.QtTest import QTest                               # noqa: E402
from PyQt6.QtWidgets import QApplication, QMessageBox        # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import sras_compute as compute                              # noqa: E402
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX           # noqa: E402
from sras_viewer import RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX  # noqa: E402
import tools.make_test_sras as gen                          # noqa: E402

_failures: list[str] = []


def check(name: str, ok: bool, detail: str = ""):
    print(f"  {'PASS' if ok else 'FAIL'}  {name}" + (f"  — {detail}" if detail else ""))
    if not ok:
        _failures.append(name)


def pump(ms: int = 250):
    """Run the event loop for a while so queued signals and worker threads
    make progress."""
    loop = QEventLoop()
    QTimer.singleShot(ms, loop.quit)
    loop.exec()


def wait_until(pred, timeout_ms: int = 20000, step: int = 100) -> bool:
    waited = 0
    while waited < timeout_ms:
        if pred():
            return True
        pump(step)
        waited += step
    return pred()


def main():
    app = QApplication(sys.argv)
    errors: list[str] = []

    tmpdir = Path(tempfile.mkdtemp(prefix="sras_gui_"))
    path = Path(sys.argv[1]) if len(sys.argv) > 1 else tmpdir / "gui.sras"
    if len(sys.argv) <= 1:
        gen.write(path, n_angles=4, seed=11, samples_per_frame=256)

    print(f"\nloading {path.name}")
    win = SrasViewerWindow()
    win.show()
    # Capture anything the app reports as an error via the status bar.
    win.statusBar().messageChanged.connect(
        lambda m: errors.append(m) if m and "error" in m.lower() else None)

    win._load_file(str(path))
    check("file loaded", wait_until(lambda: win._sras is not None))
    s = win._sras
    check("parsed as v6", s.version == 6, f"v{s.version}")
    check("defaults to CH4", win.combo_channel.currentIndex() == CH4_IDX)
    check("image displayed", win._current_image is not None)
    check("angle spinbox ranges over all angles",
          win.spin_angle.maximum() == s.n_angles - 1)
    check("scan info populated",
          win._info["Angles"].text() == f"Angles: {s.n_angles}",
          win._info["Angles"].text())

    print("\nbackground DC precompute (all angles)")
    ok = wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
                                and (a, CH3_IDX) in win._dc_cache
                                for a in range(s.n_angles)))
    check("every angle cached for CH3 and CH4", ok,
          f"{len(win._dc_cache)} entries")
    check("status label reports completion",
          "ready for all angles" in win.lbl_dc_precompute.text(),
          win.lbl_dc_precompute.text())

    print("\nangle switching (DC, should be served from cache)")
    for a in range(s.n_angles):
        win.spin_angle.setValue(a)
        win._on_view_changed()
        pump(60)
        expected = win._sras.image_shape(a)
        check(f"angle {a} shows its own geometry {expected}",
              win._current_image.shape == expected,
              str(win._current_image.shape))
    check("no compute job needed for cached DC angles",
          not win._job_running("compute"))

    print("\nchannel switching")
    win.spin_angle.setValue(0)
    win._on_view_changed()
    pump(60)
    win.combo_channel.setCurrentIndex(CH3_IDX)
    check("CH3 displayed", wait_until(lambda: win._current_ch == CH3_IDX))

    win.combo_channel.setCurrentIndex(CH1_IDX)
    check("CH1 (FFT) computed", wait_until(
        lambda: win._current_ch == CH1_IDX and not win._job_running("compute")))
    check("FFT result cached", len(win._fft_cache) > 0, f"{len(win._fft_cache)} keys")
    rf_img = win._current_image
    check("FFT image is non-degenerate", len(np.unique(rf_img)) > 1,
          f"{len(np.unique(rf_img))} unique values")

    print("\nvelocity mode (pure post-multiply, no recompute)")
    n_fft_before = len(win._fft_cache)
    win.combo_channel.setCurrentIndex(VELOCITY_MODE_IDX)
    check("velocity displayed", wait_until(
        lambda: win._current_ch == VELOCITY_MODE_IDX and not win._job_running("compute")))
    grating = win.spin_grating_um.value()
    check("velocity == freq x grating",
          np.allclose(win._current_image, rf_img * grating, atol=1e-3))
    check("velocity reused the cached FFT", len(win._fft_cache) == n_fft_before,
          f"{n_fft_before} -> {len(win._fft_cache)}")
    check("grating spinbox visible in velocity mode", win.grp_velocity.isVisible())

    print("\nthreshold change (genuine cache-key change)")
    win.combo_channel.setCurrentIndex(CH1_IDX)
    wait_until(lambda: not win._job_running("compute"))
    dc4 = win._dc_cache[(0, CH4_IDX)]
    win.spin_threshold_mv.setValue(float(np.median(dc4)))
    win._on_threshold_changed()
    check("recomputed at new threshold", wait_until(
        lambda: not win._job_running("compute") and len(win._fft_cache) > n_fft_before))
    check("masking zeroed some pixels",
          int((win._current_image == 0).sum()) > 0,
          f"{int((win._current_image == 0).sum())} of {win._current_image.size}")

    print("\nbackground subtraction toggle")
    n_before = len(win._fft_cache)
    win.chk_bg_sub.setChecked(False)
    check("recomputed without bg-sub", wait_until(
        lambda: not win._job_running("compute") and len(win._fft_cache) > n_before))
    win.chk_bg_sub.setChecked(True)
    pump(200)
    check("returning to bg-sub was a cache hit (no recompute)",
          not win._job_running("compute"))

    print("\nROI")
    x = s.x_axis_mm(0)
    y = s.y_positions_mm(0)
    roi = RoiQuad.from_bbox(float(x[1]), float(y[1]),
                            float(x[-2]), float(y[-2]))
    win.image_canvas.set_roi(roi)
    pump(120)
    check("ROI registered", win.image_canvas.get_roi() is not None)
    check("pixel count reported",
          "pixels inside" in win.lbl_roi_npix.text()
          and win.lbl_roi_npix.text() != "pixels inside: —",
          win.lbl_roi_npix.text())
    npix = int(win.lbl_roi_npix.text().split(":")[1])
    check("ROI pixel count is plausible",
          0 < npix <= win._current_image.size, f"{npix}")
    check("Export ROI enabled", win.btn_export_roi.isEnabled())

    csv_path = tmpdir / "roi.csv"
    with patch("sras_viewer.QFileDialog.getSaveFileName",
               return_value=(str(csv_path), "")):
        win._on_export_roi_csv()
    check("ROI CSV written", csv_path.exists())
    if csv_path.exists():
        body = [l for l in csv_path.read_text().splitlines() if not l.startswith("#")]
        check("ROI CSV has header + one line per pixel",
              len(body) == npix + 1, f"{len(body)} lines for {npix} pixels")

    img_csv = tmpdir / "img.csv"
    with patch("sras_viewer.QFileDialog.getSaveFileName",
               return_value=(str(img_csv), "")):
        win._on_export_csv()
    check("image CSV written", img_csv.exists())
    if img_csv.exists():
        arr = np.loadtxt(img_csv, delimiter=",")
        check("image CSV round-trips the displayed image",
              arr.shape == win._current_image.shape
              and np.allclose(arr, win._current_image, rtol=1e-5, atol=1e-4))

    print("\nROI survives angle and channel switches")
    win.spin_angle.setValue(1)
    win._on_view_changed()
    wait_until(lambda: not win._job_running("compute"))
    check("ROI still present after angle switch",
          win.image_canvas.get_roi() is not None)
    win.combo_channel.setCurrentIndex(CH4_IDX)
    wait_until(lambda: win._current_ch == CH4_IDX)
    check("ROI still present after channel switch",
          win.image_canvas.get_roi() is not None)

    print("\nangle alignment (Fusion)")
    win.spin_angle.setValue(0)
    win._on_view_changed()
    wait_until(lambda: not win._job_running("compute"))
    check("alignment action enabled", win._alignment_act.isEnabled())
    win._on_angle_alignment()
    check("alignment completed", wait_until(
        lambda: win._alignment_result is not None and not win._job_running("align"),
        timeout_ms=60000))
    if win._alignment_result is not None:
        r = win._alignment_result
        check("transform for every angle", len(r.per_angle) == s.n_angles)
        check("canvas is at least as large as any single angle",
              all(r.canvas_shape[0] >= int(s.n_rows[a])
                  and r.canvas_shape[1] >= int(s.n_frames[a])
                  for a in range(s.n_angles)), str(r.canvas_shape))
        check("reference angle has zero shift",
              r.per_angle[r.ref_angle_idx].shift_mm == (0.0, 0.0))
        check("Aligned View auto-enabled and checked",
              win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked())
        pump(200)
        check("displayed image is on the alignment canvas",
              win.image_canvas._img_shape == r.canvas_shape,
              f"{win.image_canvas._img_shape} vs {r.canvas_shape}")

        win.chk_aligned_view.setChecked(False)
        pump(200)
        check("unchecking returns to the raw per-angle grid",
              win.image_canvas._img_shape == s.image_shape(0),
              str(win.image_canvas._img_shape))

    print("\nmanual alignment (Fusion)")
    check("manual alignment action enabled", win._manual_align_act.isEnabled())

    # --- Alignment pivot is a signal-weighted centroid, not the raw bbox --
    # center, and is independent of any DC threshold (so a threshold that
    # happens to leave a real angle's binary mask empty can't silently
    # degrade the pivot back to the bbox center).
    corner_signal = np.zeros(s.image_shape(0), dtype=np.float32)
    corner_signal[0, 0] = 1.0   # single spike -> weighted centroid is exact
    expected_corner = (float(s.x_axis_mm(0)[0]), float(s.y_positions_mm(0)[0]))
    centroid = compute._signal_centroid_mm(s, 0, corner_signal)
    check("signal-weighted centroid of a single spike pixel is that pixel exactly",
          np.allclose(centroid, expected_corner), f"{centroid} vs {expected_corner}")
    bbox_center = compute._bbox_center_mm(s, 0)
    check("signal centroid differs from the raw scan-window bbox center",
          not np.allclose(centroid, bbox_center),
          f"centroid {centroid} vs bbox center {bbox_center}")

    # compute_pivot_points_mm should reuse a pre-computed dc4_mv dict rather
    # than recomputing from the real DC4 image (which has no such spike and
    # would give a different answer if silently recomputed).
    reused_pivot = compute.compute_pivot_points_mm(s, dc4_mv={0: corner_signal})[0]
    check("compute_pivot_points_mm reuses a pre-computed dc4_mv dict",
          np.allclose(reused_pivot, expected_corner))

    # A perfectly flat signal carries no information to weight by, so it
    # falls back to the bbox center rather than producing a NaN/degenerate
    # centroid.
    flat_signal = np.full(s.image_shape(0), 5.0, dtype=np.float32)
    flat_centroid = compute._signal_centroid_mm(s, 0, flat_signal)
    check("a perfectly flat signal falls back to the bbox center",
          np.allclose(flat_centroid, bbox_center))

    # --- Rotation sign convention: negative of the raw angles_deg delta ----
    check("_theta_deg negates the raw angles_deg delta (GR stage's positive "
          "angle is the opposite rotational sense from this module's CCW "
          "math convention)",
          all(np.isclose(compute._theta_deg(s, a, 0),
                         -(float(s.angles_deg[a]) - float(s.angles_deg[0])))
              for a in range(s.n_angles)))

    # --- Open: must NOT seed from the still-live automatic AlignmentResult --
    # The automatic result's translation comes from FFT phase correlation --
    # the very thing manual mode exists to work around -- so manual mode
    # must start from identity (centroids coincide, zero shift) regardless
    # of whatever the automatic run last computed. Only a previously *saved
    # manual* alignment (sidecar) should ever seed this dialog.
    win._on_manual_alignment()
    check("dialog opened", win._manual_align_dialog is not None)
    dlg = win._manual_align_dialog
    check("mask prep needed no background worker (already DC-cached)",
          not win._job_running("manual_align_masks"))
    check("no manual sidecar yet -> dialog starts at identity, not the "
          "automatic result",
          all(dlg._angle_params[a] == compute.ManualAngleParams()
              for a in range(s.n_angles)))

    # --- Reference angle is locked -------------------------------------------
    dlg.combo_active_angle.setCurrentIndex(dlg._ref_angle_idx)
    pump(30)
    before_ref = dlg._angle_params[dlg._ref_angle_idx]
    dlg._on_nudge_translate(1, 0, False)
    dlg._on_nudge_rotate(1, False)
    check("reference angle group disabled", not dlg.grp_manual_adjust.isEnabled())
    check("reference angle untouched by nudge attempts",
          dlg._angle_params[dlg._ref_angle_idx] == before_ref)

    # --- Nudging a real angle (fine + coarse, translate + rotate) -----------
    active = 1 if s.n_angles > 1 else 0
    dlg.combo_active_angle.setCurrentIndex(active)
    pump(30)
    before = dlg._angle_params[active].shift_mm
    dlg._on_nudge_translate(1, 0, False)   # fine +X
    fine_step = dlg.spin_step_translate_mm.value()
    check("fine translate nudge moved shift_x by exactly one fine step",
          abs(dlg._angle_params[active].shift_mm[0] - (before[0] + fine_step)) < 1e-9)

    before = dlg._angle_params[active].shift_mm
    dlg._on_nudge_translate(0, -1, True)   # coarse -Y
    coarse_step = fine_step * dlg.spin_step_multiplier.value()
    check("coarse translate nudge uses the multiplier",
          abs(dlg._angle_params[active].shift_mm[1] - (before[1] - coarse_step)) < 1e-9)

    before_rot = dlg._angle_params[active].rotation_deg
    dlg._on_nudge_rotate(1, False)
    check("rotate nudge changed rotation_deg",
          dlg._angle_params[active].rotation_deg != before_rot)
    check("preview canvas rebuilt for every angle after a rotation nudge",
          len(dlg._preview_layers) == s.n_angles)

    # --- Real key-event wiring (proves keyPressEvent -> signal -> slot) -----
    before = dlg._angle_params[active].shift_mm
    QTest.keyClick(dlg.canvas, Qt.Key.Key_Right)
    check("a real Right-arrow key event nudged shift_x",
          dlg._angle_params[active].shift_mm[0] > before[0])

    # --- Auto De-rotate: rotation only, translation untouched ---------------
    shift_before_derotate = dlg._angle_params[active].shift_mm
    dlg._on_auto_derotate()
    expected_theta = compute._theta_deg(s, active, dlg._ref_angle_idx)
    check("auto de-rotate set the known analytic angle",
          abs(dlg._angle_params[active].rotation_deg - expected_theta) < 1e-6)
    check("auto de-rotate left translation untouched",
          dlg._angle_params[active].shift_mm == shift_before_derotate)
    check("reference angle stays identity after auto de-rotate",
          dlg._angle_params[dlg._ref_angle_idx].rotation_deg == 0.0)

    # --- Save -----------------------------------------------------------------
    dlg._on_save()
    sidecar = compute.sidecar_path(s.path)
    check("sidecar file written", sidecar.exists())
    raw = json.loads(sidecar.read_text()) if sidecar.exists() else {}
    check("sidecar schema_version is current",
          raw.get("schema_version") == compute._SIDECAR_SCHEMA_VERSION)
    check("sidecar per_angle round-trips the dialog's resolved params",
          all(raw.get("per_angle", {}).get(str(a), {}).get("rotation_deg")
              == dlg._angle_params[a].rotation_deg for a in range(s.n_angles)))
    check("main window's alignment_result replaced by the manual build",
          win._alignment_result is not None
          and win._alignment_result.per_angle[active].rotation_deg
          == dlg._angle_params[active].rotation_deg)
    check("Aligned View auto-enabled after Save",
          win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked())

    # --- An old-schema sidecar (pre-pivot/sign fix) is treated as absent ------
    stale = dict(raw)
    stale["schema_version"] = compute._SIDECAR_SCHEMA_VERSION - 1
    sidecar.write_text(json.dumps(stale))
    check("a sidecar with an old schema_version is not loaded",
          compute.load_manual_alignment(s) is None)
    sidecar.write_text(json.dumps(raw))   # restore for the rest of this section

    # --- Clear (with confirmation) --------------------------------------------
    with patch("sras_viewer.QMessageBox.question",
               return_value=QMessageBox.StandardButton.Yes):
        dlg._on_clear()
    check("sidecar file deleted", not sidecar.exists())
    check("dialog params reset to identity",
          all(dlg._angle_params[a] == compute.ManualAngleParams()
              for a in range(s.n_angles)))
    check("main window alignment_result cleared", win._alignment_result is None)
    check("Aligned View disabled after Clear",
          not win.chk_aligned_view.isEnabled() and not win.chk_aligned_view.isChecked())

    dlg.close()
    pump(150)
    check("dialog reference released on close", win._manual_align_dialog is None)

    # --- Sidecar auto-restore on next load ------------------------------------
    win._on_manual_alignment()
    dlg = win._manual_align_dialog
    dlg.combo_active_angle.setCurrentIndex(active)
    pump(30)
    dlg._on_auto_derotate()
    dlg._on_nudge_translate(1, 1, True)
    saved_rotation = dlg._angle_params[active].rotation_deg
    saved_shift = dlg._angle_params[active].shift_mm
    dlg._on_save()
    dlg.close()
    pump(150)

    old_sras_id = id(win._sras)
    win._load_file(str(path))   # reload the same file fresh
    check("file reloaded", wait_until(
        lambda: win._sras is not None and id(win._sras) != old_sras_id))
    s = win._sras
    check("manual dialog force-closed by a reload", win._manual_align_dialog is None)
    check("reload restores the saved manual alignment automatically",
          win._alignment_result is not None)
    if win._alignment_result is not None:
        check("restored rotation matches what was saved",
              abs(win._alignment_result.per_angle[active].rotation_deg
                  - saved_rotation) < 1e-9)
        check("restored shift matches what was saved",
              win._alignment_result.per_angle[active].shift_mm == saved_shift)
        check("Aligned View auto-checked after restoring a saved alignment",
              win.chk_aligned_view.isChecked())

    print("\npixel inspector")
    win.chk_aligned_view.setChecked(False)
    pump(100)
    win._on_pixel_clicked(0, 0)
    pump(150)
    check("waveform hint hidden after a click", win.lbl_wave_hint.isHidden())
    win.combo_channel.setCurrentIndex(CH1_IDX)
    wait_until(lambda: not win._job_running("compute"))
    win._on_pixel_clicked(1, 1)
    pump(150)
    check("RF waveform panel rendered",
          len(win.wave_canvas.ax_wave.lines) > 0,
          f"{len(win.wave_canvas.ax_wave.lines)} lines")

    print("\nshutdown")
    win.close()
    pump(400)
    check("all background jobs released", len(win._jobs) == 0,
          f"{list(win._jobs)}")

    print()
    unexpected = [e for e in errors if e]
    if unexpected:
        print(f"status-bar errors seen: {unexpected}")
        _failures.append("status-bar errors")

    if _failures:
        print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures))
        return 1
    print("All GUI checks passed.")
    return 0


if __name__ == "__main__":
    sys.exit(main())