from __future__ import annotations

import io
import threading
from pathlib import Path

from adam.config import ConfigManager
from adam.executor import ToolContext
from adam.registry import ToolSpec
from adam.tools import ddpm_adapter


def _context(root: Path, logs: list[str]) -> ToolContext:
    running = threading.Event()
    running.set()
    return ToolContext(
        root=root,
        job_id="DDPMTEST",
        tool=ToolSpec("ddpm_trainer", "DDPM Trainer", "test", "Training", "train_ddpm"),
        cancel_event=threading.Event(),
        run_event=running,
        progress_callback=lambda *_args: None,
        log_callback=logs.append,
    )


def test_resolution_change_branches_from_pipeline_instead_of_resuming_checkpoint(tmp_path: Path, monkeypatch) -> None:
    trainer_root = tmp_path / "DDPM"
    model = trainer_root / "output" / "Anime"
    checkpoint = model / "checkpoint-40"
    dataset = tmp_path / "dataset"
    for folder in (checkpoint / "unet", model / "unet", model / "scheduler", dataset):
        folder.mkdir(parents=True)
    (trainer_root / "train.py").write_text("# fake trainer", encoding="utf-8")
    (model / "model_index.json").write_text("{}", encoding="utf-8")
    (checkpoint / "optimizer.bin").write_bytes(b"optimizer")
    (checkpoint / "scheduler.bin").write_bytes(b"scheduler")
    (checkpoint / "unet" / "diffusion_pytorch_model.safetensors").write_bytes(b"weights")
    (checkpoint / "unet" / "config.json").write_text('{"sample_size": 64}', encoding="utf-8")
    for index in range(2):
        (dataset / f"{index}.png").write_bytes(b"image")

    ConfigManager(tmp_path).update({"tool_folders": {"ddpm_trainer": str(trainer_root)}})
    monkeypatch.setattr(ddpm_adapter.importlib.util, "find_spec", lambda _name: object())
    commands: list[list[str]] = []

    class FakeProcess:
        stdout = io.StringIO("PROGRESS_JSON:{\"event\":\"done\"}\n")
        returncode = 0

        def poll(self):
            return 0

        def terminate(self):
            self.returncode = -15

    def fake_popen(command, **_kwargs):
        commands.append(command)
        return FakeProcess()

    monkeypatch.setattr(ddpm_adapter.subprocess, "Popen", fake_popen)

    logs: list[str] = []
    result = ddpm_adapter.train_ddpm(
        _context(tmp_path, logs),
        dataset_dir=str(dataset),
        model_name="Anime",
        epochs=5,
        output_dir=str(model),
        resume_from=str(checkpoint),
        resolution=256,
    )

    command = commands[0]
    assert "--pretrained_model_path" in command
    assert str(model.resolve()) in command
    assert "--resume_from_checkpoint" not in command
    assert result["output_folder"] != str(model.resolve())
    assert any("Changing DDPM resolution from 64px to 256px" in line for line in logs)