#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   formats.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Dataset export format validation utilities.
"""

from vi.api.resources.datasets.types import DatasetExportFormat, DatasetType

# Full dataset export (with assets) - ONLY ViFull is supported
DATASET_EXPORT_FORMATS: set[DatasetExportFormat] = {
    DatasetExportFormat.VI_FULL,
}

PHRASE_GROUNDING_ANNOTATION_FORMATS: set[DatasetExportFormat] = {
    DatasetExportFormat.CSV_FOUR_CORNER,
    DatasetExportFormat.CSV_WIDTH_HEIGHT,
    DatasetExportFormat.COCO,
    DatasetExportFormat.PASCAL_VOC,
    DatasetExportFormat.YOLO_DARKNET,
    DatasetExportFormat.YOLO_KERAS_PYTORCH,
    DatasetExportFormat.VI_JSONL,
    DatasetExportFormat.VI_TFRECORD,
}

VQA_ANNOTATION_FORMATS: set[DatasetExportFormat] = {
    DatasetExportFormat.VI_JSONL,
    DatasetExportFormat.VI_TFRECORD,
}

FREEFORM_ANNOTATION_FORMATS: set[DatasetExportFormat] = {
    DatasetExportFormat.VI_JSONL,
}


def get_supported_formats(
    dataset_type: DatasetType,
    annotations_only: bool = False,
) -> set[DatasetExportFormat]:
    """Get supported export formats for a dataset type and export mode.

    Args:
        dataset_type: The type of dataset (PHRASE_GROUNDING, VQA, or FREEFORM).
        annotations_only: If True, returns annotation-only formats.
            If False, returns full dataset export formats (only ViFull).

    Returns:
        Set of supported DatasetExportFormat values for the given configuration.

    Example:
        >>> # Full dataset export - only ViFull
        >>> formats = get_supported_formats(
        ...     DatasetType.PHRASE_GROUNDING, annotations_only=False
        ... )
        >>> DatasetExportFormat.VI_FULL in formats
        True
        >>> len(formats)
        1

        >>> # Annotation-only for phrase grounding - multiple formats
        >>> formats = get_supported_formats(
        ...     DatasetType.PHRASE_GROUNDING, annotations_only=True
        ... )
        >>> DatasetExportFormat.COCO in formats
        True
        >>> DatasetExportFormat.VI_FULL in formats
        False

        >>> # Annotation-only for VQA - limited formats
        >>> formats = get_supported_formats(DatasetType.VQA, annotations_only=True)
        >>> DatasetExportFormat.COCO in formats
        False
        >>> DatasetExportFormat.VI_JSONL in formats
        True

    """
    if annotations_only:
        if dataset_type == DatasetType.PHRASE_GROUNDING:
            return PHRASE_GROUNDING_ANNOTATION_FORMATS
        if dataset_type == DatasetType.VQA:
            return VQA_ANNOTATION_FORMATS
        if dataset_type == DatasetType.FREEFORM:
            return FREEFORM_ANNOTATION_FORMATS
    else:
        # Full dataset export - ONLY ViFull is supported
        return DATASET_EXPORT_FORMATS


def validate_export_format(
    export_format: DatasetExportFormat,
    dataset_type: DatasetType,
    annotations_only: bool = False,
) -> tuple[bool, str | None]:
    """Validate if an export format is compatible with a dataset type.

    Args:
        export_format: The export format to validate.
        dataset_type: The type of dataset (PHRASE_GROUNDING, VQA, or FREEFORM).
        annotations_only: If True, validates for annotation-only export.
            If False, validates for full dataset export.

    Returns:
        Tuple of (is_valid, error_message). If valid, error_message is None.
        If invalid, error_message contains a helpful description.

    Example:
        >>> is_valid, error = validate_export_format(
        ...     DatasetExportFormat.COCO, DatasetType.VQA, annotations_only=True
        ... )
        >>> is_valid
        False
        >>> "VQA datasets only support" in error
        True

    """
    supported_formats = get_supported_formats(dataset_type, annotations_only)

    if export_format in supported_formats:
        return True, None

    # Generate helpful error message
    if annotations_only:
        # Annotation-only export
        if export_format == DatasetExportFormat.VI_FULL:
            return False, (
                f"Format '{export_format.value}' is only supported for full dataset exports. "
                f"For annotation-only exports, choose from: {', '.join(sorted([f.value for f in supported_formats]))}."
            )

        if dataset_type == DatasetType.VQA:
            return False, (
                f"Format '{export_format.value}' is not supported for VQA annotation exports. "
                f"VQA datasets only support ViJsonl and ViTfrecord formats for annotation-only exports."
            )

        if dataset_type == DatasetType.FREEFORM:
            return False, (
                f"Format '{export_format.value}' is not supported for freeform annotation exports. "
                f"Freeform datasets only support ViJsonl format for annotation-only exports."
            )

        return False, (
            f"Format '{export_format.value}' is not supported for phrase grounding annotation exports."
        )
    else:
        # Full dataset export - only VI_FULL is valid
        return False, (
            f"Format '{export_format.value}' is only supported for annotation-only exports. "
            f"For full dataset exports, only ViFull format is supported."
        )


def get_format_suggestion(
    dataset_type: DatasetType,
    annotations_only: bool = False,
) -> str:
    """Get a formatted string of supported formats for error messages.

    Args:
        dataset_type: The type of dataset (PHRASE_GROUNDING, VQA, or FREEFORM).
        annotations_only: If True, returns annotation-only format suggestions.

    Returns:
        Formatted string listing supported formats.

    Example:
        >>> suggestion = get_format_suggestion(DatasetType.VQA, annotations_only=True)
        >>> "ViJsonl" in suggestion
        True

    """
    supported = get_supported_formats(dataset_type, annotations_only)
    format_names = sorted([f.value for f in supported])

    if annotations_only:
        prefix = f"annotation-only exports for {dataset_type.value} datasets"
    else:
        prefix = "full dataset exports"

    return f"Supported format(s) for {prefix}: {', '.join(format_names)}"
