"""pydantic による設定スキーマと設定 YAML の読み込み.

学習・モデル関連の設定は後続のステップで `ExperimentConfig` にフィールドを追加して拡張する．
"""

from __future__ import annotations

import os
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import yaml
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from tias_judge_model.data.classes import SEVERITIES

# データのルートを上書きする環境変数．サーバと手元で YAML を書き換えずに済ませるため．
DATA_ROOT_ENV_VAR = "TIAS_DATA_ROOT"
# 実験 ID に入れる git の短縮ハッシュ．サーバには .git を送らないため，tias-remote run が
# 手元で求めてこの環境変数で渡す（SPEC_01「実験の出力」・SPEC_03「run」）．
GIT_COMMIT_ENV_VAR = "TIAS_GIT_COMMIT"


class _StrictModel(BaseModel):
    """未知のキー（綴り間違い）を拒否し，生成後の変更を禁止する基底モデル."""

    model_config = ConfigDict(extra="forbid", frozen=True)


class DataConfig(_StrictModel):
    """データの場所に関する設定.

    Attributes:
        root: データのルートディレクトリ（`images/` と manifest を含む）．`~` は展開される．
            相対パスはカレントディレクトリ基準で解釈する．
        manifest: manifest CSV のパス．相対パスは `root` からの相対として解釈する．
        severities: 学習・評価に使う程度（`strong`・`mild`・`normal` の部分集合）．
            分割の後に適用する（SPEC_02「程度（ややあり）の扱い」）．値は `SEVERITIES` の
            順（strong → mild → normal）にそろえて保持する．
    """

    root: Path
    manifest: Path = Path("manifest.csv")
    severities: tuple[str, ...] = SEVERITIES

    @field_validator("root", "manifest", mode="after")
    @classmethod
    def _expand_user(cls, value: Path) -> Path:
        return value.expanduser()

    @field_validator("severities", mode="after")
    @classmethod
    def _check_severities(cls, value: tuple[str, ...]) -> tuple[str, ...]:
        if not value:
            raise ValueError("data.severities が空である（strong・mild・normal から 1 つ以上）")
        unknown = [s for s in value if s not in SEVERITIES]
        if unknown:
            raise ValueError(
                f"data.severities に未知の値がある: {', '.join(map(repr, unknown))}"
                f"（有効値: {', '.join(SEVERITIES)}）"
            )
        if len(set(value)) != len(value):
            raise ValueError("data.severities に重複がある")
        # 記載順の違いで config.yaml の記録が揺れないよう，定義順にそろえる．
        return tuple(s for s in SEVERITIES if s in value)

    @property
    def images_dir(self) -> Path:
        """画像ディレクトリ（`<root>/images`）."""
        return self.root / "images"

    @property
    def manifest_path(self) -> Path:
        """manifest CSV の解決済みパス."""
        return self.manifest if self.manifest.is_absolute() else self.root / self.manifest


class SplitConfig(_StrictModel):
    """患者単位分割の設定.

    Attributes:
        seed: 分割の乱数 seed．
        val_ratio: 単一分割で val に回す患者の比率．
        test_ratio: 単一分割で test に回す患者の比率．
        n_folds: 交差検証の fold 数．
    """

    seed: int = 42
    val_ratio: float = Field(default=0.15, gt=0.0, lt=1.0)
    test_ratio: float = Field(default=0.15, gt=0.0, lt=1.0)
    n_folds: int = Field(default=5, ge=2)

    @model_validator(mode="after")
    def _check_ratio_sum(self) -> SplitConfig:
        # train が残らない比率は分割として成立しないため拒否する．
        if self.val_ratio + self.test_ratio >= 1.0:
            raise ValueError(
                "val_ratio と test_ratio の合計は 1 未満でなければならない"
                f"（現在 {self.val_ratio} + {self.test_ratio}）"
            )
        return self


class ExperimentConfig(_StrictModel):
    """実験設定のルート.

    Attributes:
        data: データの場所．
        split: 分割の設定．
    """

    data: DataConfig
    split: SplitConfig = Field(default_factory=SplitConfig)


def load_config(path: Path, environ: Mapping[str, str] | None = None) -> ExperimentConfig:
    """設定 YAML を読み込み，検証済みの設定を返す.

    環境変数 `TIAS_DATA_ROOT` が空でない値で設定されていれば，YAML の `data.root` より優先する．

    Args:
        path: 設定 YAML のパス．
        environ: 参照する環境変数．None なら `os.environ` を使う．

    Returns:
        検証済みの設定．

    Raises:
        ValueError: YAML のトップレベルがマッピングでない場合．
        pydantic.ValidationError: 設定値が不正な場合．
    """
    env = os.environ if environ is None else environ
    with path.open(encoding="utf-8") as f:
        raw: Any = yaml.safe_load(f)
    if raw is None:
        raw = {}
    if not isinstance(raw, dict):
        raise ValueError(f"設定 YAML のトップレベルはマッピングでなければならない: {path}")

    env_root = env.get(DATA_ROOT_ENV_VAR, "")
    if env_root:
        data = raw.get("data") or {}
        if not isinstance(data, dict):
            raise ValueError(f"設定 YAML の data はマッピングでなければならない: {path}")
        raw = {**raw, "data": {**data, "root": env_root}}

    return ExperimentConfig.model_validate(raw)