"""tias-train cv の結合テスト.

仕様: SPEC_01「CLI」「実験の出力」「回収ホワイトリスト」・SPEC_02「出力」．

- runs/<実験ID>/ に SPEC_02「出力」の表のファイルが揃う
- metrics.json に部分集合ごとの指標・mild の内訳・所要時間・実行環境がある
- history.csv にエポックの所要秒数の列がある
- 実験 ID は `YYYYmmdd-HHMMSS_<設定ファイル名>_<git短縮ハッシュ>`．標準出力の最終行が実験 ID
- 回収対象の出力（ログを含む）に患者 ID・ファイル名が出ない
- split.json・predictions.csv・checkpoints は回収ホワイトリストの対象外（checkpoints は明示時のみ）
- 「ややあり除外」でも held_out は全程度で評価される

合成データ（scripts/make_synthetic_dataset.py）のファイル名を目立つ文字列に付け替えて使う．
"""

from __future__ import annotations

import contextlib
import csv
import importlib.util
import io
import json
import logging
import re
import sys
from pathlib import Path
from types import ModuleType

import numpy as np
import pytest
import yaml

import tias_judge_model.cli.train as cli
from tias_judge_model.data.manifest import build_manifest_records, write_manifest
from tias_judge_model.remote.whitelist import is_pull_allowed

pytestmark = pytest.mark.slow

SCRIPT_PATH = Path(__file__).resolve().parents[2] / "scripts" / "make_synthetic_dataset.py"
SECRET = "SECRETPT"
GIT = "abc1234"
N_FOLDS = 5
PER_DIR = 6
CONFIG_NAME = "cv_tiny_nomild"
RECOVERABLE = ["config.yaml", "metrics.json", "history.csv", "confusion_matrix.csv", "log.txt"]


def _load_script() -> ModuleType:
    spec = importlib.util.spec_from_file_location("make_synthetic_dataset_train", SCRIPT_PATH)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


@pytest.fixture(scope="module")
def run(tmp_path_factory: pytest.TempPathFactory) -> dict:
    tmp = tmp_path_factory.mktemp("train_cli")
    data_root = tmp / "dataset"
    _load_script().generate(data_root, per_dir=PER_DIR, seed=0)
    for n, png in enumerate(sorted((data_root / "images").rglob("*.png"))):
        png.rename(png.with_name(f"{SECRET}{n:04d}.png"))
    write_manifest(build_manifest_records(data_root / "images"), data_root / "manifest.csv")

    config = tmp / f"{CONFIG_NAME}.yaml"
    config.write_text(
        yaml.safe_dump(
            {
                "data": {"root": data_root.as_posix(), "severities": ["strong", "normal"]},
                "split": {"seed": 3, "n_folds": N_FOLDS},
                "model": {"name": "resnet18", "pretrained": False},
                "train": {
                    "image_size": 32,
                    "batch_size": 8,
                    "max_epochs": 2,
                    "warmup_epochs": 1,
                    "inner_val_ratio": 0.2,
                    "num_workers": 0,
                    "amp": False,
                },
            }
        ),
        encoding="utf-8",
    )
    out_root = tmp / "runs"

    mp = pytest.MonkeyPatch()
    mp.delenv("TIAS_DATA_ROOT", raising=False)
    mp.setenv("TIAS_GIT_COMMIT", GIT)
    root_logger = logging.getLogger()
    old_level = root_logger.level
    root_logger.setLevel(logging.INFO)
    out, err = io.StringIO(), io.StringIO()
    stderr_handler = logging.StreamHandler(err)
    root_logger.addHandler(stderr_handler)
    try:
        with contextlib.redirect_stdout(out):
            rc = cli.main(["cv", "--config", str(config), "--out-root", str(out_root)])
    finally:
        root_logger.removeHandler(stderr_handler)
        root_logger.setLevel(old_level)
        mp.undo()
    run_dirs = [p for p in out_root.iterdir() if p.is_dir()]
    return {
        "rc": rc,
        "stdout": out.getvalue(),
        "stderr": err.getvalue(),
        "run_dirs": run_dirs,
        "run_dir": run_dirs[0] if run_dirs else None,
        "n_images": 9 * PER_DIR,
    }


def test_exit_code_and_single_run_dir(run: dict) -> None:
    assert run["rc"] == 0
    assert len(run["run_dirs"]) == 1


def test_run_id_format_and_last_stdout_line(run: dict) -> None:
    run_id = run["run_dir"].name
    assert re.fullmatch(rf"\d{{8}}-\d{{6}}_{CONFIG_NAME}_{GIT}", run_id)
    lines = [ln for ln in run["stdout"].splitlines() if ln.strip()]
    assert lines[-1] == run_id


def test_output_files_exist(run: dict) -> None:
    d = run["run_dir"]
    for name in [*RECOVERABLE, "split.json", "predictions.csv"]:
        assert (d / name).is_file(), name
    pngs = list((d / "figures").glob("*.png"))
    assert pngs
    names = " ".join(p.name for p in pngs)
    assert "confusion" in names  # 混同行列
    assert len(pngs) >= 2  # 学習曲線 ＋ 混同行列
    ckpts = sorted(p.name for p in (d / "checkpoints").glob("*.pt"))
    assert ckpts == [f"fold{k}.pt" for k in range(N_FOLDS)]
    assert (d / "log.txt").stat().st_size > 0


def test_config_yaml_records_git_and_seed(run: dict) -> None:
    doc = yaml.safe_load((run["run_dir"] / "config.yaml").read_text(encoding="utf-8"))
    text = json.dumps(doc, ensure_ascii=False)
    assert GIT in text
    # 解決済み設定（model・train を含む）と seed が記録される
    assert doc["config"]["split"]["seed"] == 3
    assert doc["config"]["model"]["name"] == "resnet18"
    assert doc["config"]["train"]["image_size"] == 32
    assert doc["seeds"]["split"] == 3


def _metrics(run: dict) -> dict:
    return json.loads((run["run_dir"] / "metrics.json").read_text(encoding="utf-8"))


def test_metrics_has_subsets_per_fold_and_mean_std(run: dict) -> None:
    m = _metrics(run)
    assert len(m["folds"]) == N_FOLDS
    for fold in m["folds"]:
        assert set(fold["subsets"]) >= {"all", "common", "mild"}
        assert fold["best_epoch"] >= 1
        for subset in ("all", "common", "mild"):
            s = fold["subsets"][subset]
            assert {"macro_f1", "accuracy", "per_class", "confusion_matrix"} <= set(s)
    for subset in ("all", "common", "mild"):
        ms = m["mean_std"][subset]
        assert {"mean", "std"} <= set(ms["macro_f1"])
        assert {"mean", "std"} <= set(ms["accuracy"])


def test_held_out_evaluated_on_all_severities(run: dict) -> None:
    # severities=[strong, normal] でも held_out は全程度（合計 = manifest 全件）
    m = _metrics(run)
    n_all = sum(f["subsets"]["all"]["n"] for f in m["folds"])
    n_mild = sum(f["subsets"]["mild"]["n"] for f in m["folds"])
    assert n_all == run["n_images"]
    assert n_mild == 4 * PER_DIR  # mild は淡紅以外の 4 クラス


def test_confusion_matrix_sum(run: dict) -> None:
    m = _metrics(run)
    all_cm = np.array(m["confusion_matrix_sum"]["all"])
    common_cm = np.array(m["confusion_matrix_sum"]["common"])
    assert all_cm.shape == (5, 5)
    assert all_cm.sum() == run["n_images"]
    assert common_cm.sum() == run["n_images"] - 4 * PER_DIR
    # 行（正解）の合計はクラスごとの枚数（淡紅は normal のみ，他は strong + mild）
    assert all_cm.sum(axis=1).tolist() == [
        2 * PER_DIR,
        PER_DIR,
        2 * PER_DIR,
        2 * PER_DIR,
        2 * PER_DIR,
    ]

    with (run["run_dir"] / "confusion_matrix.csv").open(encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    subsets = {r["subset"] for r in rows}
    assert subsets == {"common", "all"}
    all_rows = [r for r in rows if r["subset"] == "all"]
    assert [r["true"] for r in all_rows] == ["pale", "light_red", "red", "dark_red", "purple"]
    csv_matrix = [[int(v) for k, v in r.items() if k.startswith("pred_")] for r in all_rows]
    assert csv_matrix == all_cm.tolist()


def test_mild_prediction_breakdown_present(run: dict) -> None:
    m = _metrics(run)
    breakdown = m["pooled"]["mild"]["prediction_breakdown"]
    overall = breakdown["overall"]
    assert overall["n"] == 4 * PER_DIR
    total = overall["light_red"] + overall["original"] + overall["other"]
    assert total == pytest.approx(1.0)
    assert set(breakdown["per_class"]) == {"pale", "red", "dark_red", "purple"}
    for fold in m["folds"]:
        assert "prediction_breakdown" in fold["subsets"]["mild"]


def test_metrics_has_timing_and_environment(run: dict) -> None:
    m = _metrics(run)
    t = m["timing"]
    assert t["total_seconds"] > 0
    assert t["data_loading_seconds"] >= 0
    assert t["total_seconds"] >= t["data_loading_seconds"]
    for fold in m["folds"]:
        ft = fold["timing"]
        assert ft["train_seconds"] >= 0
        assert ft["eval_seconds"] >= 0
        assert 1 <= ft["epochs_run"] <= 2
        assert ft["mean_epoch_seconds"] >= 0
    env = m["environment"]
    assert env["device"]
    assert env["model_name"] == "resnet18"
    assert env["num_parameters"] > 0


def test_history_has_epoch_seconds(run: dict) -> None:
    with (run["run_dir"] / "history.csv").open(encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    assert rows
    cols = set(rows[0])
    assert {"fold", "epoch", "train_loss", "val_loss", "val_macro_f1", "lr"} <= cols
    seconds_cols = [c for c in cols if "second" in c]
    assert seconds_cols
    assert all(float(r[seconds_cols[0]]) >= 0 for r in rows)
    assert {int(r["fold"]) for r in rows} == set(range(N_FOLDS))


def test_recoverable_outputs_have_no_patient_ids(run: dict) -> None:
    d = run["run_dir"]
    for name in RECOVERABLE:
        assert SECRET not in (d / name).read_text(encoding="utf-8"), name
    for png in (d / "figures").glob("*.png"):
        assert SECRET.encode() not in png.read_bytes(), png.name
    assert SECRET not in run["stdout"]
    assert SECRET not in run["stderr"]


def test_split_and_predictions_contain_ids_and_are_not_pulled(run: dict) -> None:
    d = run["run_dir"]
    # 個票（患者 ID を含む）はサーバ上にだけ置き，回収しない
    assert SECRET in (d / "split.json").read_text(encoding="utf-8")
    assert SECRET in (d / "predictions.csv").read_text(encoding="utf-8")
    for name in ("split.json", "predictions.csv"):
        assert not is_pull_allowed(name, with_checkpoints=False)
        assert not is_pull_allowed(name, with_checkpoints=True)


def test_predictions_cover_all_severities(run: dict) -> None:
    with (run["run_dir"] / "predictions.csv").open(encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    assert len(rows) == run["n_images"]
    assert {r["severity"] for r in rows} == {"strong", "mild", "normal"}


def test_every_output_matches_pull_whitelist_as_spec(run: dict) -> None:
    d = run["run_dir"]
    recoverable = set(RECOVERABLE)
    for path in d.rglob("*"):
        if not path.is_file():
            continue
        rel = path.relative_to(d).as_posix()
        if rel in recoverable or (rel.startswith("figures/") and rel.endswith(".png")):
            assert is_pull_allowed(rel, with_checkpoints=False), rel
        elif rel.startswith("checkpoints/") and rel.endswith(".pt"):
            assert not is_pull_allowed(rel, with_checkpoints=False), rel
            assert is_pull_allowed(rel, with_checkpoints=True), rel
        else:
            assert not is_pull_allowed(rel, with_checkpoints=True), rel
            assert rel in {"split.json", "predictions.csv"}, f"想定外の出力: {rel}"