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

仕様（実装サマリーの設計を仕様として扱う．出力の規則は SPEC_01「実験の出力」）:

- 出力: config.yaml・log.txt・metrics.json・confusion_matrix.csv・
  figures/confusion_matrix_{all,common}.png・split.json・predictions.csv．
  history.csv・learning_curves.png・checkpoints は出さない
- metrics.json: environment.model_name が color_logistic / color_hist_gb，fold ごとに
  selected_params と candidates，best_epoch は None，color_classifier（特徴名を含む），
  timing.feature_extraction_seconds
- 同じ seed なら `cv` と split.json の分割が一致する
- 回収対象（ログ・標準出力を含む）に患者 ID・ファイル名が出ない
- 最後に実験 ID を標準出力へ出す
- `sampler: balanced`・`color` 節の欠落は，実験ディレクトリを作る前に終了コード 1 で止まる
"""

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.color_features import NUM_COLOR_FEATURES, color_feature_names
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
SPLIT_SEED = 3
COLOR_NAME = "cv_color_tiny"
CNN_NAME = "cv_cnn_tiny"
RECOVERABLE = ["config.yaml", "metrics.json", "confusion_matrix.csv", "log.txt"]
C_VALUES = [0.1, 1.0, 10.0]


def _load_script() -> ModuleType:
    spec = importlib.util.spec_from_file_location("make_synthetic_dataset_color_cli", 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


def _base_config(data_root: Path) -> dict:
    return {
        "data": {"root": data_root.as_posix(), "severities": ["strong", "normal"]},
        "split": {"seed": SPLIT_SEED, "n_folds": N_FOLDS},
        "model": {"name": "resnet18", "pretrained": False},
        "train": {
            "image_size": 32,
            "batch_size": 8,
            "max_epochs": 1,
            "warmup_epochs": 0,
            "inner_val_ratio": 0.2,
            "num_workers": 0,
            "amp": False,
        },
    }


def _invoke(argv: list[str]) -> tuple[int, str, str]:
    root_logger = logging.getLogger()
    old_level = root_logger.level
    root_logger.setLevel(logging.INFO)
    out, err = io.StringIO(), io.StringIO()
    handler = logging.StreamHandler(err)
    root_logger.addHandler(handler)
    try:
        with contextlib.redirect_stdout(out):
            rc = cli.main(argv)
    finally:
        root_logger.removeHandler(handler)
        root_logger.setLevel(old_level)
    return rc, out.getvalue(), err.getvalue()


@pytest.fixture(scope="module")
def run(tmp_path_factory: pytest.TempPathFactory) -> dict:
    tmp = tmp_path_factory.mktemp("train_color_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")

    color_config = tmp / f"{COLOR_NAME}.yaml"
    color_config.write_text(
        yaml.safe_dump(
            {
                **_base_config(data_root),
                "color": {
                    "classifier": "logistic",
                    "max_side": 64,
                    "logistic": {"c_values": C_VALUES, "max_iter": 2000},
                },
            }
        ),
        encoding="utf-8",
    )
    cnn_config = tmp / f"{CNN_NAME}.yaml"
    cnn_config.write_text(yaml.safe_dump(_base_config(data_root)), encoding="utf-8")

    color_root, cnn_root = tmp / "runs_color", tmp / "runs_cnn"
    mp = pytest.MonkeyPatch()
    mp.delenv("TIAS_DATA_ROOT", raising=False)
    mp.setenv("TIAS_GIT_COMMIT", GIT)
    try:
        rc, stdout, stderr = _invoke(
            ["color-cv", "--config", str(color_config), "--out-root", str(color_root)]
        )
        cnn_rc, _, _ = _invoke(["cv", "--config", str(cnn_config), "--out-root", str(cnn_root)])
    finally:
        mp.undo()
    run_dirs = [p for p in color_root.iterdir() if p.is_dir()]
    cnn_dirs = [p for p in cnn_root.iterdir() if p.is_dir()]
    return {
        "tmp": tmp,
        "data_root": data_root,
        "rc": rc,
        "stdout": stdout,
        "stderr": stderr,
        "run_dirs": run_dirs,
        "run_dir": run_dirs[0] if run_dirs else None,
        "cnn_rc": cnn_rc,
        "cnn_dir": cnn_dirs[0] if cnn_dirs else None,
        "n_images": 9 * PER_DIR,
    }


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


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_is_last_stdout_line(run: dict) -> None:
    run_id = run["run_dir"].name
    assert re.fullmatch(rf"\d{{8}}-\d{{6}}_{COLOR_NAME}_{GIT}", run_id)
    lines = [ln for ln in run["stdout"].splitlines() if ln.strip()]
    assert lines[-1] == run_id


def test_expected_outputs_exist(run: dict) -> None:
    d = run["run_dir"]
    for name in [*RECOVERABLE, "split.json", "predictions.csv"]:
        assert (d / name).is_file(), name
    for name in ("confusion_matrix_all.png", "confusion_matrix_common.png"):
        assert (d / "figures" / name).is_file(), name
    assert (d / "log.txt").stat().st_size > 0


def test_cnn_only_outputs_are_absent(run: dict) -> None:
    d = run["run_dir"]
    assert not (d / "history.csv").exists()
    assert not (d / "figures" / "learning_curves.png").exists()
    assert not (d / "checkpoints").exists()
    assert not list(d.rglob("*.pt"))


def test_config_yaml_records_command_and_seed(run: dict) -> None:
    doc = yaml.safe_load((run["run_dir"] / "config.yaml").read_text(encoding="utf-8"))
    assert GIT in json.dumps(doc, ensure_ascii=False)
    assert doc["config"]["split"]["seed"] == SPLIT_SEED
    assert doc["config"]["color"]["classifier"] == "logistic"
    assert doc["seeds"]["split"] == SPLIT_SEED
    assert "color-cv" in json.dumps(doc, ensure_ascii=False)


def test_metrics_environment_and_color_section(run: dict) -> None:
    m = _metrics(run)
    assert m["environment"]["model_name"] == "color_logistic"
    assert m["best_epoch"] is None
    cc = m["color_classifier"]
    assert cc["classifier"] == "logistic"
    assert cc["feature_names"] == list(color_feature_names())
    assert len(cc["feature_names"]) == NUM_COLOR_FEATURES
    assert cc["candidates"] == [{"C": c} for c in C_VALUES]
    t = m["timing"]
    assert t["feature_extraction_seconds"] >= 0
    assert t["total_seconds"] >= t["feature_extraction_seconds"]
    assert m["imbalance"] == {"class_weighting": "none", "sampler": "random"}


def test_metrics_folds_have_selection(run: dict) -> None:
    m = _metrics(run)
    assert len(m["folds"]) == N_FOLDS
    for fold in m["folds"]:
        assert fold["best_epoch"] is None
        cands = fold["candidates"]
        assert [c["params"] for c in cands] == [{"C": c} for c in C_VALUES]
        assert fold["selected_params"] in [c["params"] for c in cands]
        best = max(c["val_macro_f1"] for c in cands)
        first_best = next(c for c in cands if c["val_macro_f1"] == best)
        assert fold["selected_params"] == first_best["params"]
        for subset in ("all", "common", "mild"):
            assert {"macro_f1", "accuracy", "per_class", "confusion_matrix"} <= set(
                fold["subsets"][subset]
            )


def test_held_out_evaluated_on_all_severities(run: dict) -> None:
    m = _metrics(run)
    assert sum(f["subsets"]["all"]["n"] for f in m["folds"]) == run["n_images"]
    assert sum(f["subsets"]["mild"]["n"] for f in m["folds"]) == 4 * PER_DIR
    cm = np.array(m["confusion_matrix_sum"]["all"])
    assert cm.shape == (5, 5)
    assert cm.sum() == run["n_images"]


def test_predictions_cover_manifest_with_five_probabilities(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_split_json_matches_cnn_cv_with_same_seed(run: dict) -> None:
    assert run["cnn_rc"] == 0
    color_split = json.loads((run["run_dir"] / "split.json").read_text(encoding="utf-8"))
    cnn_split = json.loads((run["cnn_dir"] / "split.json").read_text(encoding="utf-8"))
    assert color_split == cnn_split
    assert len(color_split["folds"]) == 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_every_output_follows_pull_whitelist(run: dict) -> None:
    d = run["run_dir"]
    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
        else:
            assert rel in {"split.json", "predictions.csv"}, f"想定外の出力: {rel}"
            assert not is_pull_allowed(rel, with_checkpoints=True), rel


def _write(path: Path, doc: dict) -> Path:
    path.write_text(yaml.safe_dump(doc), encoding="utf-8")
    return path


@pytest.mark.parametrize("case", ["balanced", "no_color"])
def test_invalid_color_config_stops_before_creating_run_dir(run: dict, case: str) -> None:
    tmp, data_root = run["tmp"], run["data_root"]
    doc = _base_config(data_root)
    if case == "balanced":
        doc["train"]["sampler"] = "balanced"
        doc["color"] = {"classifier": "logistic"}
    config = _write(tmp / f"cv_bad_{case}.yaml", doc)
    out_root = tmp / f"runs_bad_{case}"
    rc, _, stderr = _invoke(["color-cv", "--config", str(config), "--out-root", str(out_root)])
    assert rc == 1
    assert not out_root.exists() or not any(out_root.iterdir())
    assert SECRET not in stderr


def test_hist_gb_run_records_model_name(run: dict) -> None:
    tmp, data_root = run["tmp"], run["data_root"]
    doc = {
        **_base_config(data_root),
        "color": {
            "classifier": "hist_gb",
            "max_side": None,
            "hist_gb": {"learning_rates": [0.1], "max_depths": [2, None], "max_iter": 10},
        },
    }
    config = _write(tmp / "cv_color_gb_tiny.yaml", doc)
    out_root = tmp / "runs_gb"
    rc, _, _ = _invoke(["color-cv", "--config", str(config), "--out-root", str(out_root)])
    assert rc == 0
    (run_dir,) = [p for p in out_root.iterdir() if p.is_dir()]
    m = json.loads((run_dir / "metrics.json").read_text(encoding="utf-8"))
    assert m["environment"]["model_name"] == "color_hist_gb"
    assert m["color_classifier"]["candidates"] == [
        {"learning_rate": 0.1, "max_depth": 2},
        {"learning_rate": 0.1, "max_depth": None},
    ]
    for fold in m["folds"]:
        assert len(fold["candidates"]) == 2