import base64
from typing import Tuple

import cv2
import numpy as np

from .base import DetectResult

# 按严重程度着色，BGR格式
SEVERITY_COLORS: dict = {
    "fatal":  (0,   0,   220),
    "major":  (0,   128, 255),
    "minor":  (0,   215, 255),
    "rework": (255, 165,   0),
    "none":   (180, 180, 180),
}


def annotate(image: np.ndarray, result: DetectResult) -> np.ndarray:
    img = image.copy()
    for d in result.defects:
        color: Tuple = SEVERITY_COLORS.get(d.severity, (180, 180, 180))
        x1, y1, x2, y2 = [int(v) for v in d.box_xyxy]
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
        label = f"{d.class_name_zh} {d.confidence:.2f}"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(img, (x1, y1 - th - 6), (x1 + tw + 4, y1), color, -1)
        cv2.putText(img, label, (x1 + 2, y1 - 4),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
    return img


def to_base64(image: np.ndarray, quality: int = 85) -> str:
    ok, buf = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, quality])
    if not ok:
        raise RuntimeError("图像编码失败")
    return base64.b64encode(buf.tobytes()).decode()