# -*- coding: utf-8 -*-
"""Triton gRPC 推理客户端，封装 infer(mode, image) / infer_all(image)。"""
import os
import tempfile
from typing import Dict, Optional, Union

import numpy as np
import tritonclient.grpc as grpcclient

from gshield_pic_sdk._image import to_rgb_ndarray
from gshield_pic_sdk.exceptions import GshieldConnectionError, GshieldInferError, GshieldValueError
from gshield_pic_sdk.grpc_config import get_grpc_config, MODES_ALL
from gshield_pic_sdk.grpc.inference import (
    preprocess_image_raw,
    preprocess_numpy,
    run_task_with_client,
)

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


def _default_label_dir() -> str:
    pkg_dir = os.path.dirname(os.path.abspath(__file__))
    labels = os.path.join(pkg_dir, "labels")
    return labels if os.path.isdir(labels) else ""


class GrpcInferClient:
    """Triton gRPC 推理：按 mode 调用，支持 path/bytes/ndarray。"""

    def __init__(
        self,
        triton_url: str,
        label_dir: Optional[str] = None,
        grpc_config: Optional[Dict] = None,
        timeout: int = 30,
    ):
        self.triton_url = triton_url
        self.label_dir = label_dir or _default_label_dir()
        if not self.label_dir or not os.path.isdir(self.label_dir):
            raise GshieldValueError(
                "label_dir must be provided and point to a directory with label .txt files"
            )
        self._config = grpc_config or get_grpc_config()
        self.timeout = timeout

    def _ensure_path(self, image: ImageInput) -> str:
        """将 image 转为本地路径，供内部 inference 使用。"""
        if isinstance(image, str) and os.path.isfile(image):
            return image
        rgb = to_rgb_ndarray(image)
        # 写入临时文件（inference 里 preprocess_image_raw 和 infer_flag 需要路径读图）
        fd, path = tempfile.mkstemp(suffix=".jpg")
        try:
            import cv2
            bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
            ok, buf = cv2.imencode(".jpg", bgr)
            if not ok:
                raise GshieldValueError("failed to encode image")
            os.write(fd, buf.tobytes())
        finally:
            os.close(fd)
        return path

    def infer(self, mode: str, image: ImageInput) -> str:
        """
        单模型推理。
        :param mode: 模型名（与 grpc_config 一致，如 fire, politics, porn, flag）
        :param image: 路径、bytes 或 ndarray
        :return: summary_label 字符串
        """
        modes_list = self._config.get("modes_list", [])
        if mode != "all" and mode not in modes_list:
            raise GshieldValueError(f"unknown mode: {mode}, supported: {modes_list}")
        path = self._ensure_path(image)
        try:
            raw = preprocess_image_raw(path)
            input_224 = preprocess_numpy(raw, target_size=(224, 224))
            input_640 = preprocess_numpy(raw, target_size=(640, 640))
        except Exception as e:
            raise GshieldInferError(f"preprocess failed: {e}") from e
        try:
            client = grpcclient.InferenceServerClient(url=self.triton_url)
            if not client.is_server_live():
                raise GshieldConnectionError(f"Triton server not live: {self.triton_url}")
        except Exception as e:
            raise GshieldConnectionError(f"Triton connection failed: {e}") from e
        created_temp = not (isinstance(image, str) and os.path.isfile(image))
        temp_path = path if created_temp else None
        try:
            summary_label, _ = run_task_with_client(
                client,
                mode,
                path,
                input_224,
                input_640,
                self.label_dir,
                self._config["label_map"],
                self._config["data_types"],
            )
            return summary_label
        except Exception as e:
            raise GshieldInferError(f"infer failed: {e}") from e
        finally:
            if temp_path and os.path.isfile(temp_path):
                try:
                    os.remove(temp_path)
                except OSError:
                    pass

    def infer_all(self, image: ImageInput) -> Dict[str, str]:
        """对 MODES_ALL 中每个 mode 调用 infer，返回 {mode: label}。"""
        modes_all = self._config.get("modes_all", MODES_ALL)
        result = {}
        for m in modes_all:
            result[m] = self.infer(m, image)
        return result
