#!/usr/bin/env python3
"""Background workers for the SRAS viewer.

Every worker is a plain QObject moved onto its own QThread by
SrasViewerWindow._run_worker, exposing signals only. Workers must never touch
GUI-thread-owned state (the display caches in particular) — they take
everything they need through their constructor and hand results back by signal.
"""

import os
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from concurrent.futures.process import BrokenProcessPool
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from matplotlib.backends.backend_agg import FigureCanvasAgg
from matplotlib.figure import Figure
from PyQt6.QtCore import QObject, pyqtSignal

import sras_compute as compute
from sras_align_export import write_aligned_sras
from sras_compute import cache_file, compute_rf_image, dc_image_mv
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile

# Concurrency caps. Batch conversion runs one process per file, and each of
# those processes threads internally, so the two must be divided rather than
# both set to the core count. Files also commonly sit on one external drive,
# where a dozen concurrent readers is slower than a few — hence the low
# default, overridable from the environment.
_BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min(
    4, os.cpu_count() or 2)

# Spawning a pool costs roughly a second of interpreter startup (each child
# re-imports the entry module). That is noise against a multi-GB scan but
# dominates a batch of small files, where it would make the job *slower* —
# so below this total size the batch just runs in the worker thread.
_BATCH_POOL_MIN_BYTES = int(os.environ.get("SRAS_BATCH_POOL_MIN_MB", 512)) * 1024 * 1024


class CancellableWorker(QObject):
    """A worker whose compute polls stop() between row chunks.

    Without this a shutdown has to wait out whatever is in flight, and on a
    large scan a single angle is ~40 s — far too long to block closing the
    window. Chunk-level polling bounds the wait to one chunk instead.
    """

    def __init__(self):
        super().__init__()
        self._stop = False

    def stop(self):
        self._stop = True

    def _stopped(self) -> bool:
        return self._stop


class _PooledWorker(CancellableWorker):
    """Fans a per-item computation across a thread pool, emitting each result
    from this worker's own thread as it lands (never from a pool thread).

    Subclasses provide _plan() -> n_workers (stashing whatever per-run
    context they need), _items(), _one(item) -> result, and _emit(result).
    On stop(): queued items are dropped, in-flight ones are not waited for —
    that is what keeps closing the window responsive on a large scan.
    """
    finished = pyqtSignal()
    error = pyqtSignal(str)

    def run(self):
        try:
            pool = ThreadPoolExecutor(max_workers=max(1, self._plan()))
            try:
                futures = [pool.submit(self._one, it) for it in self._items()]
                for fut in as_completed(futures):
                    if self._stop:
                        break
                    self._emit(fut.result())
            finally:
                pool.shutdown(wait=not self._stop, cancel_futures=True)
            self.finished.emit()
        except Exception as exc:
            self.error.emit(str(exc))


class LoadWorker(QObject):
    finished = pyqtSignal(object)   # SrasFile | None
    error = pyqtSignal(str)

    def __init__(self, path: str):
        super().__init__()
        self._path = path

    def run(self):
        try:
            self.finished.emit(SrasFile(self._path))
        except Exception as exc:
            self.error.emit(str(exc))
            self.finished.emit(None)


class ComputeWorker(CancellableWorker):
    """Computes one displayable image for (angle, channel).

    For CH1/Velocity (FFT-derived) channels, the FFT is only run for pixels
    whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked pixels are
    left at 0 MHz without ever being FFT'd, since that's the expensive part of
    a scan. If the DC4 image for this angle is already known, pass it in as
    *dc4_mv* to skip re-reading the CH4 channel from disk entirely.

    Emits a plain ``np.ndarray`` already in display units.
    """
    finished = pyqtSignal(object)
    error = pyqtSignal(str)

    def __init__(self, sras: SrasFile, angle_idx: int, ch_idx: int,
                 apply_bg_sub: bool = True, n_fft: int | None = None,
                 dc_threshold_mv: float = 0.0,
                 dc4_mv: np.ndarray | None = None,
                 is_fft_mode: bool = False):
        super().__init__()
        self._sras = sras
        self._angle = angle_idx
        self._ch = ch_idx
        self._apply_bg_sub = apply_bg_sub
        self._n_fft = n_fft
        self._dc_threshold = dc_threshold_mv
        self._dc4_mv = dc4_mv
        self._is_fft_mode = is_fft_mode

    def run(self):
        try:
            if self._is_fft_mode:
                img = compute_rf_image(
                    self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
                    apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
                    dc4_mv=self._dc4_mv, should_stop=self._stopped)
            else:
                img = dc_image_mv(self._sras, self._angle, self._ch,
                                  should_stop=self._stopped)
            # On cancellation the image is only partly filled, so hand back
            # None rather than something that would be cached as real. The
            # signal still fires either way — it is what quits the thread.
            self.finished.emit(None if self._stop else img)
        except Exception as exc:
            self.error.emit(str(exc))


class DcPrecomputeWorker(_PooledWorker):
    """Computes CH3/CH4 DC images for every angle in the background.

    DC images are cheap (a per-waveform mean, no FFT) compared to the
    CH1/Velocity FFT, so precomputing them for the whole file right after load
    makes switching angles instant while on a DC channel, and also means the
    FFT masking step (which needs a DC4 image) rarely has to wait on anything.
    """
    angle_done = pyqtSignal(int, np.ndarray, np.ndarray)   # angle_idx, dc3_mv, dc4_mv

    def __init__(self, sras: SrasFile):
        super().__init__()
        self._sras = sras
        self._angle_budget = 0

    def _plan(self) -> int:
        n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
        return n_workers

    def _items(self):
        return range(self._sras.n_angles)

    def _one(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
        # max_workers=1 *and* a budget share: this call is one of several
        # concurrent angles, and both the thread count and the buffer size
        # have to be divided (see compute.plan_angle_level).
        kw = dict(max_workers=1, budget=self._angle_budget,
                  should_stop=self._stopped)
        return (a,
                dc_image_mv(self._sras, a, CH3_IDX, **kw),
                dc_image_mv(self._sras, a, CH4_IDX, **kw))

    def _emit(self, result):
        self.angle_done.emit(*result)


class BatchCacheWorker(QObject):
    """Batch-computes and stores DC or FFT images into each of *paths*'s v7
    CACH tail, in place — converting v6 sources to v7 on first use, or updating
    an existing v7 file's cache blocks without disturbing whatever the other
    block already holds.

    *mode* is ``"dc"`` (CH3/CH4 mean images), ``"fft"`` (CH1 peak-frequency
    images, unmasked — masking is applied at display time, same as v5's PREC
    convention), or ``"fft_rowavg"`` (same-row, distance-weighted CH1
    averaging before the FFT — needs *dc_threshold_mv* and a positive
    *row_avg_n*; see ``sras_compute.cache_file``).

    Both FFT modes cache at *pad_factor*, which the caller sets from the
    viewer's own padding — a cache stored at a pad the user is not viewing
    at is one the display can never use.

    Files are processed one per subprocess: they are fully independent, each
    opens its own memmap and writes only its own bytes, and only path strings
    and scalars cross the process boundary. Emits ``progress(int)`` (0–100 by
    files completed), ``file_done(str, str)`` (path, error message or "") so
    one file's failure doesn't abort the batch, and ``finished()``.
    """
    progress = pyqtSignal(int)
    file_done = pyqtSignal(str, str)
    finished = pyqtSignal()

    def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool,
                dc_threshold_mv: float | None = None, row_avg_n: int = 0,
                pad_factor: int = 1):
        super().__init__()
        self._paths = paths
        self._mode = mode
        self._apply_bg_sub = apply_bg_sub
        self._dc_threshold = dc_threshold_mv
        self._row_avg_n = row_avg_n
        self._pad_factor = pad_factor

    def _report(self, path: str, err: str, done: int, total: int):
        self.file_done.emit(path, err)
        self.progress.emit(int(done / max(1, total) * 100))

    def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]:
        """Process the batch across *n_procs* subprocesses. Returns the paths
        that never got a real answer because the pool itself died, so the
        caller can retry them in-process.

        Under spawn each child re-imports the entry module, so the batch must
        survive that going wrong (an unguarded __main__, a frozen build, a
        sandbox that forbids subprocesses) rather than reporting every file as
        failed — hence the retry list instead of a per-file error.
        """
        # Each child threads internally; divide the machine rather than
        # letting every process claim every core.
        per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs)
        unresolved: list[str] = []
        done = 0

        with ProcessPoolExecutor(max_workers=n_procs) as executor:
            futures = {
                executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
                                compute.get_fft_backend(), per_proc_workers,
                                pad_factor=self._pad_factor,
                                dc_threshold_mv=self._dc_threshold,
                                row_avg_n=self._row_avg_n): p
                for p in paths
            }
            for fut in as_completed(futures):
                path = futures[fut]
                try:
                    err = fut.result()
                except BrokenProcessPool:
                    unresolved.append(path)
                    continue
                except Exception as exc:
                    err = str(exc)
                done += 1
                self._report(path, err, done, len(paths))

        return unresolved

    def _run_inline(self, paths: list[str], done: int, total: int):
        """Fallback / single-file path: compute in this thread. Still uses the
        full core count internally, since nothing else is competing."""
        for path in paths:
            try:
                err = cache_file(path, self._mode, self._apply_bg_sub,
                                 compute.get_fft_backend(),
                                 compute.default_max_workers(),
                                 pad_factor=self._pad_factor,
                                 dc_threshold_mv=self._dc_threshold,
                                 row_avg_n=self._row_avg_n)
            except Exception as exc:
                err = str(exc)
            done += 1
            self._report(path, err, done, total)

    def _worth_pooling(self, paths: list[str]) -> bool:
        if len(paths) < 2:
            return False
        total = 0
        for p in paths:
            try:
                total += os.path.getsize(p)
            except OSError:
                pass          # unreadable files are reported by cache_file
        return total >= _BATCH_POOL_MIN_BYTES

    def run(self):
        paths = self._paths
        n_procs = max(1, min(_BATCH_MAX_PROCS, len(paths)))

        if not self._worth_pooling(paths):
            self._run_inline(paths, 0, len(paths))
            self.finished.emit()
            return

        try:
            unresolved = self._run_pooled(paths, n_procs)
        except Exception:
            # The pool could not be created or collapsed wholesale.
            unresolved = list(paths)

        if unresolved:
            self._run_inline(unresolved, len(paths) - len(unresolved), len(paths))

        self.finished.emit()


@dataclass
class ExportChannel:
    """One row of a batch image export: which raw channel to read, whether
    to apply the Velocity post-multiply, and the fixed display range/labels
    to render it with.

    Kept free of sras_viewer's display constants (CH_LABELS, VELOCITY_MODE_IDX,
    etc.) so BatchExportWorker has no dependency on the GUI module — the
    caller resolves labels/units/filename tags once, up front.
    """
    ch_idx: int        # CH1_IDX, CH3_IDX, or CH4_IDX -- which raw data to read
    is_velocity: bool  # True only for the derived Velocity map (post-multiply of CH1 freq)
    vmin: float
    vmax: float
    label: str          # e.g. "CH3 -- Bias A (DC mean)"
    unit: str           # colorbar units, e.g. "mV"
    tag: str            # filename tag: "CH1", "CH3", "CH4", "VEL"


def _render_map_png(img: np.ndarray, extent: list[float], *, cmap: str,
                    vmin: float, vmax: float, title: str, colorbar_label: str,
                    out_path: str):
    """Render one map image to *out_path* with a fixed vmin/vmax, using a
    headless Agg canvas so this never touches the GUI thread's interactive
    matplotlib backend. Layout mirrors ImageCanvas.show_image."""
    fig = Figure(figsize=(7, 5), tight_layout=True)
    FigureCanvasAgg(fig)
    ax = fig.add_subplot(111)
    im = ax.imshow(img, aspect="auto", origin="upper", extent=extent,
                   cmap=cmap, vmin=vmin, vmax=vmax, interpolation="nearest")
    cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    if colorbar_label:
        cb.set_label(colorbar_label)
    ax.set_xlabel("X (mm)")
    ax.set_ylabel("Y (mm)")
    ax.set_title(title)
    fig.savefig(out_path, dpi=150)


class BatchExportWorker(CancellableWorker):
    """Renders and saves one PNG per (angle, selected channel) for an
    already-open SrasFile, with each channel's vmin/vmax held fixed across
    every angle so the colorbar is directly comparable image to image.

    Takes the SrasFile object directly (like ComputeWorker/DcPrecomputeWorker)
    rather than a path -- this runs against the file already open in the GUI,
    not an arbitrary batch of files, so there is no need to reopen it in a
    subprocess the way BatchCacheWorker does.

    Emits progress(int) (0-100 over all angle x channel pairs), file_done(str,
    str) (output path, error message or ""), and finished() -- same shape as
    BatchCacheWorker.
    """
    progress = pyqtSignal(int)
    file_done = pyqtSignal(str, str)
    finished = pyqtSignal()

    def __init__(self, sras: SrasFile, channels: list[ExportChannel],
                 output_dir: str, prefix: str, *, cmap: str,
                 apply_bg_sub: bool, dc_threshold_mv: float,
                 n_fft: int | None, grating_um: float):
        super().__init__()
        self._sras = sras
        self._channels = channels
        self._output_dir = Path(output_dir)
        self._prefix = prefix
        self._cmap = cmap
        self._apply_bg_sub = apply_bg_sub
        self._dc_threshold_mv = dc_threshold_mv
        self._n_fft = n_fft
        self._grating_um = grating_um

    def _angle_extent(self, angle_idx: int) -> list[float]:
        s = self._sras
        x_axis = s.x_axis_mm(angle_idx)
        y_axis = s.y_positions_mm(angle_idx)
        dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm
        dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
        return [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
                y_axis[-1] + dy / 2, y_axis[0] - dy / 2]

    def run(self):
        try:
            s = self._sras
            n_angles = s.n_angles
            total = n_angles * len(self._channels)
            done = 0
            needs_freq = any(c.ch_idx == CH1_IDX for c in self._channels)

            for angle_idx in range(n_angles):
                if self._stop:
                    break
                extent = self._angle_extent(angle_idx)
                angle_deg = s.angles_deg[angle_idx]

                freq_mhz = None
                if needs_freq:
                    freq_mhz = compute_rf_image(
                        s, angle_idx, dc_threshold_mv=self._dc_threshold_mv,
                        apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
                        should_stop=self._stopped)

                for channel in self._channels:
                    if self._stop:
                        break
                    out_path = (self._output_dir
                               / f"{self._prefix}_angle{angle_idx:02d}_{channel.tag}.png")
                    try:
                        if channel.ch_idx == CH1_IDX:
                            img = (freq_mhz * self._grating_um if channel.is_velocity
                                  else freq_mhz)
                        else:
                            img = dc_image_mv(s, angle_idx, channel.ch_idx,
                                              should_stop=self._stopped)
                        title = f"{channel.tag}  |  {angle_deg:.1f}°"
                        colorbar_label = (f"{channel.label} ({channel.unit})"
                                         if channel.unit else channel.label)
                        _render_map_png(
                            img, extent, cmap=self._cmap,
                            vmin=channel.vmin, vmax=channel.vmax,
                            title=title, colorbar_label=colorbar_label,
                            out_path=str(out_path))
                        self.file_done.emit(str(out_path), "")
                    except Exception as exc:
                        self.file_done.emit(str(out_path), str(exc))
                    done += 1
                    self.progress.emit(int(done / max(1, total) * 100))

            self.finished.emit()
        except Exception as exc:
            self.file_done.emit("", str(exc))
            self.finished.emit()


class Ch4MaskWorker(_PooledWorker):
    """Fetches each requested angle's CH4 (Bias B) DC image in mV, for the
    alignment wizard's initial threshold-mask stack.

    Reuses dc_image_mv, which prefers a stored v5/v7 cache over recomputing
    from raw waveforms, so this only does real work for a file that hasn't
    gone through the v7 "Convert" batch step and for angles the main
    window's own DcPrecomputeWorker (which runs automatically right after
    every file load) hasn't reached yet. In the common case — the user opens
    Fusion -> Manual Alignment after DC precompute has already finished —
    *angle_indices* is empty and this worker is never even constructed (see
    CorrelatePage._start_mask_prep).
    """
    angle_done = pyqtSignal(int, np.ndarray)   # angle_idx, dc4_mv

    def __init__(self, sras: SrasFile, angle_indices: list[int]):
        super().__init__()
        self._sras = sras
        self._angles = angle_indices
        self._budget = 0

    def _plan(self) -> int:
        n_workers, self._budget = compute.plan_angle_level(self._sras)
        return n_workers

    def _items(self):
        return self._angles

    def _one(self, a: int) -> tuple[int, np.ndarray]:
        return a, dc_image_mv(self._sras, a, CH4_IDX,
                              max_workers=1, budget=self._budget)

    def _emit(self, result):
        self.angle_done.emit(*result)


class CrossCorrelateWorker(_PooledWorker):
    """Rigid registration (rotation + translation, never scale) of each of
    *angle_indices* against *ref_angle_idx*, for the alignment wizard's
    Run/Re-run Correlation button.

    Runs on a background thread — registering a real many-angle,
    high-resolution scan takes long enough that doing it on the GUI thread
    would visibly freeze the dialog. Rotation is *searched*, not taken from the
    stage's reported angle: see compute.register_angle_to_reference, which
    seeds from that angle but scores both of its signs and refines from there.
    dc4_mv is the dialog's own already-in-memory per-angle CH4 image — this
    worker does no fetching of its own.
    """
    # angle_idx, rotation_deg, shift_x_mm, shift_y_mm, score, source
    angle_done = pyqtSignal(int, float, float, float, float, str)

    def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
                dc4_mv: dict[int, np.ndarray], *, reg_kwargs: dict | None = None):
        """*reg_kwargs* is splatted into register_angle_to_reference — every
        registration setting the wizard exposes (sources, threshold, search
        width, seed, signs, refine, grid sizes) travels in it, so this class
        holds no opinion about which knobs exist and exposing another needs no
        change here."""
        super().__init__()
        self._sras = sras
        self._ref = ref_angle_idx
        self._angles = angle_indices
        self._dc4_mv = dc4_mv
        self._reg_kwargs = dict(reg_kwargs or {})

    def _plan(self) -> int:
        return compute.registration_workers(self._sras)

    def _items(self):
        return self._angles

    def _one(self, a: int) -> tuple[int, compute.RigidFit]:
        return a, compute.register_angle_to_reference(
            self._sras, a, self._ref, self._dc4_mv, **self._reg_kwargs)

    def _emit(self, result):
        a, fit = result
        self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
                             fit.shift_mm[1], fit.score, fit.source)


class AlignedExportWorker(CancellableWorker):
    """Writes the aligned, cropped .sras on a background thread.

    Unlike every other worker here this one produces a *file*, which changes
    what cancellation has to mean: write_aligned_sras stages into a ".part"
    sibling and removes it when should_stop() fires, so a cancelled or crashed
    export leaves nothing behind. That matters more than it sounds — a
    truncated .sras is not detectably broken, since the v6 parser reads a short
    file as an aborted scan and opens it happily.

    Cancellation is polled per output row chunk, the same granularity
    CancellableWorker's docstring justifies, so closing the window never waits
    on a multi-gigabyte write.
    """
    progress = pyqtSignal(int)                  # 0-100
    finished = pyqtSignal(str, str)             # written path ("" = none), error

    def __init__(self, sras: SrasFile, result, out_path: str):
        super().__init__()
        self._sras = sras
        self._result = result
        self._out_path = out_path

    def run(self):
        try:
            written = write_aligned_sras(
                self._sras, self._result, self._out_path,
                progress_cb=self.progress.emit, should_stop=self._stopped)
            if self._stopped():
                self.finished.emit("", "")      # cancelled: no file, no error
            else:
                self.finished.emit(str(written), "")
        except Exception as exc:
            self.finished.emit("", str(exc))