import time

from ..media import input_frames
from ..models.ben2 import infer_foreground, infer_foreground_frames


def run(context, parameters):
    """Run BEN2 and write transparent foreground images into the Job directory."""
    total_started = time.perf_counter()
    context.progress(0.02, "Starting Environment")
    context.check_cancelled()

    output_directory = context.directory
    output_directory.mkdir(parents=True, exist_ok=True)
    device = parameters["device"]
    model_manager = context.resource("model_manager")
    model_directory = model_manager.directory(parameters["model"])

    with input_frames(context, parameters) as images:
        context.progress(0.12, f"Loading BEN2 for {device.upper()}")
        if len(images) == 1:
            context.progress(0.35, "Removing background")
            foreground, load_ms, inference_ms = infer_foreground(
                images[0],
                model_directory,
                device,
                model_cache=model_manager,
                cancel_check=context.check_cancelled,
                preserve_input_alpha=True,
            )
            context.check_cancelled()
            context.progress(0.92, "Writing transparent PNG")
            output_path = output_directory / "image.0001.png"
            foreground.convert("RGBA").save(
                output_path,
                format="PNG",
                compress_level=1,
            )
            if not output_path.is_file():
                raise RuntimeError("BEN2 did not write a foreground image")
        else:
            load_ms, inference_ms = infer_foreground_frames(
                images,
                model_directory,
                device,
                output_directory,
                context.progress,
                0.12,
                0.7,
                model_cache=model_manager,
                cancel_check=context.check_cancelled,
                output_prefix="image",
                preserve_input_alpha=True,
            )

    context.check_cancelled()
    total_ms = (time.perf_counter() - total_started) * 1000.0
    frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)"
    message = (
        f"Background removed{frame_text} in {inference_ms:.0f} ms "
        f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
    )
    context.log(message)
    context.progress(1.0, message)
    return {"image": "."}