import time

from ..media import input_frames
from ..media.input import open_image
from ..models.ben2 import infer_alpha, 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"])
    output_kind = parameters.get("output_kind", "foreground")
    if output_kind not in {"foreground", "alpha"}:
        raise ValueError(f"Unknown background removal output: {output_kind}")

    with input_frames(context, parameters) as images:
        context.progress(0.12, f"Loading BEN2 for {device.upper()}")
        if output_kind == "alpha":
            import numpy as np

            if len(images) != 1:
                raise ValueError("Alpha output requires one still image")
            with open_image(images[0]) as source:
                alpha, load_ms, inference_ms = infer_alpha(
                    source, model_directory, device, model_cache=model_manager,
                    cancel_check=context.check_cancelled,
                )
            context.check_cancelled()
            np.save(output_directory / "alpha.npy", alpha, allow_pickle=False)
        elif 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 {"alpha": "alpha.npy"} if output_kind == "alpha" else {"image": "."}