# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import contextlib
import importlib
import re
from pathlib import Path
from types import SimpleNamespace

import torch

from torchtitan.config import ConfigLoader


flux_infer = importlib.import_module("torchtitan.models.flux.inference.infer")


def test_inference_launcher_default_config_loads():
    launcher = (
        Path(__file__).parents[3] / "torchtitan/models/flux/run_infer.sh"
    ).read_text()

    def get_default(name: str) -> str:
        match = re.search(
            rf'^{name}=\$\{{{name}:-"([^"]+)"\}}$', launcher, re.MULTILINE
        )
        assert match is not None
        return match.group(1)

    config = ConfigLoader().load(
        [
            "--module",
            get_default("MODULE"),
            "--config",
            get_default("CONFIG"),
        ]
    )

    assert config.inference is not None


def test_inference_runs_model_in_engine_context(monkeypatch, tmp_path):
    prompts_path = tmp_path / "prompts.txt"
    prompts_path.write_text("a prompt\n")

    context_active = False

    @contextlib.contextmanager
    def spmd_context():
        nonlocal context_active
        assert not context_active
        context_active = True
        try:
            yield
        finally:
            context_active = False

    engine = SimpleNamespace(
        device=torch.device("cpu"),
        model_parts=[object()],
        load_checkpoint=lambda: None,
        parallelism_context=SimpleNamespace(
            activate_spmd=lambda **kwargs: spmd_context()
        ),
        config=SimpleNamespace(debug=SimpleNamespace(spmd_typechecking=False)),
    )
    trainer = SimpleNamespace(
        engine=engine,
        _dtype=torch.float32,
        autoencoder=object(),
        t5_encoder=object(),
        clip_encoder=object(),
    )
    config = SimpleNamespace(
        inference=SimpleNamespace(
            prompts_path=str(prompts_path),
            local_batch_size=1,
            img_size=16,
            save_img_folder="images",
            sampling=SimpleNamespace(
                enable_classifier_free_guidance=False,
                denoising_steps=1,
                classifier_free_guidance_scale=1.0,
            ),
        ),
        checkpointer=SimpleNamespace(load_step=-1),
        tokenizer=SimpleNamespace(build=lambda: object()),
        dump_folder=str(tmp_path),
    )

    monkeypatch.setenv("WORLD_SIZE", "1")
    monkeypatch.setenv("RANK", "0")
    monkeypatch.setattr(flux_infer, "FluxTrainer", lambda config: trainer)
    monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0)
    monkeypatch.setattr(torch.distributed, "destroy_process_group", lambda: None)
    monkeypatch.setattr(flux_infer, "save_image", lambda **kwargs: None)

    def generate_image(**kwargs):
        assert context_active
        return torch.zeros(1, 3, 16, 16)

    monkeypatch.setattr(flux_infer, "generate_image", generate_image)

    flux_infer.inference(config)
    assert not context_active
