"""pull の回収ホワイトリスト（SPEC_01「回収ホワイトリスト」）の判定と，受け取った tar の検査.

回収できるのは `runs/<RUN_ID>/` 配下で名前が一致するファイルのみ．拡張子や中身では判断せず，
列挙外はすべて回収しない．判定は 2 か所で使う（二重の防御．SPEC_03「pull」）．

1. サーバの一覧（`build_pull_list_command` の出力）から送らせるファイルを絞り込む（`plan_pull`）
2. 手元で受け取った tar を展開前に全メンバー検査する（`check_pull_members`）

エラーメッセージ・戻り値の表示用情報には件数だけを載せ，ファイル名は出さない
（患者情報を含む名前の可能性があるため）．
"""

from __future__ import annotations

import re
import tarfile
from collections.abc import Iterable
from dataclasses import dataclass

# runs/<RUN_ID>/ 直下で回収するファイル名
PULL_ROOT_FILES = frozenset(
    {
        "config.yaml",
        "history.csv",
        "metrics.json",
        "confusion_matrix.csv",
        "log.txt",
        "dataset_summary.json",
        # TIASjudge で使う判定モデル（tias-train color-export．係数と集計値のみ）
        "color_model.json",
    }
)
PULL_FIGURES_DIR = "figures"
PULL_FIGURE_SUFFIX = ".png"
PULL_CHECKPOINTS_DIR = "checkpoints"
PULL_CHECKPOINT_SUFFIX = ".pt"

# find -printf '%y' の通常ファイル
_FIND_TYPE_FILE = "f"
_WINDOWS_DRIVE_RE = re.compile(r"^[A-Za-z]:")


@dataclass(frozen=True)
class PullPlan:
    """サーバから送らせるファイルの計画.

    Attributes:
        files: 送らせるファイル（`runs/` 基準の POSIX パス．表示しない）．
        skipped: 回収しなかったファイルの件数（リンク等の通常ファイル以外を含む）．
        skipped_checkpoints: そのうち `--with-checkpoints` を付ければ回収できた件数．
    """

    files: tuple[str, ...]
    skipped: int
    skipped_checkpoints: int


def validate_run_id(run_id: str) -> str:
    """RUN_ID（`runs/` 直下のディレクトリ名）を検証して返す.

    Args:
        run_id: 実験 ID．

    Returns:
        検証済みの RUN_ID．

    Raises:
        ValueError: 空・`/`・`\\`・`..`・制御文字を含む，`.`/`-` で始まる場合．
    """
    if (
        not run_id
        or "/" in run_id
        or "\\" in run_id
        or ".." in run_id
        or run_id.startswith((".", "-"))
        or any(ord(c) < 0x20 or c == "\x7f" for c in run_id)
        or run_id != run_id.strip()
    ):
        raise ValueError(
            f"RUN_ID が不正である: {run_id!r}（runs/ 直下のディレクトリ名を指定する．"
            "/・\\・.. は使えない）"
        )
    return run_id


def _is_safe_part(part: str) -> bool:
    """パスの 1 要素が `..`・`.`・空・バックスラッシュ・ドライブ文字を含まないか."""
    return part not in ("", ".", "..") and "\\" not in part and not _WINDOWS_DRIVE_RE.match(part)


def _is_safe_relative(path: str) -> bool:
    """POSIX 区切りの相対パスとして安全か（絶対パス・パス外へ出る要素を含まない）."""
    return not path.startswith("/") and all(_is_safe_part(part) for part in path.split("/"))


def is_pull_allowed(rel_path: str, *, with_checkpoints: bool) -> bool:
    """`runs/<RUN_ID>/` 基準の相対パスが回収ホワイトリストに一致するかを返す.

    Args:
        rel_path: `runs/<RUN_ID>/` 基準の POSIX 相対パス．
        with_checkpoints: `checkpoints/*.pt` も回収対象にするか．

    Returns:
        回収してよければ True．
    """
    if not _is_safe_relative(rel_path):
        return False
    parts = rel_path.split("/")
    if len(parts) == 1:
        return parts[0] in PULL_ROOT_FILES
    if len(parts) != 2:
        return False
    directory, name = parts
    if directory == PULL_FIGURES_DIR:
        return name.endswith(PULL_FIGURE_SUFFIX) and len(name) > len(PULL_FIGURE_SUFFIX)
    if directory == PULL_CHECKPOINTS_DIR and with_checkpoints:
        return name.endswith(PULL_CHECKPOINT_SUFFIX) and len(name) > len(PULL_CHECKPOINT_SUFFIX)
    return False


def _strip_run_prefix(path: str, run_id: str) -> str | None:
    """`<RUN_ID>/<rel>` から `<rel>` を取り出す（RUN_ID 配下でなければ None）."""
    prefix = f"{run_id}/"
    return path[len(prefix) :] if path.startswith(prefix) else None


def parse_find_listing(data: bytes) -> list[tuple[str, str]]:
    """`build_pull_list_command` の出力（`<種類> <パス>` の NUL 区切り）を分解する.

    Linux のファイル名は任意のバイト列なので，往復で失われないよう surrogateescape で扱う．

    Args:
        data: サーバの出力．

    Returns:
        (種類, `runs/` 基準のパス) の一覧．

    Raises:
        ValueError: 書式が崩れている場合（件数のみ示す）．
    """
    entries: list[tuple[str, str]] = []
    for number, record in enumerate((r for r in data.split(b"\0") if r), start=1):
        kind, sep, path = record.decode("utf-8", "surrogateescape").partition(" ")
        if not sep or not kind or not path:
            raise ValueError(f"サーバのファイル一覧の {number} 件目を解釈できない")
        entries.append((kind, path))
    return entries


def plan_pull(
    listing: Iterable[tuple[str, str]], run_id: str, *, with_checkpoints: bool
) -> PullPlan:
    """サーバのファイル一覧から，送らせるファイル（ホワイトリストに一致する通常ファイル）を選ぶ.

    Args:
        listing: `parse_find_listing` の結果．
        run_id: 検証済みの RUN_ID．
        with_checkpoints: `checkpoints/*.pt` も回収するか．

    Returns:
        送らせるファイルと，回収しなかった件数．
    """
    files: list[str] = []
    skipped = 0
    skipped_checkpoints = 0
    for kind, path in listing:
        rel = _strip_run_prefix(path, run_id)
        if kind == _FIND_TYPE_FILE and rel is not None:
            if is_pull_allowed(rel, with_checkpoints=with_checkpoints):
                files.append(path)
                continue
            if not with_checkpoints and is_pull_allowed(rel, with_checkpoints=True):
                skipped_checkpoints += 1
        skipped += 1
    return PullPlan(tuple(sorted(files)), skipped, skipped_checkpoints)


def check_pull_members(
    members: Iterable[tarfile.TarInfo], run_id: str, *, with_checkpoints: bool
) -> int:
    """受け取った tar の全メンバーを展開前に検査する.

    1 つでも問題があれば例外を送出する（呼び出し側は何も展開しない）．受け入れるのは
    `<RUN_ID>/<ホワイトリストに一致する名前>` の通常ファイルのみで，ディレクトリ・リンク
    （シンボリック・ハード）・特殊ファイルは拒否する（サーバはファイル名だけを tar に渡すため，
    正常な応答にディレクトリのメンバーは含まれない）．

    Args:
        members: tar のメンバー（`TarFile.getmembers()`）．
        run_id: 検証済みの RUN_ID．
        with_checkpoints: `checkpoints/*.pt` を受け入れるか．

    Returns:
        受け入れたメンバーの件数．

    Raises:
        ValueError: ホワイトリスト外の名前・パス外に出る名前（`..`・絶対パス・ドライブ文字・
            バックスラッシュ）・リンク・通常ファイル以外が含まれる場合（件数のみを示す）．
    """
    accepted = 0
    outside = 0
    links = 0
    not_regular = 0
    not_allowed = 0
    for member in members:
        if not _is_safe_relative(member.name):
            outside += 1
        elif member.issym() or member.islnk():
            links += 1
        elif not member.isreg():
            not_regular += 1
        else:
            rel = _strip_run_prefix(member.name, run_id)
            if rel is None or not is_pull_allowed(rel, with_checkpoints=with_checkpoints):
                not_allowed += 1
            else:
                accepted += 1
    if outside or links or not_regular or not_allowed:
        raise ValueError(
            "受け取ったアーカイブに回収できないメンバーが含まれるため，何も展開せず中止した"
            f"（ホワイトリスト外 {not_allowed} 件・パス外 {outside} 件・リンク {links} 件・"
            f"通常ファイル以外 {not_regular} 件）"
        )
    return accepted