"""tar ストリームの書き出し（push）と，回収した tar の検査・展開（pull）.

Windows には rsync が無いため，Python の tarfile でストリームを作り `ssh ... 'tar -x ...'` の
標準入力へ流す（push）．pull は ssh の標準出力をいったん一時ファイルに受け取り，
全メンバーを検査してから展開する（ストリームのまま展開すると，問題のメンバーに当たる前の
ファイルが展開済みになってしまうため．SPEC_03「pull」の二重の防御）．
"""

from __future__ import annotations

import tarfile
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO

from tias_judge_model.remote.files import PushEntry
from tias_judge_model.remote.whitelist import check_pull_members

_ENCODING = "utf-8"
_COPY_CHUNK = 1024 * 1024


@dataclass(frozen=True)
class ExtractResult:
    """展開の結果.

    Attributes:
        files: 展開したファイル数．
        bytes: 展開したファイルの合計バイト数．
    """

    files: int
    bytes: int


def _normalize_tarinfo(info: tarfile.TarInfo) -> tarfile.TarInfo:
    """所有者とパーミッションを送信元 OS に依存しない値にそろえる.

    Windows の stat は所有者を持たず実行ビットも拡張子から推測した値なので，
    展開結果がそれに引きずられないよう固定する（実行ビットが立っているものだけ 755）．
    """
    info.uid = info.gid = 0
    info.uname = info.gname = ""
    info.mode = 0o755 if info.isdir() or info.mode & 0o111 else 0o644
    return info


def write_tar_stream(entries: Sequence[PushEntry], fileobj: BinaryIO) -> int:
    """送信対象を非圧縮の tar ストリームとして fileobj へ書き出す.

    PAX 形式で書く（日本語ファイル名を UTF-8 で確実に運ぶ）．fileobj は閉じない
    （呼び出し側が ssh の標準入力を閉じて EOF を伝える）．

    Args:
        entries: 送信対象．
        fileobj: 書き込み先のバイナリストリーム（ssh の標準入力）．

    Returns:
        書き出したファイルの合計バイト数．

    Raises:
        OSError: 書き込み先が閉じられた（ssh が異常終了した）場合など．
    """
    done = 0
    with tarfile.open(
        fileobj=fileobj, mode="w|", format=tarfile.PAX_FORMAT, encoding=_ENCODING, dereference=True
    ) as tar:
        for entry in entries:
            tar.add(entry.path, arcname=entry.arcname, recursive=False, filter=_normalize_tarinfo)
            done += entry.size
    return done


def copy_stream(source: BinaryIO, destination: BinaryIO) -> int:
    """source を終端まで読み，destination へ書き写す.

    Args:
        source: 読み出し元（ssh の標準出力）．
        destination: 書き込み先（一時ファイル）．

    Returns:
        書き写したバイト数．
    """
    total = 0
    while chunk := source.read(_COPY_CHUNK):
        destination.write(chunk)
        total += len(chunk)
    return total


def extract_pull_archive(
    archive: BinaryIO, dest: Path, run_id: str, *, with_checkpoints: bool
) -> ExtractResult:
    """回収した tar（シーク可能なファイル）の全メンバーを検査し，問題が無ければ dest へ展開する.

    検査（`check_pull_members`）で 1 つでも問題があれば，何も展開せずに例外を送出する．
    検査を通ったメンバーは相対パスの通常ファイルのみだが，念のため Python 3.12 の
    `filter="data"`（展開先の外・特殊ファイル・危険なパーミッションを拒否）も重ねる．

    Args:
        archive: tar を書き込んだ一時ファイル（読み出し位置は問わない）．
        dest: 展開先（手元の `runs/`．`<RUN_ID>/...` がその下にできる．同名ファイルは上書き）．
        run_id: 検証済みの RUN_ID．
        with_checkpoints: `checkpoints/*.pt` を受け入れるか．

    Returns:
        展開したファイル数と合計バイト数．

    Raises:
        ValueError: 回収できないメンバーが含まれる場合（何も展開しない）．
        tarfile.TarError: アーカイブが壊れている場合（何も展開しない）．
        OSError: 書き込みに失敗した場合．
    """
    archive.seek(0)
    with tarfile.open(fileobj=archive, mode="r:", encoding=_ENCODING) as tar:
        members = tar.getmembers()
        check_pull_members(members, run_id, with_checkpoints=with_checkpoints)
        dest.mkdir(parents=True, exist_ok=True)
        tar.extractall(dest, members=members, filter="data")
    return ExtractResult(files=len(members), bytes=sum(member.size for member in members))