"""Typed containers and coercion helpers for inference outputs."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .exceptions import InferenceRunnerError


@dataclass
class ClassificationInferenceOutput:
    """Standard output container for classification-style runners."""

    predictions: Any
    outputs: Any
    targets: Any
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class DetectionInferenceOutput:
    """Standard output container for detection/segmentation-style runners."""

    outputs: Any
    targets: Any
    predictions: Any | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


def coerce_classification_output(raw: Any) -> ClassificationInferenceOutput:
    """Convert raw runner output to ``ClassificationInferenceOutput``.

    Args:
        raw: Runner output object or tuple-like result.

    Returns:
        Normalized ``ClassificationInferenceOutput`` instance.
    """
    if isinstance(raw, ClassificationInferenceOutput):
        return raw
    if isinstance(raw, tuple) and len(raw) >= 3:
        return ClassificationInferenceOutput(
            predictions=raw[0],
            outputs=raw[1],
            targets=raw[2],
            metadata={"raw_extra": raw[3:]},
        )
    raise InferenceRunnerError(
        "Classification runner must return (predictions, outputs, targets) "
        "or ClassificationInferenceOutput."
    )


def coerce_detection_output(raw: Any) -> DetectionInferenceOutput:
    """Convert raw runner output to ``DetectionInferenceOutput``.

    Args:
        raw: Runner output object or tuple-like result.

    Returns:
        Normalized ``DetectionInferenceOutput`` instance.
    """
    if isinstance(raw, DetectionInferenceOutput):
        return raw
    if isinstance(raw, tuple) and len(raw) >= 2:
        if len(raw) == 2:
            return DetectionInferenceOutput(outputs=raw[0], targets=raw[1])
        return DetectionInferenceOutput(
            predictions=raw[0],
            outputs=raw[1],
            targets=raw[2],
            metadata={"raw_extra": raw[3:]},
        )
    raise InferenceRunnerError(
        "Detection runner must return (outputs, targets) or "
        "(predictions, outputs, targets) or DetectionInferenceOutput."
    )

