"""
Mask Exporter

Exports segmentation mask annotations as PNG binary images.
Each label gets a separate PNG where filled pixels are the label color
and background is transparent.

Requires: numpy and Pillow (PIL)
"""

import os
import logging
from typing import Optional, Tuple, List

from .base import BaseExporter, ExportContext, ExportResult
from .cv_utils import (
    extract_image_annotations,
    get_image_dimensions,
    get_image_filename,
    build_category_mapping,
    decode_rle,
)

logger = logging.getLogger(__name__)


class MaskExporter(BaseExporter):
    format_name = "mask_png"
    description = "Segmentation masks as PNG images (requires Pillow)"
    file_extensions = [".png"]

    def can_export(self, context: ExportContext) -> Tuple[bool, str]:
        # Check for Pillow
        try:
            from PIL import Image
        except ImportError:
            return False, "Pillow (PIL) is required for mask export. Install with: pip install Pillow"

        has_image_schema = any(
            s.get("annotation_type") == "image_annotation"
            for s in context.schemas
        )
        if not has_image_schema:
            return False, "No image_annotation schema found in config"

        return True, ""

    def export(self, context: ExportContext, output_path: str,
               options: Optional[dict] = None) -> ExportResult:
        from PIL import Image

        options = options or {}
        warnings = []
        files_written = []

        os.makedirs(output_path, exist_ok=True)
        category_map = build_category_mapping(context.annotations, context.schemas)

        # Assign colors to categories
        default_colors = [
            (255, 0, 0), (0, 255, 0), (0, 0, 255),
            (255, 255, 0), (255, 0, 255), (0, 255, 255),
            (128, 0, 0), (0, 128, 0), (0, 0, 128),
            (128, 128, 0),
        ]
        category_colors = {}
        for name, idx in category_map.items():
            category_colors[name] = default_colors[idx % len(default_colors)]

        masks_exported = 0

        for ann in context.annotations:
            instance_id = ann.get("instance_id", "")
            item = context.items.get(instance_id, {})
            img_anns = extract_image_annotations(ann)
            if not img_anns:
                continue

            width, height = get_image_dimensions(item)
            if width <= 0 or height <= 0:
                # Try to get from mask RLE size
                for _, objects in img_anns:
                    for obj in objects:
                        if obj.get("type") == "mask" and "rle" in obj:
                            size = obj["rle"].get("size", [])
                            if len(size) == 2:
                                height, width = size
                                break
                    if width > 0:
                        break

            if width <= 0 or height <= 0:
                warnings.append(f"No dimensions for {instance_id}, skipping masks")
                continue

            file_name = get_image_filename(item) or instance_id
            raw_stem = os.path.splitext(os.path.basename(file_name))[0]
            stem = "".join(c if c.isalnum() or c in "-_." else "_" for c in raw_stem)

            for schema_name, objects in img_anns:
                for obj in objects:
                    if obj.get("type") != "mask":
                        continue

                    label = obj.get("label", "unknown")
                    rle = obj.get("rle", {})
                    if not rle.get("counts"):
                        continue

                    mask_data = decode_rle(rle, width, height)
                    color = category_colors.get(label, (255, 255, 255))

                    # Create RGBA image
                    img = Image.new("RGBA", (width, height), (0, 0, 0, 0))
                    pixels = img.load()

                    for i, val in enumerate(mask_data):
                        if val:
                            y = i // width
                            x = i % width
                            if x < width and y < height:
                                pixels[x, y] = (color[0], color[1], color[2], 200)

                    safe_label = "".join(c if c.isalnum() or c in "-_." else "_" for c in label)
                    mask_file = os.path.join(output_path, f"{stem}_{safe_label}_mask.png")
                    img.save(mask_file)
                    files_written.append(mask_file)
                    masks_exported += 1

        return ExportResult(
            success=True,
            format_name=self.format_name,
            files_written=files_written,
            warnings=warnings,
            stats={"num_masks": masks_exported},
        )