"""Export status tracking wrapper for the Matrice ActionTracker."""

from __future__ import annotations

import logging
import time
from typing import Any

logger = logging.getLogger(__name__)

# Step codes used by the Matrice platform for export actions.
# These match the codes already used across the existing export pipelines
# (ml_yolov8_detection, ml_hf_detection_mscoco, ml_pytorch_vision_*).
_STEP_PIPELINE_ACK = "MDL_EXP_ACK"       # Pipeline acknowledged
_STEP_PIPELINE_START = "MDL_EXP_STR"     # Pipeline started
_STEP_PIPELINE_COMPLETE = "MDL_EXP_CMPL" # Pipeline completed

_STEP_FORMAT_ACK = "MDL_EXPT_ACK"        # Per-format export acknowledged
_STEP_FORMAT_RESULT = "MDL_EXPT"         # Per-format export result


class ExportTracker:
    """Wraps a Matrice :class:`ActionTracker` for export status reporting.

    If no *action_tracker* is provided, all reporting methods become no-ops so
    callers never need to guard against ``None``.

    Parameters
    ----------
    action_tracker:
        An instance of :class:`matrice.action_tracker.ActionTracker` (or any
        object that exposes an ``update_status(stepCode, status, description)``
        method).  ``None`` is acceptable -- all calls will be silently skipped.
    """

    def __init__(self, action_tracker: Any | None = None) -> None:
        self.tracker = action_tracker
        self._start_time: float | None = None
        self._format_start_times: dict[str, float] = {}

    # ------------------------------------------------------------------ #
    # Pipeline-level reporting
    # ------------------------------------------------------------------ #

    def report_pipeline_start(self, formats: list[str]) -> None:
        """Report that the export pipeline has been acknowledged and started.

        Sends two status updates:

        1. ``MDL_EXP_ACK`` (OK) -- pipeline acknowledged.
        2. ``MDL_EXP_STR`` (OK) -- pipeline started.

        Parameters
        ----------
        formats:
            The list of target format names that will be exported.
        """
        self._start_time = time.time()
        description_list = ", ".join(formats)

        if self.tracker is None:
            logger.info("Export pipeline starting for formats: %s", description_list)
            return

        try:
            self.tracker.update_status(
                _STEP_PIPELINE_ACK,
                "OK",
                f"Model Export Acknowledged (formats: {description_list})",
            )
            self.tracker.update_status(
                _STEP_PIPELINE_START,
                "OK",
                f"Model Export Started (formats: {description_list})",
            )
        except Exception as exc:
            logger.warning("Failed to report pipeline start: %s", exc)

    def report_pipeline_complete(self, results: dict[str, dict[str, Any]]) -> None:
        """Report that the export pipeline has finished.

        Computes a summary of successes / failures and sends
        ``MDL_EXP_CMPL`` with status ``SUCCESS`` if at least one format
        succeeded, or ``ERROR`` if all failed.

        Parameters
        ----------
        results:
            The results dict returned by :meth:`ExportPipeline.export`.
        """
        elapsed = ""
        if self._start_time is not None:
            elapsed = f" in {time.time() - self._start_time:.1f}s"

        successes = [f for f, r in results.items() if r.get("status") == "success"]
        failures = [f for f, r in results.items() if r.get("status") != "success"]

        if successes:
            status = "SUCCESS"
            desc = f"Model Export Completed{elapsed} ({len(successes)} succeeded"
            if failures:
                desc += f", {len(failures)} failed: {', '.join(failures)}"
            desc += ")"
        else:
            status = "ERROR"
            desc = f"Model Export Failed{elapsed} (all {len(failures)} format(s) failed)"

        if self.tracker is None:
            logger.info(desc)
            return

        try:
            self.tracker.update_status(_STEP_PIPELINE_COMPLETE, status, desc)
        except Exception as exc:
            logger.warning("Failed to report pipeline complete: %s", exc)

    # ------------------------------------------------------------------ #
    # Per-format reporting
    # ------------------------------------------------------------------ #

    def report_export_start(self, fmt: str) -> None:
        """Report that export of a specific format has been acknowledged.

        Parameters
        ----------
        fmt:
            The canonical format name (e.g. ``"onnx"``).
        """
        self._format_start_times[fmt] = time.time()

        if self.tracker is None:
            logger.info("Export started for format: %s", fmt)
            return

        try:
            self.tracker.update_status(
                _STEP_FORMAT_ACK,
                "OK",
                f"{fmt} export acknowledged",
            )
        except Exception as exc:
            logger.warning("Failed to report export start for %s: %s", fmt, exc)

    def report_export_success(
        self,
        fmt: str,
        path: str,
        validation: dict[str, Any] | None = None,
    ) -> None:
        """Report that a format was exported successfully.

        Parameters
        ----------
        fmt:
            The canonical format name.
        path:
            File-system path to the exported artefact.
        validation:
            Optional validation results dict.
        """
        elapsed = ""
        start = self._format_start_times.pop(fmt, None)
        if start is not None:
            elapsed = f" in {time.time() - start:.1f}s"

        desc = f"{fmt} export successful{elapsed} -> {path}"
        if validation:
            desc += f" (validation: {validation.get('status', 'n/a')})"

        if self.tracker is None:
            logger.info(desc)
            return

        try:
            self.tracker.update_status(_STEP_FORMAT_RESULT, "SUCCESS", desc)
        except Exception as exc:
            logger.warning("Failed to report export success for %s: %s", fmt, exc)

    def report_export_failure(self, fmt: str, error: str) -> None:
        """Report that a format export failed.

        Parameters
        ----------
        fmt:
            The canonical format name.
        error:
            A human-readable error description.
        """
        self._format_start_times.pop(fmt, None)
        desc = f"{fmt} export failed: {error}"

        if self.tracker is None:
            logger.error(desc)
            return

        try:
            self.tracker.update_status(_STEP_FORMAT_RESULT, "ERROR", desc)
        except Exception as exc:
            logger.warning("Failed to report export failure for %s: %s", fmt, exc)
