"""撮影画像 → 舌の検出 → 舌の分割 → マスク適用 → 舌色 5 分類をまとめた入口（Issue #13）．

処理手順・閾値・クラス順は TIASJudgeModel の引き渡し仕様
（`docs/04_SPEC/SPEC_04_TIASjudge引き渡し仕様.md`）に従う．各段の計算は参照実装のコピー
（`detection.py`・`resunet.py`・`pipeline.py`・`color_model.py`）を呼ぶだけで，本モジュールは
設定値の固定・モデルの遅延ロード・失敗の判定・結果の組み立てだけを持つ．

1. 舌の検出（`Detection_Normal.onnx`）: スコアが 0.5 を超える枠のうち最高の 1 つ
2. 舌の分割（`ResUNet.onnx`）: 枠の中の舌マスク → 元の画像サイズのマスク
3. 撮影画像（BGR）を RGB に並べ替え，マスクの外を (0, 0, 0) にする
4. 色特徴（107 次元）→ ロジスティック回帰（`color_model.json`）: 5 クラスの確率

入力は **TIAS の機材で撮影した色補正なしの画像**に限る（変換行列は通さない）．
検出できない・マスクが空の場合は判定せず，その旨を結果の `status` で返す（例外にしない）．
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from enum import Enum
from pathlib import Path

import cv2
import numpy as np
from numpy.typing import NDArray
from onnxruntime.capi import onnxruntime_pybind11_state as ort_errors

from tiasjudge.core.ml.color_model import (
    ColorModel,
    downscale_nearest,
    load_color_model,
    predict_proba,
)
from tiasjudge.core.ml.detection import OnnxDetector, Rect
from tiasjudge.core.ml.diagnosis import MlDiagnosis, build_ml_diagnosis
from tiasjudge.core.ml.pipeline import SegmentationParams, segment_in_rect
from tiasjudge.core.ml.resunet import ResUNetRunner

logger = logging.getLogger(__name__)

#: onnxruntime がモデルの読込失敗（壊れた ONNX・非対応の演算等）で送出する例外．
#: いずれも `Exception` の直接の派生で `RuntimeError` ではないため，列挙して捕捉する
_ORT_LOAD_ERRORS: tuple[type[Exception], ...] = (
    ort_errors.Fail,
    ort_errors.InvalidArgument,
    ort_errors.InvalidGraph,
    ort_errors.InvalidProtobuf,
    ort_errors.NoModel,
    ort_errors.NoSuchFile,
    ort_errors.NotImplemented,
    ort_errors.RuntimeException,
)

# --- モデルファイル名（`data/models/` に置く．来歴は `data/models/README.md`）---

#: 舌の検出モデル（DeepTIAS 2.7 の `Detection_Normal.pb` を ONNX に変換したもの）
DETECTION_MODEL_FILE_NAME = "Detection_Normal.onnx"
#: 舌の分割モデル（DeepTIAS 2.7 同梱の ResUNet）
SEGMENTATION_MODEL_FILE_NAME = "ResUNet.onnx"
#: 舌色の判定モデル（色特徴＋ロジスティック回帰）
COLOR_MODEL_FILE_NAME = "color_model.json"

# --- 設定値（引き渡し仕様「1. 舌の検出」「2. 舌の分割」．DeepTIAS 2.7 と同じ値）---

#: 検出の採用スコアの下限（これを **超える** 枠のみ採用．0.5 ちょうどは不採用）
DETECTION_SCORE_THRESHOLD = 0.5
#: 検出モデルへ渡す画像の色の順（TF の物体検出 API の既定は RGB）
DETECTION_INPUT_COLOR = "rgb"
#: 検出モデルの入出力名（`:0` の有無は参照実装が吸収する）
DETECTION_INPUT_NAME = "image_tensor"
DETECTION_BOXES_NAME = "detection_boxes"
DETECTION_SCORES_NAME = "detection_scores"
#: ResUNet の入力の一辺
RESUNET_INPUT_SIZE = 256
#: 分割の後処理の設定（TEST_07 の評価 IoU 0.958 はこの値）
SEGMENTATION_PARAMS = SegmentationParams(foreground_channel=0, min_area=500, mask_threshold=127)

# --- Git LFS のポインタファイルの判定（LFS を取得しない checkout ではモデルの代わりに置かれる）---

#: LFS のポインタファイルの先頭（git-lfs の仕様 v1）
LFS_POINTER_PREFIX = b"version https://git-lfs.github.com/spec/v1"
#: ポインタファイルとみなす大きさの上限（実物は 130 バイト前後．git-lfs は 1024 未満を扱う）
LFS_POINTER_MAX_BYTES = 1024


class MlModelError(Exception):
    """判定モデルを使えないことを知らせるエラーの基底（GUI がエラーダイアログで案内する）．"""


class ModelUnavailableError(MlModelError):
    """モデルファイルが無い，または Git LFS のポインタファイルのままである．"""


class ModelLoadError(MlModelError):
    """モデルファイルはあるが読み込めない（壊れている・形式が合わない等）．"""


def is_lfs_pointer(path: Path) -> bool:
    """ファイルが Git LFS のポインタファイル（実体を取得していない状態）かを返す．

    Args:
        path: 調べるファイル．

    Returns:
        ポインタファイルなら True（ファイルが無い・読めない場合は False）．
    """
    try:
        if path.stat().st_size >= LFS_POINTER_MAX_BYTES:
            return False
        with path.open("rb") as file:
            head = file.read(len(LFS_POINTER_PREFIX))
    except OSError:
        return False
    return head == LFS_POINTER_PREFIX


def is_model_file_available(path: Path) -> bool:
    """モデルファイルの実体があるか（存在し，LFS のポインタファイルでない）を返す．

    CI 等で LFS を取得しない checkout ではポインタファイルが置かれるため，存在だけでは
    判定できない．テストのスキップ条件にも使える．

    Args:
        path: モデルファイルのパス．

    Returns:
        実体があれば True．
    """
    return path.is_file() and not is_lfs_pointer(path)


@dataclass(frozen=True)
class ModelPaths:
    """判定に使う 3 つのモデルファイルのパス．

    Attributes:
        detection: 舌の検出モデル（`Detection_Normal.onnx`）．
        segmentation: 舌の分割モデル（`ResUNet.onnx`）．
        color: 舌色の判定モデル（`color_model.json`）．
    """

    detection: Path
    segmentation: Path
    color: Path

    @classmethod
    def in_directory(cls, directory: Path) -> ModelPaths:
        """既定のファイル名で `directory` 直下のパスを組み立てる．

        Args:
            directory: モデルファイルを置いたディレクトリ（`data/models` 等）．

        Returns:
            3 つのモデルファイルのパス．
        """
        return cls(
            detection=directory / DETECTION_MODEL_FILE_NAME,
            segmentation=directory / SEGMENTATION_MODEL_FILE_NAME,
            color=directory / COLOR_MODEL_FILE_NAME,
        )

    def all_paths(self) -> tuple[Path, Path, Path]:
        """3 つのパスを（検出・分割・判定）の順で返す．"""
        return (self.detection, self.segmentation, self.color)


#: 使えない理由（利用者向け）
UNAVAILABLE_REASON_MISSING = "ファイルが無い"
UNAVAILABLE_REASON_LFS_POINTER = "Git LFS のポインタファイルのままで，実体が無い"


@dataclass(frozen=True)
class UnavailableModel:
    """使えないモデルファイル 1 つ．

    Attributes:
        path: モデルファイルのパス．
        reason: 使えない理由（利用者向け．`UNAVAILABLE_REASON_*`）．
    """

    path: Path
    reason: str

    def __str__(self) -> str:
        """「<パス>（理由）」の形の説明．"""
        return f"{self.path}（{self.reason}）"


def find_unavailable_models(paths: ModelPaths) -> list[UnavailableModel]:
    """使えないモデルファイル（無い・LFS のポインタのまま）を返す（すべて使えれば空）．

    ファイルの存在と先頭の数十バイトだけを調べる（ONNX は読み込まない．起動時の判定用）．

    Args:
        paths: モデルファイルのパス．

    Returns:
        使えないモデルファイルの並び（検出・分割・判定の順）．
    """
    problems: list[UnavailableModel] = []
    for path in paths.all_paths():
        if not path.is_file():
            problems.append(UnavailableModel(path, UNAVAILABLE_REASON_MISSING))
        elif is_lfs_pointer(path):
            problems.append(UnavailableModel(path, UNAVAILABLE_REASON_LFS_POINTER))
    return problems


class JudgeStatus(Enum):
    """判定の結果の種類．"""

    #: 判定できた
    SUCCEEDED = "succeeded"
    #: 舌を検出できなかった（スコアが閾値を超える枠が無い）
    NOT_DETECTED = "not_detected"
    #: 舌の分割に失敗した（マスクの舌の画素が 0）
    EMPTY_MASK = "empty_mask"


@dataclass(frozen=True)
class ClassProbability:
    """1 クラスの予測確率．

    Attributes:
        class_id: クラス ID（`color_model.json` の `classes` の `id`．0 淡白〜4 紫）．
        slug: クラスの識別子（`pale`・`light_red`・`red`・`dark_red`・`purple`）．
        display_name: 表示名（日本語．`color_model.json` の `display_name`）．
        probability: 予測確率（0〜1．5 クラスの和は 1）．
    """

    class_id: int
    slug: str
    display_name: str
    probability: float


@dataclass(frozen=True)
class MlJudgeResult:
    """判定の結果．

    Attributes:
        status: 結果の種類．
        detection_score: 検出の最高スコア（検出モデルの出力が空なら None）．
        rect: 採用した舌の枠（画素．検出できなければ None）．
        mask: `(H, W)` の bool の舌マスク（元の画像サイズ．検出できなければ None）．
        probabilities: 5 クラスの予測確率（クラス ID 順．判定しなかった場合は空）．
    """

    status: JudgeStatus
    detection_score: float | None
    rect: Rect | None
    mask: NDArray[np.bool_] | None
    probabilities: tuple[ClassProbability, ...]

    @property
    def succeeded(self) -> bool:
        """判定できたか．"""
        return self.status is JudgeStatus.SUCCEEDED

    @property
    def top(self) -> ClassProbability | None:
        """最大の確率のクラス（判定結果）．判定しなかった場合は None．

        同点は ID の小さいクラス（`np.argmax` と同じ．引き渡し仕様「4. 舌色の判定」）．
        """
        if not self.probabilities:
            return None
        values = np.array([item.probability for item in self.probabilities])
        return self.probabilities[int(np.argmax(values))]

    @property
    def mask_pixel_count(self) -> int:
        """舌マスクの画素数（マスクが無ければ 0）．"""
        return 0 if self.mask is None else int(np.count_nonzero(self.mask))

    def to_diagnosis(self) -> MlDiagnosis:
        """既存の診断結果と同じ形（5 色の割合・上位 2 色）にする（画面表示用）．

        Raises:
            ValueError: 判定しなかった結果（検出できない・マスクが空）の場合．
        """
        if not self.succeeded:
            raise ValueError(f"判定していない結果は変換できない（{self.status.value}）")
        return build_ml_diagnosis({item.slug: item.probability for item in self.probabilities})


def apply_mask(bgr: NDArray[np.uint8], mask: NDArray[np.bool_]) -> NDArray[np.uint8]:
    """撮影画像を RGB に並べ替え，マスクの外を (0, 0, 0) にする（引き渡し仕様「3.」）．

    Args:
        bgr: 撮影画像（`(H, W, 3)`, uint8．OpenCV で読んだ BGR のまま）．
        mask: `(H, W)` の bool の舌マスク．

    Returns:
        `(H, W, 3)` の uint8 の RGB 画像（舌の外は黒）．
    """
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    rgb[~mask] = 0
    return rgb


def _has_tongue_pixels(rgb_masked: NDArray[np.uint8], model: ColorModel) -> bool:
    """判定モデルが特徴を取れる（背景の黒以外の画素が残る）かを返す．

    `predict_proba` は内部で `max_side` へ縮小してから黒以外の画素を舌とみなすため，
    同じ縮小をしてから確かめる（ごく小さなマスクが縮小で消える場合も判定しない）．
    """
    arr = rgb_masked
    if model.max_side is not None:
        arr = downscale_nearest(arr, model.max_side)
    return bool(np.any(arr != 0))


class TongueColorJudge:
    """撮影画像から舌色 5 分類の確率を求める（検出 → 分割 → マスク適用 → 判定）．

    モデル（ONNX セッション・判定モデル）は生成時には読み込まず，初回の `judge`（または
    `load`）で読み込んで以後は再利用する（起動を遅くしないため）．

    Args:
        paths: モデルファイルのパス．
    """

    def __init__(self, paths: ModelPaths) -> None:
        self._paths = paths
        self._detector: OnnxDetector | None = None
        self._runner: ResUNetRunner | None = None
        self._color_model: ColorModel | None = None

    @property
    def paths(self) -> ModelPaths:
        """モデルファイルのパス．"""
        return self._paths

    @property
    def is_loaded(self) -> bool:
        """モデルを読み込み済みか．"""
        return self._color_model is not None

    def load(self) -> None:
        """モデルを読み込む（読み込み済みなら何もしない）．

        Raises:
            ModelUnavailableError: モデルファイルが無い，または LFS のポインタファイルの場合．
            ModelLoadError: モデルファイルを読み込めない場合．
        """
        if self.is_loaded:
            return
        problems = find_unavailable_models(self._paths)
        if problems:
            raise ModelUnavailableError(
                "使えないモデルファイルがある:\n" + "\n".join(map(str, problems))
            )

        logger.info("舌色判定のモデルを読み込みます: %s", self._paths)
        try:
            detector = OnnxDetector(
                self._paths.detection,
                score_threshold=DETECTION_SCORE_THRESHOLD,
                input_color=DETECTION_INPUT_COLOR,
                input_name=DETECTION_INPUT_NAME,
                boxes_name=DETECTION_BOXES_NAME,
                scores_name=DETECTION_SCORES_NAME,
            )
            runner = ResUNetRunner(self._paths.segmentation, RESUNET_INPUT_SIZE)
            color_model = load_color_model(self._paths.color)
        except (ValueError, OSError, *_ORT_LOAD_ERRORS) as error:
            # 参照実装の入出力名・形状の不一致と JSON の形式不正は ValueError（JSON の構文誤りの
            # `json.JSONDecodeError` も ValueError の派生）．壊れた ONNX は onnxruntime の例外
            logger.warning("舌色判定のモデルを読み込めません: %s", error)
            raise ModelLoadError(f"モデルファイルを読み込めない: {error}") from error

        self._detector = detector
        self._runner = runner
        self._color_model = color_model

    def judge(self, bgr: NDArray[np.uint8]) -> MlJudgeResult:
        """撮影画像 1 枚の舌色を判定する．

        Args:
            bgr: 撮影画像（`(H, W, 3)`, uint8．OpenCV で読んだ BGR のまま．色補正なし．
                TIASjudge の変換行列は通さない）．

        Returns:
            判定の結果（検出できない・マスクが空なら `status` でそれを示し，確率は空）．

        Raises:
            ModelUnavailableError: モデルファイルが無い，または LFS のポインタファイルの場合．
            ModelLoadError: モデルファイルを読み込めない場合．
        """
        self.load()
        detector, runner, color_model = self._detector, self._runner, self._color_model
        if detector is None or runner is None or color_model is None:  # load() 後は常に揃う
            raise ModelLoadError("モデルが読み込まれていない")

        detection = detector.detect(bgr)
        if detection.rect is None:
            logger.info("舌を検出できませんでした（最高スコア %s）", detection.score)
            return MlJudgeResult(
                status=JudgeStatus.NOT_DETECTED,
                detection_score=detection.score,
                rect=None,
                mask=None,
                probabilities=(),
            )

        masks = segment_in_rect(bgr, detection.rect, runner, SEGMENTATION_PARAMS)
        mask = masks.post
        rgb_masked = apply_mask(bgr, mask)
        if not mask.any() or not _has_tongue_pixels(rgb_masked, color_model):
            logger.info(
                "舌の領域が空でした（検出スコア %.4f・枠 %s）", detection.score, detection.rect
            )
            return MlJudgeResult(
                status=JudgeStatus.EMPTY_MASK,
                detection_score=detection.score,
                rect=detection.rect,
                mask=mask,
                probabilities=(),
            )

        proba = predict_proba(rgb_masked, color_model)
        probabilities = tuple(
            ClassProbability(
                class_id=int(item["id"]),
                slug=str(item["slug"]),
                display_name=str(item["display_name"]),
                probability=float(value),
            )
            for item, value in zip(color_model.classes, proba, strict=True)
        )
        result = MlJudgeResult(
            status=JudgeStatus.SUCCEEDED,
            detection_score=detection.score,
            rect=detection.rect,
            mask=mask,
            probabilities=probabilities,
        )
        # 研究用途での再確認のため，入力（枠・マスク）と出力（確率）を INFO で残す
        # （error-handling.md「ログ」）
        logger.info(
            "舌色判定: 検出スコア %.4f・枠 %s・舌の画素 %d・確率 %s・判定 %s",
            detection.score,
            detection.rect,
            result.mask_pixel_count,
            {item.slug: item.probability for item in probabilities},
            result.top.slug if result.top is not None else None,
        )
        return result