from __future__ import annotations

import json
import os
from typing import Any, Dict, Tuple

import cv2
from loguru import logger

from datamate.core.base_op import Mapper


class test_annotation_marker(Mapper):
    def __init__(self, *args, **kwargs):
        super(test_annotation_marker, self).__init__(*args, **kwargs)
        self._marker_text = str(kwargs.get("markerText", "TEST_ANNOTATION"))
        self._marker_color = self._parse_color(kwargs.get("markerColor", "0,255,0"))
        self._marker_thickness = int(kwargs.get("markerThickness", 2) or 2)
        self._output_dir = kwargs.get("outputDir")

    @staticmethod
    def _parse_color(value: Any) -> Tuple[int, int, int]:
        if isinstance(value, (list, tuple)) and len(value) >= 3:
            try:
                bgr = tuple(int(max(min(float(item), 255), 0)) for item in value[:3])
                return bgr  # type: ignore[return-value]
            except Exception:
                return 0, 255, 0

        if isinstance(value, str):
            parts = [part.strip() for part in value.split(",")]
            if len(parts) >= 3:
                try:
                    bgr = tuple(int(max(min(float(item), 255), 0)) for item in parts[:3])
                    return bgr  # type: ignore[return-value]
                except Exception:
                    return 0, 255, 0

        return 0, 255, 0

    def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        image_path = sample.get(self.image_key) or sample.get("image")
        if not image_path or not os.path.exists(image_path):
            logger.warning("test_annotation_marker: image not found: {}", image_path)
            return sample

        image = cv2.imread(image_path)
        if image is None:
            logger.warning("test_annotation_marker: failed to read image: {}", image_path)
            return sample

        image_height, image_width = image.shape[:2]

        margin = max(min(image_width, image_height) // 20, 10)
        x1, y1 = margin, margin
        x2, y2 = image_width - margin, image_height - margin

        cv2.rectangle(
            image,
            (x1, y1),
            (x2, y2),
            self._marker_color,
            self._marker_thickness,
        )

        cv2.putText(
            image,
            self._marker_text,
            (x1, max(y1 - 10, 20)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.8,
            self._marker_color,
            max(self._marker_thickness, 1),
            cv2.LINE_AA,
        )

        if self._output_dir and os.path.exists(self._output_dir):
            output_dir = self._output_dir
        else:
            output_dir = os.path.dirname(image_path)

        images_dir = os.path.join(output_dir, "images")
        annotations_dir = os.path.join(output_dir, "annotations")
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(annotations_dir, exist_ok=True)

        base_name = os.path.basename(image_path)
        name_without_ext = os.path.splitext(base_name)[0]

        output_image_path = os.path.join(images_dir, base_name)
        output_json_path = os.path.join(annotations_dir, f"{name_without_ext}.json")

        cv2.imwrite(output_image_path, image)

        annotations = {
            "image": base_name,
            "width": image_width,
            "height": image_height,
            "marker": {
                "text": self._marker_text,
                "color_bgr": list(self._marker_color),
                "thickness": self._marker_thickness,
            },
            "detections": [
                {
                    "label": self._marker_text,
                    "class_id": -1,
                    "confidence": 1.0,
                    "bbox_xyxy": [x1, y1, x2, y2],
                    "bbox_xywh": [x1, y1, x2 - x1, y2 - y1],
                }
            ],
        }

        with open(output_json_path, "w", encoding="utf-8") as file:
            json.dump(annotations, file, indent=2, ensure_ascii=False)

        sample["output_image"] = output_image_path
        sample["annotations_file"] = output_json_path
        sample["annotations"] = annotations
        sample["detection_count"] = 1
        return sample