"""Acquire fresh object-calibration sweeps through the native preprocessor."""

from __future__ import annotations

import time

from PyQt6.QtCore import QObject, QTimer, pyqtSignal

from python_app.gui.controllers.pipeline_controller import PipelineController
from python_app.gui.object_calibration_dialog import ObjectCalibrationDialog
from python_app.gui.runtime.pipeline_pause import PipelinePause
from python_app.models.collection_identity import collection_key
from python_app.models.dataset_model import SweepCollection
from python_app.models.run_config_schema import RunConfigModel


class ObjectSweepCapture(QObject):
    """Hold a pipeline pause while collecting a fresh matrix without object correction."""

    finished = pyqtSignal()
    TIMEOUT_S = 300

    def __init__(self, pipeline: PipelineController, dialog: ObjectCalibrationDialog,
                 config: RunConfigModel, count: int, provenance: dict) -> None:
        super().__init__(dialog)
        self._pipeline, self._dialog = pipeline, dialog
        self._config, self._count, self._provenance = config, count, provenance
        self._timer = QTimer(self)
        self._timer.timeout.connect(self._poll)
        self._done = self._started = self._starting = False
        self._pending_error: str | None = None
        self._pause: PipelinePause | None = None
        self._collections: list[SweepCollection] = []
        self._seen: set[tuple[int, int]] = set()
        self._stop_generation = pipeline.stop_generation
        self._start_ns = 0
        dialog.set_capture_pending(True)
        QTimer.singleShot(0, self._start)

    def _start(self) -> None:
        if self._done:
            return
        self._starting = True
        try:
            if not self._pipeline.can_resume(self._stop_generation):
                self.cancel()
                return
            self._pause = self._pipeline.pause()
            if self._pending_error is not None:
                return
            self._start_ns = time.monotonic_ns()
            self._started = self._pipeline.start(config_override=self._config, pause=self._pause)
            if self._pending_error is not None:
                return
            if not self._started or not self._pipeline.acquisition_healthy:
                raise RuntimeError("Calibration acquisition did not start; check the preprocessing sets and log")
            self._dialog.set_status(f"Taking calibration data: 0 / {self._count} sweeps")
            self._timer.start(50)
        except Exception as exc:
            self._pending_error = str(exc)
        finally:
            self._starting = False
            if self._pending_error is not None:
                self._finish(error=self._pending_error)

    def _poll(self) -> None:
        if self._done:
            return
        try:
            if self._pause is None or not self._pause.valid:
                self.cancel()
                return
            if not self._pipeline.acquisition_healthy:
                raise RuntimeError("Calibration acquisition stopped before all sweeps were received")
            for collection in self._pipeline.history.preprocessed:
                identity = collection_key(collection)
                if collection.capture_start_ns < self._start_ns or identity in self._seen:
                    continue
                self._seen.add(identity)
                self._collections.append(collection)
                if len(self._collections) == self._count:
                    self._finish()
                    return
            self._dialog.set_status(f"Taking calibration data: {len(self._collections)} / {self._count} sweeps")
            if time.monotonic_ns() - self._start_ns > self.TIMEOUT_S * 1e9:
                raise TimeoutError("Timed out waiting for preprocessed calibration sweeps")
        except Exception as exc:
            self._finish(error=str(exc))

    def cancel(self) -> None:
        self._finish(error="Calibration capture cancelled")

    def _finish(self, *, error: str | None = None) -> None:
        if self._done:
            return
        if self._starting:
            # Stop/start drains dispatch Qt events. Keep the lease until startup
            # returns, so Cancel cannot start a second run inside the first one.
            self._pending_error = error or "Calibration capture cancelled"
            return
        self._done = True
        self._timer.stop()
        try:
            if self._started:
                self._pipeline.stop_pipeline(keep_processor=True)
            if self._pause is not None and not self._pause.valid:
                error = "Calibration capture cancelled"
            if error is None:
                self._dialog.add_capture(self._collections, self._config.to_dict(), self._provenance)
            else:
                self._dialog.set_status(error)
        except Exception as exc:
            self._dialog.set_status(f"Calibration capture failed: {exc}")
        finally:
            self._dialog.set_capture_pending(False)
            # The owner clears its job before finish() can start normal acquisition.
            self.finished.emit()
            try:
                if self._pause is not None:
                    self._pause.finish()
            finally:
                self.deleteLater()