"""Save operation for writing images to disk."""

from pathlib import Path
from typing import Any

import cv2
import numpy as np
from numpy.typing import NDArray

from axm.operations import OperationProtocol, OperationResult


class SaveOperation(OperationProtocol):
    """Save an image to disk.

    Context:
        output_path: Path to save the image.

    Config:
        format: Output format (default: inferred from path).
        quality: JPEG quality 1-100 (default: 95).
    """

    @property
    def name(self) -> str:
        """Operation name."""
        return "image.save"

    def execute(
        self,
        data: NDArray[np.uint8] | None,
        context: dict[str, Any],
        config: dict[str, Any],
    ) -> OperationResult:
        """Save image to output path.

        Args:
            data: Image array to save (RGB format).
            context: Must contain 'output_path' key.
            config: Optional format and quality settings.

        Returns:
            OperationResult with same image (for chaining).
        """
        if data is None:
            return OperationResult(
                data=None,
                success=False,
                error="No input image provided",
            )

        output_path = Path(context.get("output_path", ""))
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Convert RGB to BGR for OpenCV
        image_bgr = cv2.cvtColor(data, cv2.COLOR_RGB2BGR)

        # Get format-specific params
        quality = config.get("quality", 95)
        ext = output_path.suffix.lower()

        params: list[int] = []
        if ext in [".jpg", ".jpeg"]:
            params = [cv2.IMWRITE_JPEG_QUALITY, quality]
        elif ext == ".png":
            params = [cv2.IMWRITE_PNG_COMPRESSION, 3]

        success = cv2.imwrite(str(output_path), image_bgr, params)

        if not success:
            return OperationResult(
                data=data,
                success=False,
                error=f"Failed to save image to {output_path}",
            )

        return OperationResult(
            data=data,
            success=True,
            metadata={"saved_to": str(output_path)},
        )
