# -*- coding: utf-8 -*-
"""图片入参统一：path/bytes/ndarray 转为 HTTP 上传用 bytes 或 gRPC 用 RGB ndarray。"""
import os
from typing import Union

import cv2
import numpy as np

ImageInput = Union[str, bytes, np.ndarray]


def to_bytes(image: ImageInput) -> bytes:
    """
    将图片转为 HTTP 上传用的 bytes。
    :param image: 文件路径(str)、bytes、或 HWC BGR/RGB ndarray (uint8)
    :return: JPEG 编码的 bytes
    """
    if isinstance(image, str):
        if not os.path.isfile(image):
            raise ValueError(f"image path not found: {image}")
        with open(image, "rb") as f:
            data = f.read()
        if not data:
            raise ValueError("image file is empty")
        return data
    if isinstance(image, bytes):
        if not image:
            raise ValueError("image bytes is empty")
        return image
    if isinstance(image, np.ndarray):
        if image.size == 0:
            raise ValueError("image ndarray is empty")
        if image.ndim != 3 or image.shape[-1] != 3:
            raise ValueError("image ndarray must be HWC with 3 channels")
        # cv2.imencode 期望 BGR；若已是 RGB 则先转 BGR 再编码
        if _is_likely_rgb(image):
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        ok, buf = cv2.imencode(".jpg", image)
        if not ok:
            raise ValueError("failed to encode ndarray as jpeg")
        return buf.tobytes()
    raise ValueError(f"image must be str, bytes or ndarray, got {type(image)}")


def _is_likely_rgb(arr: np.ndarray) -> bool:
    """启发式：若均值 R > B 则视为 RGB（通常自然图 R>=B）。"""
    if arr.shape[-1] != 3:
        return False
    r, g, b = arr[..., 0].mean(), arr[..., 1].mean(), arr[..., 2].mean()
    return r >= b


def to_rgb_ndarray(image: ImageInput) -> np.ndarray:
    """
    将图片转为 HWC RGB uint8 ndarray，供 gRPC 预处理使用。
    :param image: 文件路径(str)、bytes、或 HWC BGR/RGB ndarray
    :return: HWC RGB uint8, contiguous
    """
    if isinstance(image, str):
        if not os.path.isfile(image):
            raise ValueError(f"image path not found: {image}")
        img = cv2.imread(image)
        if img is None:
            raise ValueError(f"failed to read image: {image}")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return np.ascontiguousarray(img)
    if isinstance(image, bytes):
        if not image:
            raise ValueError("image bytes is empty")
        buf = np.frombuffer(image, dtype=np.uint8)
        img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError("failed to decode image bytes")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return np.ascontiguousarray(img)
    if isinstance(image, np.ndarray):
        if image.size == 0 or image.ndim != 3 or image.shape[-1] != 3:
            raise ValueError("image ndarray must be HWC with 3 channels")
        # 约定：ndarray 按 OpenCV 惯例视为 BGR，统一转为 RGB
        out = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        out = np.ascontiguousarray(out)
        if out.dtype != np.uint8:
            out = out.astype(np.uint8)
        return out
    raise ValueError(f"image must be str, bytes or ndarray, got {type(image)}")
