"""虎牙滑块/点选验证码 OCR。"""

from __future__ import annotations

from functools import lru_cache
from io import BytesIO
from pathlib import Path
from typing import cast

import cv2
import numpy as np
import onnxruntime as ort
from PIL import Image
from scipy.optimize import linear_sum_assignment

MODEL_DIR = Path(__file__).resolve().parent / "models"


class YoloOnnx:
    """轻量 YOLO ONNX 推理封装。"""

    def __init__(
        self,
        model_path: str | Path,
        classes: list[str],
        providers: list[str] | None = None,
    ):
        providers = providers or ["CPUExecutionProvider"]
        ort.set_default_logger_severity(3)
        self.session = ort.InferenceSession(str(model_path), providers=providers)
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape[2:]
        self.names = classes
        self.scale = 1.0
        self.pad_top = 0
        self.pad_left = 0

    @staticmethod
    def xywh2xyxy(value: np.ndarray) -> np.ndarray:
        result = np.copy(value)
        result[..., 0] = value[..., 0] - value[..., 2] / 2
        result[..., 1] = value[..., 1] - value[..., 3] / 2
        result[..., 2] = value[..., 0] + value[..., 2] / 2
        result[..., 3] = value[..., 1] + value[..., 3] / 2
        return result

    def preprocess(self, image: Image.Image) -> np.ndarray:
        img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        height, width = img.shape[:2]

        self.scale = min(self.input_shape[0] / height, self.input_shape[1] / width)
        new_size = (int(width * self.scale), int(height * self.scale))
        resized = cv2.resize(img, new_size, interpolation=cv2.INTER_LINEAR)

        dh = self.input_shape[0] - new_size[1]
        dw = self.input_shape[1] - new_size[0]
        self.pad_top = dh // 2
        pad_bottom = dh - self.pad_top
        self.pad_left = dw // 2
        pad_right = dw - self.pad_left

        padded = cv2.copyMakeBorder(
            resized,
            self.pad_top,
            pad_bottom,
            self.pad_left,
            pad_right,
            cv2.BORDER_CONSTANT,
            value=(114, 114, 114),
        )
        input_tensor = padded.transpose(2, 0, 1).astype(np.float32) / 255.0
        return np.expand_dims(input_tensor, axis=0)

    @staticmethod
    def nms(
        boxes: np.ndarray, scores: np.ndarray, iou_threshold: float = 0.3
    ) -> list[int]:
        order = scores.argsort()[::-1]
        keep = []

        while order.size > 0:
            index = order[0]
            keep.append(index)

            xx1 = np.maximum(boxes[index, 0], boxes[order[1:], 0])
            yy1 = np.maximum(boxes[index, 1], boxes[order[1:], 1])
            xx2 = np.minimum(boxes[index, 2], boxes[order[1:], 2])
            yy2 = np.minimum(boxes[index, 3], boxes[order[1:], 3])

            width = np.maximum(0.0, xx2 - xx1)
            height = np.maximum(0.0, yy2 - yy1)
            intersection = width * height
            area_i = (boxes[index, 2] - boxes[index, 0]) * (
                boxes[index, 3] - boxes[index, 1]
            )
            area_j = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (
                boxes[order[1:], 3] - boxes[order[1:], 1]
            )
            iou = intersection / (area_i + area_j - intersection + 1e-7)

            inds = np.where(iou <= iou_threshold)[0]
            order = order[inds + 1]

        return keep

    def postprocess(
        self,
        outputs: list[np.ndarray],
        conf_threshold: float = 0.1,
        iou_threshold: float = 0.3,
    ) -> list[dict]:
        predictions = np.squeeze(outputs[0]).T
        class_scores = predictions[:, 4:]
        scores = np.max(class_scores, axis=1)
        class_ids = np.argmax(class_scores, axis=1)

        valid_indices = scores > conf_threshold
        predictions = predictions[valid_indices]
        scores = scores[valid_indices]
        class_ids = class_ids[valid_indices]
        if predictions.size == 0:
            return []

        boxes = predictions[:, :4].copy()
        boxes[:, 0] = (boxes[:, 0] - self.pad_left) / self.scale
        boxes[:, 1] = (boxes[:, 1] - self.pad_top) / self.scale
        boxes[:, 2] /= self.scale
        boxes[:, 3] /= self.scale
        xyxy_boxes = self.xywh2xyxy(boxes)

        keep_indices = []
        for class_id in np.unique(class_ids):
            class_indices = np.where(class_ids == class_id)[0]
            class_keep = self.nms(
                xyxy_boxes[class_indices],
                scores[class_indices],
                iou_threshold,
            )
            keep_indices.extend(class_indices[class_keep])

        results = []
        for index in keep_indices:
            class_id = int(class_ids[index])
            x1, y1, x2, y2 = map(int, xyxy_boxes[index].tolist())
            results.append(
                {
                    "label_id": class_id,
                    "label_name": self.names[class_id],
                    "confidence": float(scores[index]),
                    "box_mid_xy": [(x1 + x2) // 2, (y1 + y2) // 2],
                    "xyxy": [x1, y1, x2, y2],
                }
            )

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

    def detect(self, image: Image.Image) -> list[dict]:
        input_tensor = self.preprocess(image)
        outputs = cast(
            list[np.ndarray],
            self.session.run([self.output_name], {self.input_name: input_tensor}),
        )
        return self.postprocess(outputs)


class SimilarityOnnx:
    """点选验证码中目标图和候选字的相似度模型。"""

    def __init__(self, model_path: str | Path, providers: list[str] | None = None):
        providers = providers or ["CPUExecutionProvider"]
        ort.set_default_logger_severity(3)
        self.session = ort.InferenceSession(str(model_path), providers=providers)
        self.input_shape: tuple[int, int] = (64, 64)

    @staticmethod
    def sigmoid(value: np.ndarray) -> np.ndarray:
        return 1 / (1 + np.exp(-value))

    @staticmethod
    def _to_image(value) -> Image.Image:
        if isinstance(value, np.ndarray):
            return Image.fromarray(value)
        if isinstance(value, bytes):
            return Image.open(BytesIO(value))
        if isinstance(value, Image.Image):
            return value
        return Image.open(value)

    def _tensor(self, value) -> np.ndarray:
        image = (
            self._to_image(value)
            .convert("RGB")
            .resize((self.input_shape[1], self.input_shape[0]), 1)
        )
        array = np.array(image).astype(np.float32) / 255.0
        return np.expand_dims(np.transpose(array, (2, 0, 1)), 0)

    def score(self, image_1, image_2) -> int:
        out = self.session.run(
            None, {"x1": self._tensor(image_1), "x2": self._tensor(image_2)}
        )
        similarity = self.sigmoid(cast(np.ndarray, out[0]))[0][0]
        return int(round(similarity.item(), 2) * 100)


class HuyaCaptchaOcr:
    """封装虎牙滑块和点选识别。"""

    def __init__(
        self, model_dir: str | Path = MODEL_DIR, providers: list[str] | None = None
    ):
        model_dir = Path(model_dir)
        self.similarity = SimilarityOnnx(
            model_dir / "weights.onnx", providers=providers
        )
        self.click_model = YoloOnnx(
            model_dir / "best.onnx", classes=["target", "char"], providers=providers
        )
        self.slider_model = YoloOnnx(
            model_dir / "slider_2.onnx", classes=["slider"], providers=providers
        )

    @staticmethod
    def _open_image(image_bytes: bytes) -> Image.Image:
        return Image.open(BytesIO(image_bytes)).convert("RGB")

    def ocr_slider_box(self, image_bytes: bytes) -> list[int]:
        """识别滑块缺口框，返回 [x1, y1, x2, y2]。"""
        image = self._open_image(image_bytes)
        detections = self.slider_model.detect(image)
        boxes = [item["xyxy"] for item in detections]
        return boxes[0] if boxes else [100, 100, 100, 100]

    def ocr_click_points(self, image_bytes: bytes) -> list[list[int]]:
        """识别点选坐标，返回按目标顺序排列的点击点。"""
        image = self._open_image(image_bytes)
        detections = self.click_model.detect(image)
        results = [
            {**item, "cropped_image": image.crop(tuple(item["xyxy"]))}
            for item in detections
        ]

        char_list = [item for item in results if item.get("label_name") == "char"]
        target_list = [item for item in results if item.get("label_name") == "target"]
        target_list = [
            item for item in target_list if item["xyxy"][2] - item["xyxy"][0] > 10
        ]
        char_list.sort(key=lambda item: item["xyxy"][0])
        target_list.sort(key=lambda item: item["xyxy"][0])

        if not char_list or not target_list:
            return []

        score_matrix = np.zeros((len(target_list), len(char_list)))
        for target_index, target in enumerate(target_list):
            for char_index, char in enumerate(char_list):
                try:
                    score_matrix[target_index][char_index] = -self.similarity.score(
                        target["cropped_image"],
                        char["cropped_image"],
                    )
                except Exception:
                    score_matrix[target_index][char_index] = 1e6

        row_ind, col_ind = linear_sum_assignment(score_matrix)
        matched = []
        for target_index, char_index in zip(row_ind, col_ind):
            if score_matrix[target_index][char_index] == 1e6:
                continue
            char = char_list[char_index]
            matched.append({**char, "index": target_index + 1})

        return [item["box_mid_xy"] for item in matched]


@lru_cache(maxsize=1)
def default_ocr() -> HuyaCaptchaOcr:
    """懒加载默认 OCR，避免导入 core.huya 时立即加载 150MB 模型。"""
    return HuyaCaptchaOcr()