"""把 generate_boxes.py 的产出整理成 YOLO 训练集。"""

import csv
import json
import sys
from pathlib import Path

import pytest

pytest.importorskip("PIL")

import numpy as np  # noqa: E402
import yaml  # noqa: E402
from PIL import Image  # noqa: E402

sys.path.insert(0, "scripts")
import prepare_yolo_dataset as pyd  # noqa: E402

CLASSES = ["grease", "normal", "whitedot"]


@pytest.fixture
def source(tmp_path):
    """模拟 generate_boxes.py 的产出：images/<类别>/ + labels/<类别>/ + data.yaml。"""
    root = tmp_path / "runs" / "boxes"
    rng = np.random.default_rng(0)
    # whitedot 只有 2 张（小类），normal 里有 1 张无框（负样本）
    for class_id, (name, count) in enumerate(zip(CLASSES, (6, 5, 2), strict=True)):
        image_dir = root / "images" / name
        label_dir = root / "labels" / name
        image_dir.mkdir(parents=True)
        label_dir.mkdir(parents=True)
        for i in range(count):
            stem = f"{name}_{i}"
            Image.fromarray(rng.integers(0, 255, (32, 32, 3), dtype=np.uint8)).save(
                image_dir / f"{stem}.png"
            )
            empty = name == "normal" and i == 0
            (label_dir / f"{stem}.txt").write_text(
                "" if empty else f"{class_id} 0.5 0.5 0.2 0.2\n", encoding="utf-8"
            )
    (root / "data.yaml").write_text(
        yaml.safe_dump({"nc": len(CLASSES), "names": dict(enumerate(CLASSES))}), encoding="utf-8"
    )
    return root


def output_of(out, source_stem):
    """通过 name_map.csv 找源文件对应的输出图片——重命名后只能这么找。"""
    with (out / "name_map.csv").open(encoding="utf-8-sig", newline="") as handle:
        for row in csv.DictReader(handle):
            if Path(row["source"]).stem == source_stem:
                return out / "images" / row["name"]
    return None


def run(source, out, *extra):
    argv = ["prepare_yolo_dataset.py", "--source", str(source), "--out", str(out), *extra]
    original, sys.argv = sys.argv, argv
    try:
        return pyd.main()
    finally:
        sys.argv = original


def test_flat_layout_without_split(source, tmp_path):
    """类别目录被拍平、不划分 train/val：images/ 与 labels/ 下直接是文件，一一对应。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0

    for kind in ("images", "labels"):
        assert not [p for p in (out / kind).iterdir() if p.is_dir()]
    images = sorted(p.stem for p in (out / "images").iterdir())
    labels = sorted(p.stem for p in (out / "labels").iterdir())
    assert images == labels
    assert len(images) == 13
    assert json.loads((out / "summary.json").read_text(encoding="utf-8"))["split"] is False


def test_hard_links_share_inode(source, tmp_path):
    """用硬链接而不是复制，不占额外磁盘。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    original = source / "images" / "grease" / "grease_0.png"
    linked = output_of(out, "grease_0")
    assert linked.stat().st_ino == original.stat().st_ino


def test_per_class_counts_in_summary(source, tmp_path):
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    per_class = json.loads((out / "summary.json").read_text(encoding="utf-8"))["per_class"]
    assert per_class == {"grease": 6, "normal": 4, "(无框)": 1, "whitedot": 2}


def test_warns_that_data_is_not_split(source, tmp_path, capsys):
    """train 与 val 指向同一目录，直接训练验证指标会虚高，必须提醒。"""
    assert run(source, tmp_path / "data" / "boxes") == 0
    assert "未划分" in capsys.readouterr().out


def test_empty_labels_kept_as_negatives_by_default(source, tmp_path):
    """无框的图默认保留——YOLO 里它们是压低误检的负样本。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
    assert summary["empty_labels"] == {"normal": {"policy": "keep", "count": 1}}
    assert summary["num_images"] == 13


def test_empty_policy_drop_removes_the_image(source, tmp_path):
    out = tmp_path / "data" / "boxes"
    assert run(source, out, "--empty-policy", "drop") == 0
    summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
    assert summary["num_images"] == 12
    assert summary["empty_labels"]["normal"]["policy"] == "drop"
    assert output_of(out, "normal_0") is None


def test_empty_policy_is_per_class(source, tmp_path):
    """缺陷类的空标签丢掉，normal 的留作负样本——两种处置能共存。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out, "--empty-policy", "drop", "normal=keep") == 0
    summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
    assert summary["empty_labels"]["normal"]["policy"] == "keep"
    assert summary["num_images"] == 13


def test_empty_policy_full_writes_a_whole_image_box(source, tmp_path):
    """整图属性的类别退化成整图框：标签要重写，不能链接源文件（内容不同）。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out, "--empty-policy", "normal=full") == 0

    image = output_of(out, "normal_0")
    label = out / "labels" / f"{image.stem}.txt"
    assert label.read_text(encoding="utf-8").split() == ["1", "0.5", "0.5", "1.0", "1.0"]
    # 重写而不是硬链接，源标签必须还是空的
    assert (source / "labels" / "normal" / "normal_0.txt").read_text(encoding="utf-8") == ""

    # 有了确定类别之后，计入 normal 而不是"(无框)"
    per_class = json.loads((out / "summary.json").read_text(encoding="utf-8"))["per_class"]
    assert per_class["normal"] == 5


@pytest.mark.parametrize(
    "tokens,expected",
    [
        (["drop", "nromal=keep"], "不在 data.yaml"),   # 类别拼错
        (["normal=kep"], "未知策略"),                   # 策略拼错
        (["dorp"], "看不懂"),                           # 全局默认拼错
    ],
)
def test_empty_policy_rejects_typos(source, tmp_path, capsys, tokens, expected):
    """静默忽略拼错的类别名会让策略看起来生效了、其实没有。"""
    assert run(source, tmp_path / "out", "--empty-policy", *tokens) != 0
    assert expected in capsys.readouterr().out


def test_data_yaml_uses_absolute_path_and_source_names(source, tmp_path):
    """data.yaml 要能被 ultralytics 从任意工作目录读到，类别顺序沿用来源。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8"))
    assert Path(meta["path"]).is_absolute()
    assert meta["nc"] == 3
    assert [meta["names"][i] for i in sorted(meta["names"])] == CLASSES
    # 不划分：两者都指向扁平的 images/
    assert meta["train"] == meta["val"] == "images"
    assert (Path(meta["path"]) / meta["train"]).is_dir()


def test_name_collision_after_flattening_is_reported(source, tmp_path, capsys):
    """保留原名时，不同类别目录下的同名文件拍平后会互相覆盖——必须报错而不是静默丢数据。"""
    duplicate = source / "images" / "normal" / "grease_0.png"
    Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(duplicate)
    (source / "labels" / "normal" / "grease_0.txt").write_text("1 0.5 0.5 0.2 0.2\n", encoding="utf-8")

    assert run(source, tmp_path / "data" / "boxes", "--keep-names") != 0
    assert "冲突" in capsys.readouterr().out


def test_refuses_to_clobber_existing_output(source, tmp_path):
    """输出目录非空时要求显式 --overwrite。"""
    out = tmp_path / "data" / "boxes"
    out.mkdir(parents=True)
    (out / "keep.txt").write_text("x", encoding="utf-8")
    assert run(source, out) != 0
    assert (out / "keep.txt").exists()
    assert run(source, out, "--overwrite") == 0
    assert not (out / "keep.txt").exists()


def test_missing_label_is_an_error(source, tmp_path):
    """图片多于标签说明来源目录不完整，不能悄悄跳过。"""
    (source / "labels" / "grease" / "grease_0.txt").unlink()
    assert run(source, tmp_path / "out") != 0


def test_dry_run_writes_nothing(source, tmp_path):
    out = tmp_path / "data" / "boxes"
    assert run(source, out, "--dry-run") == 0
    assert not out.exists()


def _all_outputs(out, kind):
    return list((out / kind).iterdir())


def _add_sample(source, group, stem, suffix, class_id, *, fmt=None, mode="RGB"):
    """往来源目录里加一张指定格式的图 + 一个非空标签。"""
    size = (32, 24)
    array = np.random.default_rng(1).integers(0, 255, (size[1], size[0], 3), dtype=np.uint8)
    image = Image.fromarray(array)
    if mode != "RGB":
        image = image.convert(mode)
    path = source / "images" / group / f"{stem}{suffix}"
    image.save(path, format=fmt)
    (source / "labels" / group / f"{stem}.txt").write_text(
        f"{class_id} 0.5 0.5 0.2 0.2\n", encoding="utf-8"
    )
    return path, image


def test_all_output_images_are_png(source, tmp_path):
    """BMP / JPEG 来源都要落成真正的 PNG——后缀和文件头都得是。"""
    _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP")
    _add_sample(source, "normal", "jpg_img", ".jpg", 1, fmt="JPEG")
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0

    images = _all_outputs(out, "images")
    assert len(images) == 15
    assert {p.suffix for p in images} == {".png"}
    assert all(pyd.is_png(p) for p in images)
    # 标签跟着改名后的图片走，一一对应
    assert sorted(p.stem for p in images) == sorted(p.stem for p in _all_outputs(out, "labels"))


def test_bmp_conversion_is_lossless(source, tmp_path):
    """转换不能缩放或有损压缩，否则归一化的框就不再对应原来的像素。"""
    _, original = _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP")
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0

    converted = output_of(out, "bmp_img")
    with Image.open(converted) as image:
        assert image.size == original.size
        assert np.array_equal(np.asarray(image), np.asarray(original))


def test_png_sources_are_still_hard_linked(source, tmp_path):
    """只有非 PNG 才转换；PNG 照旧硬链接，不占额外磁盘。"""
    bmp, _ = _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP")
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0

    summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
    assert summary["link_modes"] == {"hard": 13, "convert": 1}
    converted = output_of(out, "bmp_img")
    assert converted.stat().st_ino != bmp.stat().st_ino
    assert bmp.read_bytes()[:2] == b"BM"      # 源文件没被动过


def test_png_suffix_with_non_png_content_is_converted(source, tmp_path):
    """按文件头判断而不是后缀：写着 .png 其实是 JPEG 的也要转。"""
    fake, _ = _add_sample(source, "grease", "fake", ".png", 0, fmt="JPEG")
    assert not pyd.is_png(fake)
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    output = output_of(out, "fake")
    assert pyd.is_png(output)


@pytest.mark.parametrize("mode", ["L", "RGBA", "CMYK"])
def test_conversion_handles_pixel_modes(source, tmp_path, mode):
    """灰度 / 带透明度原样保留；PNG 存不下的 CMYK 转成 RGB。"""
    _add_sample(source, "grease", "moded", ".tif", 0, fmt="TIFF", mode=mode)
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    output = output_of(out, "moded")
    with Image.open(output) as image:
        assert image.format == "PNG"
        assert image.mode == ("RGB" if mode == "CMYK" else mode)


def test_same_stem_different_suffix_collides(source, tmp_path, capsys):
    """grease_0.png 与另一目录的 grease_0.bmp 统一后缀后会撞成同一个文件，必须报错。"""
    _add_sample(source, "normal", "grease_0", ".bmp", 1, fmt="BMP")
    assert run(source, tmp_path / "data" / "boxes", "--keep-names") != 0
    assert "冲突" in capsys.readouterr().out


def test_dry_run_reports_conversion_count(source, tmp_path, capsys):
    _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP")
    assert run(source, tmp_path / "out", "--dry-run") == 0
    assert "1 张不是 PNG" in capsys.readouterr().out


def test_classes_txt_matches_data_yaml(source, tmp_path):
    """classes.txt 的行号就是 class_id，必须与 data.yaml 的 names 严格一致。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0

    lines = (out / "classes.txt").read_text(encoding="utf-8").splitlines()
    meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8"))
    assert lines == CLASSES
    assert lines == [meta["names"][i] for i in range(meta["nc"])]
    assert (out / "classes.txt").read_text(encoding="utf-8").endswith("\n")


def test_classes_txt_rejects_names_that_break_line_numbering(tmp_path):
    from datasets.yolo import write_classes_txt

    with pytest.raises(ValueError, match="换行"):
        write_classes_txt(tmp_path / "classes.txt", ["ok", "bad\nname"])
    with pytest.raises(ValueError, match="为空"):
        write_classes_txt(tmp_path / "classes.txt", ["ok", "  "])


def _stems(out):
    return sorted(p.stem for p in _all_outputs(out, "images"))


def test_files_renamed_by_class(source, tmp_path):
    """默认按 <类别>_<序号> 重新编号，序号在类内从 1 连续编到该类总数。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    assert _stems(out) == (
        [f"grease_{i:04d}" for i in range(1, 7)]
        + [f"normal_{i:04d}" for i in range(1, 6)]
        + [f"whitedot_{i:04d}" for i in range(1, 3)]
    )
    assert _stems(out) == sorted(p.stem for p in _all_outputs(out, "labels"))


def test_renamed_file_class_matches_its_label(source, tmp_path):
    """文件名里的类别必须与标签第一列一致，否则改名反倒制造了误导。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    for label in _all_outputs(out, "labels"):
        text = label.read_text(encoding="utf-8").split()
        if text:
            assert label.stem.rsplit("_", 1)[0] == CLASSES[int(text[0])]


def test_name_map_traces_back_to_source(source, tmp_path):
    """name_map.csv 覆盖每一张输出，且源路径真实存在、内容一致。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    with (out / "name_map.csv").open(encoding="utf-8-sig", newline="") as handle:
        rows = list(csv.DictReader(handle))
    assert len(rows) == 13
    for row in rows:
        output = out / "images" / row["name"]
        assert output.read_bytes() == Path(row["source"]).read_bytes()
        assert row["name"].startswith(row["class"] + "_")


def test_names_stable_across_reruns(source, tmp_path):
    """重跑同一份来源，同一张源图拿到同一个名字，已做的复核记录才对得上。"""
    first, second = tmp_path / "a", tmp_path / "b"
    assert run(source, first) == 0
    assert run(source, second) == 0
    for stem in ("grease_0", "grease_5", "normal_3", "whitedot_1"):
        assert output_of(first, stem).name == output_of(second, stem).name


def test_keep_names_preserves_source_stems(source, tmp_path):
    """--keep-names 保留原名，只是后缀统一成 .png。"""
    _add_sample(source, "grease", "原名_中文", ".bmp", 0, fmt="BMP")
    out = tmp_path / "data" / "boxes"
    assert run(source, out, "--keep-names") == 0
    assert "原名_中文" in _stems(out)
    assert output_of(out, "原名_中文").name == "原名_中文.png"


def test_kept_empty_labels_named_after_source_folder(source, tmp_path):
    """空标签没有 class_id，按来源目录名编号；与同类有框的图共用序号，不会撞名。"""
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    assert output_of(out, "normal_0").stem.startswith("normal_")
    assert len([s for s in _stems(out) if s.startswith("normal_")]) == 5


@pytest.fixture
def label_studio_out(source, tmp_path):
    pytest.importorskip("label_studio_sdk.converter.imports.yolo")
    out = tmp_path / "data" / "boxes"
    assert run(source, out, "--empty-policy", "normal=keep", "--label-studio") == 0
    return out


def _tasks(out):
    return json.loads((out / "tasks.json").read_text(encoding="utf-8"))


def test_label_studio_files_at_dataset_root(label_studio_out):
    """tasks.json 与 labeling config 直接放在数据集根目录，不再有 label_studio/ 子目录。"""
    assert (label_studio_out / "tasks.json").is_file()
    assert (label_studio_out / "tasks.label_config.xml").is_file()
    assert not (label_studio_out / "label_studio").exists()


def test_label_studio_tasks_are_predictions_with_urls(label_studio_out):
    """每张图一个任务；框作为 predictions，图片 URL 相对 DOCUMENT_ROOT。"""
    tasks = _tasks(label_studio_out)
    assert len(tasks) == 13
    for task in tasks:
        assert "annotations" not in task
        assert task["data"]["image"].startswith("/data/local-files/?d=images/")
        name = task["data"]["image"].rsplit("/", 1)[1]
        assert (label_studio_out / "images" / name).exists()


def test_label_studio_boxes_round_trip_to_yolo(label_studio_out):
    """Label Studio 的左上角百分比坐标反算回 YOLO 中心点坐标，必须与标签一致。"""
    view = label_studio_out
    checked = 0
    for task in _tasks(label_studio_out):
        name = Path(task["data"]["image"].rsplit("/", 1)[1]).stem
        lines = (view / "labels" / f"{name}.txt").read_text(encoding="utf-8").split()
        results = task["predictions"][0]["result"]
        if not lines:
            assert results == []        # 保留下来的空标签：任务里没有框
            continue
        value = results[0]["value"]
        class_id, cx, cy, w, h = lines
        assert value["rectanglelabels"] == [CLASSES[int(class_id)]]
        assert value["x"] / 100 + value["width"] / 200 == pytest.approx(float(cx))
        assert value["y"] / 100 + value["height"] / 200 == pytest.approx(float(cy))
        assert value["width"] / 100 == pytest.approx(float(w))
        checked += 1
    assert checked == 12


def test_label_studio_config_and_summary(label_studio_out):
    config = (label_studio_out / "tasks.label_config.xml").read_text(encoding="utf-8")
    for name in CLASSES:
        assert f'value="{name}"' in config
    summary = json.loads((label_studio_out / "summary.json").read_text(encoding="utf-8"))
    assert summary["label_studio"]["tasks"] == 13
    assert Path(summary["label_studio"]["document_root"]) == label_studio_out.resolve()


def test_label_studio_not_built_by_default(source, tmp_path):
    out = tmp_path / "data" / "boxes"
    assert run(source, out) == 0
    assert not (out / "tasks.json").exists()
    assert json.loads((out / "summary.json").read_text(encoding="utf-8"))["label_studio"] is None