"""
FunCaptcha 专项 ONNX 推理。
"""

from __future__ import annotations

import io
import time
from pathlib import Path

import numpy as np
from PIL import Image

from config import FUN_CAPTCHA_TASKS, INFERENCE_CONFIG
from inference.model_metadata import load_model_metadata
from inference.pipeline import _try_import_ort


class FunCaptchaRollballPipeline:
    """
    `4_3d_rollball_animals` 专项推理器。

    输入整张 challenge 图片，内部自动裁切 reference / candidates，
    再使用 Siamese ONNX 模型逐个候选打分。
    """

    def __init__(self, question: str = "4_3d_rollball_animals", models_dir: str | None = None):
        if question not in FUN_CAPTCHA_TASKS:
            raise ValueError(f"不支持的 FunCaptcha question: {question}")

        ort = _try_import_ort()
        self.question = question
        self.task_cfg = FUN_CAPTCHA_TASKS[question]
        self.models_dir = Path(models_dir or INFERENCE_CONFIG["default_models_dir"])
        self.model_path = self.models_dir / f"{self.task_cfg['artifact_name']}.onnx"
        if not self.model_path.exists():
            raise FileNotFoundError(f"未找到 FunCaptcha ONNX 模型: {self.model_path}")

        opts = ort.SessionOptions()
        opts.inter_op_num_threads = 1
        opts.intra_op_num_threads = 2
        self.session = ort.InferenceSession(
            str(self.model_path),
            sess_options=opts,
            providers=["CPUExecutionProvider"],
        )
        self.metadata = load_model_metadata(self.model_path) or {}
        self.mean = float(INFERENCE_CONFIG["normalize_mean"])
        self.std = float(INFERENCE_CONFIG["normalize_std"])
        self.answer_index_base = int(
            self.metadata.get("answer_index_base", self.task_cfg["answer_index_base"])
        )

    def solve(self, image) -> dict:
        t0 = time.perf_counter()
        challenge = self._load_image(image)
        candidates, reference = self._split_challenge(challenge)

        ref_batch = np.repeat(reference, repeats=candidates.shape[0], axis=0)
        input_names = [inp.name for inp in self.session.get_inputs()]
        if len(input_names) != 2:
            raise RuntimeError(f"专项模型输入数量异常: expected=2 got={len(input_names)}")

        logits = self.session.run(None, {
            input_names[0]: candidates,
            input_names[1]: ref_batch,
        })[0].reshape(-1)
        scores = 1.0 / (1.0 + np.exp(-logits))
        answer_idx = int(np.argmax(logits))
        selected = answer_idx + self.answer_index_base
        elapsed = (time.perf_counter() - t0) * 1000

        return {
            "type": "funcaptcha",
            "question": self.question,
            "objects": [selected],
            "scores": [round(float(score), 6) for score in scores.tolist()],
            "raw": str(selected),
            "result": str(selected),
            "time_ms": round(elapsed, 2),
        }

    def _split_challenge(self, image: Image.Image) -> tuple[np.ndarray, np.ndarray]:
        tile_w, tile_h = self.metadata.get("tile_size", self.task_cfg["tile_size"])
        ref_box = tuple(self.metadata.get("reference_box", self.task_cfg["reference_box"]))
        num_candidates = int(self.metadata.get("num_candidates", self.task_cfg["num_candidates"]))
        input_h, input_w = self.task_cfg["input_size"]

        candidates = []
        for idx in range(num_candidates):
            left = idx * tile_w
            candidate = image.crop((left, 0, left + tile_w, tile_h))
            candidates.append(self._preprocess(candidate, (input_h, input_w)))

        reference = image.crop(ref_box)
        return (
            np.concatenate(candidates, axis=0),
            self._preprocess(reference, (input_h, input_w)),
        )

    def _preprocess(self, image: Image.Image, target_size: tuple[int, int]) -> np.ndarray:
        img_h, img_w = target_size
        image = image.convert("RGB").resize((img_w, img_h), Image.BILINEAR)
        arr = np.asarray(image, dtype=np.float32) / 255.0
        arr = (arr - self.mean) / self.std
        arr = np.transpose(arr, (2, 0, 1))
        return arr.reshape(1, 3, img_h, img_w)

    @staticmethod
    def _load_image(image) -> Image.Image:
        if isinstance(image, Image.Image):
            return image
        if isinstance(image, (str, Path)):
            return Image.open(image).convert("RGB")
        if isinstance(image, bytes):
            return Image.open(io.BytesIO(image)).convert("RGB")
        raise TypeError(f"不支持的图片输入类型: {type(image)}")