from __future__ import annotations

from pathlib import Path
from typing import Any

import mlx.core as mx
import numpy as np

_VISION_AVAILABLE = False

try:
    import Cocoa
    import CoreImage
    import CoreVideo
    import Quartz
    import Vision

    _VISION_AVAILABLE = True
except ImportError:
    Vision = None
    Cocoa = None
    Quartz = None
    CoreImage = None
    CoreVideo = None


def is_vision_available() -> bool:
    return _VISION_AVAILABLE


def _require_vision() -> None:
    if not _VISION_AVAILABLE:
        raise ImportError(
            "Apple Vision framework features require macOS and the "
            "mlx-vision[vision] extra. Install with: pip install mlx-vision[vision]"
        )


def _coerce_uint8_hwc(image: mx.array) -> np.ndarray:
    array = np.asarray(image)

    if array.ndim == 2:
        array = array[:, :, np.newaxis]
    if array.ndim != 3:
        raise ValueError("Vision array inputs must have shape [H, W, C]")

    if array.dtype != np.uint8:
        array = np.clip(array, 0.0, 1.0)
        array = (array * 255.0).astype(np.uint8)

    if array.shape[2] not in (1, 3, 4):
        raise ValueError("Expected image with 1, 3, or 4 channels")

    return array


def _cfdata_from_bytes(data: bytes):
    return Cocoa.NSData.dataWithBytes_length_(data, len(data))


def _cgimage_from_array(image: mx.array):
    array = _coerce_uint8_hwc(image)
    height, width, channels = array.shape

    if channels == 1:
        array = np.repeat(array, 3, axis=2)
        channels = 3
    elif channels == 4:
        array = array[:, :, :3]
        channels = 3

    bytes_per_row = width * channels
    provider = Quartz.CGDataProviderCreateWithCFData(_cfdata_from_bytes(array.tobytes()))
    color_space = Quartz.CGColorSpaceCreateDeviceRGB()
    bitmap_info = Quartz.kCGImageAlphaNone | Quartz.kCGBitmapByteOrderDefault

    return Quartz.CGImageCreate(
        width,
        height,
        8,
        8 * channels,
        bytes_per_row,
        color_space,
        bitmap_info,
        provider,
        None,
        False,
        Quartz.kCGRenderingIntentDefault,
    )


def _make_handler(image: mx.array | str | Path):
    if isinstance(image, (str, Path)):
        url = Cocoa.NSURL.fileURLWithPath_(str(Path(image)))
        return Vision.VNImageRequestHandler.alloc().initWithURL_options_(url, {})

    cgimage = _cgimage_from_array(image)
    return Vision.VNImageRequestHandler.alloc().initWithCGImage_options_(cgimage, {})


def detect_faces(image: mx.array | str | Path) -> list[dict[str, Any]]:
    _require_vision()

    handler = _make_handler(image)
    request = Vision.VNDetectFaceRectanglesRequest.alloc().init()

    ok, err = handler.performRequests_error_([request], None)
    if not ok and err is not None:
        raise RuntimeError(str(err))

    results = []
    for observation in request.results() or []:
        bbox = observation.boundingBox()
        results.append(
            {
                "bbox": mx.array(
                    [
                        float(bbox.origin.x),
                        float(bbox.origin.y),
                        float(bbox.size.width),
                        float(bbox.size.height),
                    ]
                ),
                "confidence": float(observation.confidence()),
            }
        )
    return results


def detect_text(
    image: mx.array | str | Path,
    recognition_level: str = "accurate",
    languages: list[str] | None = None,
) -> list[dict[str, Any]]:
    _require_vision()

    if recognition_level not in {"accurate", "fast"}:
        raise ValueError('recognition_level must be "accurate" or "fast"')

    handler = _make_handler(image)
    request = Vision.VNRecognizeTextRequest.alloc().init()

    if recognition_level == "accurate":
        request.setRecognitionLevel_(Vision.VNRequestTextRecognitionLevelAccurate)
    else:
        request.setRecognitionLevel_(Vision.VNRequestTextRecognitionLevelFast)

    if languages:
        request.setRecognitionLanguages_(languages)

    ok, err = handler.performRequests_error_([request], None)
    if not ok and err is not None:
        raise RuntimeError(str(err))

    results = []
    for observation in request.results() or []:
        candidates = observation.topCandidates_(1)
        if not candidates:
            continue
        bbox = observation.boundingBox()
        results.append(
            {
                "text": str(candidates[0].string()),
                "confidence": float(candidates[0].confidence()),
                "bbox": mx.array(
                    [
                        float(bbox.origin.x),
                        float(bbox.origin.y),
                        float(bbox.size.width),
                        float(bbox.size.height),
                    ]
                ),
            }
        )
    return results


def classify_image(image: mx.array | str | Path) -> list[dict[str, Any]]:
    _require_vision()

    handler = _make_handler(image)
    request = Vision.VNClassifyImageRequest.alloc().init()

    ok, err = handler.performRequests_error_([request], None)
    if not ok and err is not None:
        raise RuntimeError(str(err))

    results = []
    for observation in request.results() or []:
        confidence = float(observation.confidence())
        if confidence > 0.01:
            results.append(
                {
                    "label": str(observation.identifier()),
                    "confidence": confidence,
                }
            )

    return sorted(results, key=lambda item: item["confidence"], reverse=True)


def saliency_map(image: mx.array | str | Path, kind: str = "attention") -> mx.array:
    _require_vision()

    if kind not in {"attention", "objectness"}:
        raise ValueError('kind must be "attention" or "objectness"')

    handler = _make_handler(image)
    if kind == "attention":
        request = Vision.VNGenerateAttentionBasedSaliencyImageRequest.alloc().init()
    else:
        request = Vision.VNGenerateObjectnessBasedSaliencyImageRequest.alloc().init()

    ok, err = handler.performRequests_error_([request], None)
    if not ok and err is not None:
        raise RuntimeError(str(err))

    results = request.results()
    if not results:
        raise ValueError("Saliency map generation returned no results")

    observation = results[0]
    pixel_buffer = observation.pixelBuffer()
    CoreVideo.CVPixelBufferLockBaseAddress(pixel_buffer, 0)

    try:
        base = CoreVideo.CVPixelBufferGetBaseAddress(pixel_buffer)
        height = CoreVideo.CVPixelBufferGetHeight(pixel_buffer)
        width = CoreVideo.CVPixelBufferGetWidth(pixel_buffer)
        row_bytes = CoreVideo.CVPixelBufferGetBytesPerRow(pixel_buffer)

        import ctypes

        row_floats = row_bytes // 4
        buffer_type = ctypes.c_float * (height * row_floats)
        buffer = buffer_type.from_address(base)
        array = np.ctypeslib.as_array(buffer).reshape(height, row_floats)[:, :width].copy()
    finally:
        CoreVideo.CVPixelBufferUnlockBaseAddress(pixel_buffer, 0)

    array = np.clip(array, 0.0, 1.0).astype(np.float32)
    return mx.array(array)


def attention_saliency_map(image: mx.array | str | Path) -> mx.array:
    return saliency_map(image, kind="attention")
