import numpy as np
import tritonclient.grpc as grpcclient
import cv2
import os
import time


def _to_summary_label(label):
    """Return 'normal' if label is None or str.strip().lower() == 'normal', else return the label string."""
    if label is None:
        return "normal"
    s = (label if isinstance(label, str) else str(label)).strip()
    return "normal" if s.lower() == "normal" else (label if isinstance(label, str) else str(label))


def load_labels(label_path):
    """从文件加载标签列表"""
    if not os.path.exists(label_path):
        return None
    try:
        with open(label_path, "r", encoding="utf-8") as f:
            labels = [line.strip() for line in f.readlines() if line.strip()]
        return labels
    except Exception:
        return None


def letterbox_resize_opencv(img, target_size=224, fill_value=128):
    """
    OpenCV implementation of letterbox resize to match PIL implementation in verify_onnx.py
    """
    h, w = img.shape[:2]
    scale = min(target_size / w, target_size / h)
    new_w, new_h = int(w * scale), int(h * scale)

    # Resize
    img_resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_CUBIC)

    # Create canvas
    canvas = np.full((target_size, target_size, 3), fill_value, dtype=np.uint8)

    # Paste centered
    x_offset = (target_size - new_w) // 2
    y_offset = (target_size - new_h) // 2

    canvas[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = img_resized

    return canvas


def preprocess_flag_classification(img_input, target_size=224):
    """
    针对 Flag Classification 的特殊预处理
    """
    if isinstance(img_input, str):
        if not os.path.exists(img_input):
            raise ValueError(f"文件不存在: {img_input}")
        img = cv2.imread(img_input)
        if img is None:
            raise ValueError(f"无法读取图片: {img_input}")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    elif isinstance(img_input, np.ndarray):
        # Assume input is BGR (standard OpenCV) since we usually pass cv2.imread or crops from it
        img = cv2.cvtColor(img_input, cv2.COLOR_BGR2RGB)
    else:
        raise ValueError("Invalid input type")

    # 1. Letterbox Resize
    img = letterbox_resize_opencv(img, target_size=target_size, fill_value=128)

    # 2. ToTensor (HWC -> CHW, 0-255 -> 0-1)
    img = img.transpose((2, 0, 1)).astype(np.float32)
    img /= 255.0

    # 3. Normalize
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(3, 1, 1)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(3, 1, 1)
    img = (img - mean) / std

    # Add batch dim
    img = np.expand_dims(img, axis=0)

    return img


def preprocess_image_raw(image_path):
    """只做读取和RGB转换，返回原始cv2对象"""
    if not os.path.exists(image_path):
        raise ValueError(f"文件不存在: {image_path}")

    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img


def preprocess_numpy(img_raw, target_size=(224, 224)):
    """对已经读取的图片做 resize 和 normalization"""
    if target_size:
        img = cv2.resize(img_raw, target_size)
    else:
        img = img_raw.copy()

    img = img.transpose((2, 0, 1))
    img = np.expand_dims(img, axis=0).astype(np.float32)
    img /= 255.0
    return img


def infer_classification(client, input_data, model_name, input_name, output_name, labels=None, data_type="FP32"):
    try:
        t2 = time.time()

        # 根据 data_type 转换数据类型
        if data_type == "FP16":
            input_data = input_data.astype(np.float16)

        # gRPC InferInput
        inputs = [grpcclient.InferInput(input_name, input_data.shape, data_type)]
        inputs[0].set_data_from_numpy(input_data)

        # gRPC InferRequestedOutput
        outputs = [grpcclient.InferRequestedOutput(output_name)]

        # gRPC Infer
        response = client.infer(model_name, inputs=inputs, outputs=outputs)
        t3 = time.time()

        logits = response.as_numpy(output_name)
        if logits.ndim == 1:
            logits = logits.reshape(1, -1)
        # 数值稳定 softmax，避免 exp 溢出
        logits_max = np.max(logits, axis=-1, keepdims=True)
        exp_logits = np.exp(logits - logits_max)
        probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
        probs = np.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)

        pred_id = int(np.argmax(probs, axis=-1).flat[0])
        confidence = float(np.asarray(probs).flat[pred_id])
        label_text = str(pred_id)
        if labels and pred_id < len(labels):
            label_text = labels[pred_id]

        return pred_id, confidence, label_text

    except Exception as e:
        pass
        return None, None, None


def infer_yolo(client, input_data, model_name, input_name, output_name, conf_thres=0.25, iou_thres=0.45, num_classes=None, labels=None, data_type="FP32"):
    try:
        t2 = time.time()

        # 根据 data_type 转换数据类型
        if data_type == "FP16":
            input_data = input_data.astype(np.float16)

        inputs = [grpcclient.InferInput(input_name, input_data.shape, data_type)]
        inputs[0].set_data_from_numpy(input_data)

        outputs = [grpcclient.InferRequestedOutput(output_name)]

        response = client.infer(model_name, inputs=inputs, outputs=outputs)
        t3 = time.time()

        det_out = response.as_numpy(output_name)

        # Handle End-to-End ONNX with built-in NMS (Batch, N, 6)
        if det_out.ndim == 3 and det_out.shape[2] == 6:
            pred = det_out[0] # (N, 6)

            scores = pred[:, 4]
            mask = scores > conf_thres

            filtered_pred = pred[mask]

            results = []
            for item in filtered_pred:
                # item: [x1, y1, x2, y2, score, class_id]
                x1, y1, x2, y2 = item[:4]
                score = item[4]
                cid = int(item[5])

                label_name = str(cid)
                if labels and cid < len(labels):
                    label_name = labels[cid]

                w = x2 - x1
                h = y2 - y1
                cx = x1 + w / 2
                cy = y1 + h / 2

                box = [cx, cy, w, h]
                results.append({"class_id": cid, "class_label": label_name, "confidence": score, "box": box})

            return results

        # Standard YOLO Output Processing
        pred = det_out[0].transpose()

        boxes_cxcywh = pred[:, :4]

        if num_classes:
            scores_mat = pred[:, 4:4+num_classes]
        else:
            scores_mat = pred[:, 4:]

        class_ids = np.argmax(scores_mat, axis=1)
        scores = np.max(scores_mat, axis=1)

        mask = scores > conf_thres

        filtered_boxes = boxes_cxcywh[mask]
        filtered_scores = scores[mask]
        filtered_ids = class_ids[mask]

        if len(filtered_boxes) > 0:
            nms_boxes = filtered_boxes.copy()
            nms_boxes[:, 0] = filtered_boxes[:, 0] - filtered_boxes[:, 2] / 2
            nms_boxes[:, 1] = filtered_boxes[:, 1] - filtered_boxes[:, 3] / 2

            indices = cv2.dnn.NMSBoxes(nms_boxes.tolist(), filtered_scores.tolist(), conf_thres, iou_thres)

            results = []
            if len(indices) > 0:
                indices = indices.flatten()
                for i in indices:
                    cid = filtered_ids[i]
                    conf = filtered_scores[i]
                    box = filtered_boxes[i] # cx, cy, w, h

                    label_name = str(cid)
                    if labels and cid < len(labels):
                        label_name = labels[cid]

                    results.append({"class_id": cid, "class_label": label_name, "confidence": conf, "box": box.tolist()})
            return results
        else:
            return []

    except Exception as e:
        pass
        return []


def infer_trt_nms_yolo(client, input_data, model_name, input_name="images", target_size=(640, 640), conf_thres=0.25, data_type="FP32", labels=None):
    """
    专门调用带有 TRT EfficientNMS 插件的模型 (gRPC 版)。
    """
    try:
        t2 = time.time()

        # 根据 data_type 转换数据类型
        if data_type == "FP16":
            input_data = input_data.astype(np.float16)

        inputs = [grpcclient.InferInput(input_name, input_data.shape, data_type)]
        inputs[0].set_data_from_numpy(input_data)

        # 定义 4 个输出请求
        outputs = [
            grpcclient.InferRequestedOutput("num_detections"),
            grpcclient.InferRequestedOutput("detection_boxes"),
            grpcclient.InferRequestedOutput("detection_scores"),
            grpcclient.InferRequestedOutput("detection_classes")
        ]

        response = client.infer(model_name, inputs=inputs, outputs=outputs)
        t3 = time.time()

        num_dets = response.as_numpy("num_detections")[0][0]
        det_boxes = response.as_numpy("detection_boxes")[0]   # [MAX_OUT, 4]
        det_scores = response.as_numpy("detection_scores")[0] # [MAX_OUT]
        det_classes = response.as_numpy("detection_classes")[0] # [MAX_OUT]

        results = []
        for i in range(num_dets):
            score = det_scores[i]
            if score < conf_thres:
                continue

            cid = int(det_classes[i])
            label_name = str(cid)
            if labels and cid < len(labels):
                label_name = labels[cid]

            box_xyxy = det_boxes[i] # [x1, y1, x2, y2]

            # 转换为 cx, cy, w, h 保持与原本 infer_yolo 一致
            x1, y1, x2, y2 = box_xyxy
            w = x2 - x1
            h = y2 - y1
            cx = x1 + w / 2
            cy = y1 + h / 2

            box_cxcywh = [float(cx), float(cy), float(w), float(h)]

            results.append({"class_id": cid, "class_label": label_name, "confidence": float(score), "box": box_cxcywh})

        return results

    except Exception as e:
        pass
        return []


def infer_porn(client, image_path, preprocessed_data_224=None, labels_nsfw=None, labels_cls=None, labels_seg=None, dtype_nsfw="FP16", dtype_cls="FP16", dtype_seg="FP32"):
    # 1. 先调用 nsfw 模型：输出 normal / nsfw，仅当为 nsfw 时再走 porn_classification
    pred_nsfw_id, conf_nsfw, _ = infer_classification(
        client, preprocessed_data_224, "nsfw", "images", "output0",
        labels=labels_nsfw,
        data_type=dtype_nsfw
    )
    if pred_nsfw_id is None:
        return "normal"

    nsfw_label_text = str(pred_nsfw_id)
    if labels_nsfw and pred_nsfw_id < len(labels_nsfw):
        nsfw_label_text = labels_nsfw[pred_nsfw_id]
    is_nsfw = nsfw_label_text.lower() == "nsfw"
    if not is_nsfw:
        return "normal"

    # 2. nsfw 为 nsfw 时，调用 porn_classification
    # labels_cls 对应 porn_classification.txt: 1=hentai, 3=porn 等
    pred_id, conf, label_text = infer_classification(
        client, preprocessed_data_224, "porn_classification", "images", "output0",
        labels=labels_cls,
        data_type=dtype_cls
    )
    if pred_id is None:
        return "normal"

    is_target = False
    if pred_id == 1 or pred_id == 3:
        is_target = True
    if labels_cls and pred_id < len(labels_cls):
        label_text = labels_cls[pred_id]
        if label_text in ["porn", "hentai"]:
            is_target = True
    else:
        label_text = str(pred_id)

    if is_target:
        infer_yolo(
            client, preprocessed_data_224, "porn_segmentation", "images", "output0",
            num_classes=4,
            labels=labels_seg,
            data_type=dtype_seg
        )
    return label_text


def infer_flag(client, image_path, preprocessed_data_640=None, labels_det=None, labels_cls=None, dtype_det="FP16", dtype_cls="FP16"):
    # 1. Detection (YOLO)
    # flag_detect uses 640 input
    results = infer_yolo(
        client, preprocessed_data_640, "flag_detect", "images", "output0",
        labels=labels_det,
        data_type=dtype_det
    )

    if not results:
        return "normal"

    # Load original image for cropping
    if not os.path.exists(image_path):
        return "normal"

    img_raw = cv2.imread(image_path)
    if img_raw is None:
        return "normal"

    h_raw, w_raw = img_raw.shape[:2]

    # Scale factors (assuming input was resized to 640x640)
    scale_x = w_raw / 640.0
    scale_y = h_raw / 640.0

    for i, res in enumerate(results):
        box = res["box"] # cx, cy, w, h
        cx, cy, w, h = box

        # Convert to x1, y1, x2, y2 on 640x640
        x1_640 = cx - w / 2
        y1_640 = cy - h / 2
        x2_640 = cx + w / 2
        y2_640 = cy + h / 2

        # Scale to original image
        x1 = int(x1_640 * scale_x)
        y1 = int(y1_640 * scale_y)
        x2 = int(x2_640 * scale_x)
        y2 = int(y2_640 * scale_y)

        # Clip to image boundaries
        x1 = max(0, x1)
        y1 = max(0, y1)
        x2 = min(w_raw, x2)
        y2 = min(h_raw, y2)

        if x2 <= x1 or y2 <= y1:
            continue

        crop = img_raw[y1:y2, x1:x2]

        # Call classification
        input_data = preprocess_flag_classification(crop, target_size=224)

        _, _, cls_label = infer_classification(
            client, input_data, "flag_classification", "pixel_values", "logits",
            labels=labels_cls,
            data_type=dtype_cls
        )
        if cls_label:
            return cls_label
    return "normal"


def run_task_with_client(client, mode, img_path, input_224, input_640, label_dir, model_label_map, model_data_types):
    """
    使用已有 client 跑单模型推理，使用传入的 label_dir、model_label_map、model_data_types。
    返回 (summary_label, duration_seconds)。
    """
    s = time.time()
    summary_label = "normal"
    try:
        def get_labels(model_name):
            label_path = os.path.join(label_dir, model_label_map.get(model_name, f"{model_name}.txt"))
            return load_labels(label_path)

        labels = get_labels(mode)
        dtype = model_data_types.get(mode, "FP32")

        if mode == "fire":
            res = infer_yolo(client, input_640, "fire", "images", "output0", labels=labels, data_type=dtype)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "politics":
            _, _, label = infer_classification(client, input_224, "politics", "images", "output0", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "qrcode":
            res = infer_yolo(client, input_640, "qrcode", "images", "output0", labels=labels, data_type=dtype)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "prohibitlogo":
            res = infer_yolo(client, input_640, "prohibitlogo", "images", "output0", labels=labels, data_type=dtype)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "chinamap":
            _, _, label = infer_classification(client, input_224, "chinamap", "input", "output", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "discomfort":
            _, _, label = infer_classification(client, input_224, "discomfort", "input", "output", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "discomfort_excrement":
            _, _, label = infer_classification(client, input_224, "discomfort_excrement", "input", "output", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "gamble":
            _, _, label = infer_classification(client, input_224, "gamble", "input", "output", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "tattoo":
            _, _, label = infer_classification(client, input_224, "tattoo", "input", "output", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "tiananmen":
            _, _, label = infer_classification(client, input_224, "tiananmen", "input", "output", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "trypophobia":
            _, _, label = infer_classification(client, input_224, "trypophobia", "images", "output0", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "blood":
            _, _, label = infer_classification(client, input_224, "blood", "input", "output", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "military":
            res = infer_trt_nms_yolo(client, input_640, "military", "images", target_size=(640, 640), data_type=dtype, labels=labels)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "flag_detect":
            res = infer_yolo(client, input_640, "flag_detect", "images", "output0", labels=labels, data_type=dtype)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "flag_classification":
            _, _, label = infer_classification(client, input_224, "flag_classification", "pixel_values", "logits", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "porn_classification":
            _, _, label = infer_classification(client, input_224, "porn_classification", "images", "output0", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "nsfw":
            _, _, label = infer_classification(client, input_224, "nsfw", "images", "output0", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "emblem":
            res = infer_yolo(client, input_640, "emblem", "images", "output0", labels=labels, data_type=dtype)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "religion_symbol":
            res = infer_yolo(client, input_640, "religion_symbols", "images", "output0", labels=labels, data_type=dtype)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "religion_activity":
            _, _, label = infer_classification(client, input_224, "religion_activity", "images", "output0", labels=labels, data_type=dtype)
            summary_label = _to_summary_label(label)

        elif mode == "religion_clothes":
            res = infer_yolo(client, input_640, "religion_clothes", "images", "output0", labels=labels, data_type=dtype)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "media1":
            res = infer_yolo(client, input_640, "media_version1", "images", "output0", labels=labels, data_type=dtype)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "media2":
            res = infer_yolo(client, input_640, "media_version2", "images", "output0", labels=labels, data_type=dtype)
            summary_label = "normal" if not res else res[0]["class_label"]

        elif mode == "porn":
            labels_nsfw = get_labels("nsfw")
            labels_cls = get_labels("porn_classification")
            out = infer_porn(client, img_path,
                            preprocessed_data_224=input_224,
                            labels_nsfw=labels_nsfw,
                            labels_cls=labels_cls,
                            dtype_nsfw=model_data_types.get("nsfw", "FP16"),
                            dtype_cls=model_data_types.get("porn_classification", "FP16"))
            summary_label = _to_summary_label(out if out else "normal")

        elif mode == "flag":
            labels_det = get_labels("flag_detect")
            labels_cls = get_labels("flag_classification")
            out = infer_flag(client, img_path,
                            preprocessed_data_640=input_640,
                            labels_det=labels_det,
                            labels_cls=labels_cls,
                            dtype_det=model_data_types.get("flag_detect", "FP16"),
                            dtype_cls=model_data_types.get("flag_classification", "FP16"))
            summary_label = _to_summary_label(out if out else "normal")

        duration = time.time() - s
        return summary_label, duration
    except Exception as e:
        return "normal", time.time() - s
