"""
End-to-end smoke test: synthetic DFF layout -> short CV train -> inference.

Run from the classifier package root:

    cd classifier && conda activate drl && python -m unittest tests.smoke_test -v
"""
import io
import subprocess
import sys
import unittest
from pathlib import Path

from PIL import Image

_CLASSIFIER_ROOT = Path(__file__).resolve().parents[1]
_SMOKE_CFG = _CLASSIFIER_ROOT / "configs" / "smoke" / "smoke.json"


def _write_synthetic_dff(data_root: Path, *, n_identities: int = 16) -> None:
    """Minimal DeepFakeFace-style tree: wiki + three fake sources, shared basenames per identity."""
    sources = ("wiki", "inpainting", "text2img", "insight")
    for i in range(n_identities):
        stem = f"id{i:03d}"
        for src in sources:
            d = data_root / src / f"person_{i:03d}"
            d.mkdir(parents=True, exist_ok=True)
            buf = io.BytesIO()
            Image.new(
                "RGB",
                (96, 96),
                color=(min(20 + i * 11, 255), 80, 120 if src == "wiki" else 40),
            ).save(buf, format="JPEG")
            # Same filename stem across all sources → one CV group per identity (matches DFF).
            (d / f"{stem}.jpg").write_bytes(buf.getvalue())


class SmokeTrainInferTests(unittest.TestCase):
    def test_local_smoke_train_then_inference(self):
        import tempfile

        with tempfile.TemporaryDirectory() as td:
            tmp = Path(td)
            data_dir = tmp / "data"
            out_root = tmp / "outputs"
            _write_synthetic_dff(data_dir)

            sys.path.insert(0, str(_CLASSIFIER_ROOT))
            import run as train_run

            train_run.main(
                str(_SMOKE_CFG),
                data_dir_override=str(data_dir),
                output_root=str(out_root),
                use_gpu=False,
            )

            models_dir = out_root / "models"
            ck_fold0 = models_dir / "smoke_fold0_best.pt"
            if not ck_fold0.is_file():
                ck_fold0 = models_dir / "smoke_fold0_final.pt"
            self.assertTrue(
                ck_fold0.is_file(),
                f"Expected fold-0 checkpoint under {models_dir}",
            )

            probe = tmp / "probe.jpg"
            probe.write_bytes((data_dir / "wiki" / "person_000" / "id000.jpg").read_bytes())

            cmd = [
                sys.executable,
                str(_CLASSIFIER_ROOT / "tools" / "inference.py"),
                str(probe),
                str(_SMOKE_CFG),
                "--checkpoint",
                str(ck_fold0),
            ]
            proc = subprocess.run(
                cmd,
                cwd=str(_CLASSIFIER_ROOT),
                capture_output=True,
                text=True,
                timeout=300,
            )
            self.assertEqual(proc.returncode, 0, proc.stderr + proc.stdout)
            self.assertIn("P(fake)", proc.stdout)


if __name__ == "__main__":
    unittest.main()