"""tools/import_studio_export.py: images once in the pool, symlinked per dataset."""
import json
from pathlib import Path

import pytest

from tools import import_studio_export as imp

YAML = "path: /x\ntrain: train/images\nval: val/images\n\nnames:\n  0: loop_all\n  1: pin\n"


def make_export(root: Path, name: str, train=("tpi_1", "tpi_2"), val=("tpi_3",), missing_label=()):
    exp = root / name
    for split, keys in (("train", train), ("val", val)):
        (exp / split / "images").mkdir(parents=True)
        (exp / split / "labels").mkdir(parents=True)
        for k in keys:
            (exp / split / "images" / f"{k}.jpg").write_bytes(b"\xff\xd8" + k.encode() * 20)
            if k not in missing_label:
                (exp / split / "labels" / f"{k}.txt").write_text("0 0.1 0.1 0.2 0.1 0.2 0.2\n" if k != "tpi_2" else "")
    (exp / "dataset.yaml").write_text(YAML)
    (exp / "manifest.csv").write_text("key,tpid,split\n" + "\n".join(f"{k},{k[4:]},train" for k in train) + "\n")
    return exp


def test_import_builds_pool_and_symlinks(tmp_path):
    repo = tmp_path / "repo"; repo.mkdir()
    exp = make_export(tmp_path, "export_a", missing_label=("tpi_3",))
    src = imp.import_export(exp, "2026-10-02_alc", root=repo)
    ds = repo / "data" / "datasets" / "2026-10-02_alc"
    pool = repo / "data" / "pool" / "images"
    assert sorted(p.name for p in pool.iterdir()) == ["tpi_1.jpg", "tpi_2.jpg", "tpi_3.jpg"]
    link = ds / "images" / "train" / "tpi_1.jpg"
    assert link.is_symlink() and link.resolve() == (pool / "tpi_1.jpg").resolve()
    assert (ds / "labels" / "train" / "tpi_1.txt").read_text().startswith("0 0.1")
    assert (ds / "labels" / "val" / "tpi_3.txt").read_text() == ""      # missing label -> background
    assert src["counts"]["train"] == {"images": 2, "labels": 2, "empty": 1}
    assert src["counts"]["val"] == {"images": 1, "labels": 0, "empty": 1}
    assert src["new_pool_images"] == 3 and src["classes"] == {0: "loop_all", 1: "pin"}
    assert "path: ." in (ds / "dataset.yaml").read_text() and (ds / "manifest.csv").exists()
    assert json.loads((ds / "source.json").read_text())["name"] == "2026-10-02_alc"


def test_second_dataset_reuses_the_pool_and_names_are_immutable(tmp_path):
    repo = tmp_path / "repo"; repo.mkdir()
    imp.import_export(make_export(tmp_path, "a"), "2026-10-02_alc", root=repo)
    src = imp.import_export(make_export(tmp_path, "b", train=("tpi_1", "tpi_2", "tpi_9"), val=("tpi_3",)),
                            "2026-10-09_alc", root=repo)
    assert src["new_pool_images"] == 1
    assert len(list((repo / "data" / "pool" / "images").iterdir())) == 4
    with pytest.raises(SystemExit):
        imp.import_export(make_export(tmp_path, "c"), "2026-10-02_alc", root=repo)


def test_gold_import(tmp_path):
    repo = tmp_path / "repo"; repo.mkdir()
    gold = tmp_path / "gold"
    (gold / "images").mkdir(parents=True); (gold / "labels").mkdir()
    (gold / "images" / "tpi_7.jpg").write_bytes(b"\xff\xd8gold")
    (gold / "labels" / "tpi_7.txt").write_text("1 0.5 0.5 0.6 0.5 0.6 0.6\n")
    (gold / "gold_manifest.json").write_text(json.dumps({"keys": ["tpi_7"]}))
    src = imp.import_export(make_export(tmp_path, "a"), "2026-10-02_alc", root=repo, gold_dir=gold)
    assert src["gold"]["images"] == 1 and src["gold"]["manifest_sha256"]
    assert (repo / "data" / "gold" / "images" / "tpi_7.jpg").is_symlink()
    assert (repo / "data" / "gold" / "labels" / "tpi_7.txt").exists()
    assert (repo / "data" / "pool" / "images" / "tpi_7.jpg").exists()