"""
QgsTask wrapper for running probability calculations in a background thread.

This allows the QGIS UI to remain responsive during long calculations
and provides proper progress reporting through the QGIS task manager.
"""
from typing import Any, Callable, Optional
from qgis.core import QgsTask, QgsMessageLog, Qgis
from qgis.PyQt.QtCore import pyqtSignal


class CalculationTask(QgsTask):
    """
    Background task for running drifting model calculations.

    This task runs the expensive probability hole calculations in a separate thread,
    allowing the QGIS UI to remain responsive. Progress updates are sent via signals
    and displayed in the QGIS task manager.

    Signals:
        progress_updated: Emitted with (completed, total, message) when progress changes
        calculation_finished: Emitted when calculation completes successfully
        calculation_failed: Emitted with error message if calculation fails
    """

    # Custom signals for progress and completion
    progress_updated = pyqtSignal(int, int, str)  # (completed, total, message)
    calculation_finished = pyqtSignal(object)  # (calc object with results)
    calculation_failed = pyqtSignal(str)  # (error message)

    def __init__(
        self,
        description: str,
        calc_object: Any,
        data: dict[str, Any],
    ):
        """
        Initialize the calculation task.

        Args:
            description: Readable description for the task
            calc_object: The Calculation object with run_drifting_model method
            data: Dictionary of data passed to run_drifting_model
        """
        super().__init__(description, QgsTask.Flag.CanCancel)
        self.calc = calc_object
        self.data = data
        self.exception: Optional[Exception] = None
        self.error_msg: Optional[str] = None

        # Store reference to progress callback
        self._progress_callback: Optional[Callable[[int, int, str], None]] = None

    def set_progress_callback(self, callback: Callable[[int, int, str], None]) -> None:
        """
        Set a callback function for progress updates.

        Args:
            callback: Function that takes (completed, total, message)
        """
        self._progress_callback = callback

    def _update_description(self, text: str) -> None:
        """Update the task description shown in the QGIS task manager."""
        self.setDescription(f"OMRAT: {text}")

    def _make_progress_wrapper(self) -> Callable[[int, int, str], bool]:
        def progress_wrapper(completed: int, total: int, message: str) -> bool:
            if self.isCanceled():
                return False
            if total > 0:
                self.setProgress(int((completed / total) * 100))
            self._update_description(message)
            self.progress_updated.emit(completed, total, message)
            return True
        return progress_wrapper

    def _run_consequence(self) -> None:
        self._update_description("Computing oil-spill consequences...")
        try:
            from compute.consequence import compute_catastrophe_exceedance
            consequence_block = self.data.get('consequence') or {}
            self.calc.consequence_result = compute_catastrophe_exceedance(
                consequence_block,
                drifting_report=getattr(self.calc, 'drifting_report', None),
                powered_grounding_report=getattr(self.calc, 'powered_grounding_report', None),
                powered_allision_report=getattr(self.calc, 'powered_allision_report', None),
                collision_report=getattr(self.calc, 'collision_report', None),
            )
        except Exception as exc:
            QgsMessageLog.logMessage(f'Consequence calculation failed: {exc}', 'OMRAT', Qgis.MessageLevel.Warning)
            self.calc.consequence_result = None

    def _run_all_phases(self) -> None:
        self._update_description("Drifting model - preparing...")
        self.calc.run_drifting_model(self.data)
        if self.isCanceled():
            return
        self._update_description("Ship-ship collisions...")
        self.setProgress(0)
        self.calc.run_ship_collision_model(self.data)
        if self.isCanceled():
            return
        self._update_description("Powered grounding...")
        self.setProgress(0)
        self.calc.run_powered_grounding_model(self.data)
        if self.isCanceled():
            return
        self._update_description("Powered allision...")
        self.setProgress(0)
        self.calc.run_powered_allision_model(self.data)
        if not self.isCanceled():
            self._run_consequence()

    def run(self) -> bool:
        """Execute the calculation in a background thread."""
        QgsMessageLog.logMessage('Starting calculation...', 'OMRAT', Qgis.MessageLevel.Info)
        try:
            from compute.data_preparation import apply_traffic_scaling
            apply_traffic_scaling(self.data)
            if hasattr(self.calc, 'set_progress_callback'):
                self.calc.set_progress_callback(self._make_progress_wrapper())
            self._run_all_phases()
            if self.isCanceled():
                QgsMessageLog.logMessage('Calculation was cancelled by user', 'OMRAT', Qgis.MessageLevel.Warning)
                return False
            self._update_description("Complete")
            self.setProgress(100)
            QgsMessageLog.logMessage('All calculations completed successfully', 'OMRAT', Qgis.MessageLevel.Success)
            return True
        except Exception as e:
            self.exception = e
            self.error_msg = str(e)
            QgsMessageLog.logMessage(
                f'Calculation failed with error: {self.error_msg}', 'OMRAT', Qgis.MessageLevel.Critical)
            return False

    def finished(self, result: bool) -> None:
        """
        Called when the task completes (success or failure).

        This method runs in the main thread and can safely interact with Qt widgets.

        Args:
            result: True if successful, False if failed or cancelled
        """
        if result:
            # Success - emit signal with calculation object
            self.calculation_finished.emit(self.calc)
            QgsMessageLog.logMessage(
                'Calculation task finished successfully',
                'OMRAT',
                Qgis.MessageLevel.Success
            )
        elif self.isCanceled():
            # Cancelled by user
            QgsMessageLog.logMessage(
                'Calculation task was cancelled',
                'OMRAT',
                Qgis.MessageLevel.Warning
            )
        else:
            # Failed with error
            error_msg = self.error_msg or "Unknown error"
            self.calculation_failed.emit(error_msg)
            QgsMessageLog.logMessage(
                f'Calculation task failed: {error_msg}',
                'OMRAT',
                Qgis.MessageLevel.Critical
            )

    def cancel(self) -> None:
        """
        Called when the user cancels the task.

        This sets the cancellation flag which is checked during the calculation.
        """
        QgsMessageLog.logMessage(
            'Cancelling calculation task...',
            'OMRAT',
            Qgis.MessageLevel.Warning
        )
        super().cancel()
