from ..geometry import build_surface, foreground_mesh_mask
from ..media import copy_color_images
from ..models import geometry_adapter, moge2
from ..outputs.surface import write_normal_texture, write_surface_archive
from .common import input_frames
from .foreground import write_optional


def inference_progress_message(model_name, index, total):
    message = f"Estimating {model_name} geometry"
    if total > 1:
        message += f" ({index + 1}/{total})"
    return message


def run(context, parameters):
    context.progress(0.02, "Preparing surface input")
    context.directory.mkdir(parents=True, exist_ok=True)
    model_manager = context.resource("model_manager")
    adapter = geometry_adapter(parameters["model"])
    with input_frames(context, parameters) as images:
        if len(images) != 1:
            raise ValueError("Convert to Surface currently supports one image")
        context.progress(0.12, "Estimating surface geometry")
        if adapter is moge2:
            prediction = adapter.infer(
                model_manager,
                parameters,
                images,
                include_points=True,
                progress_callback=lambda index, total: context.progress(
                    0.12 + 0.58 * index / total,
                    inference_progress_message("MoGe-2", index, total),
                ),
                cancel_check=context.check_cancelled,
            )
        else:
            prediction = adapter.infer(
                model_manager,
                parameters,
                images,
                progress_callback=lambda index, total: context.progress(
                    0.12 + 0.58 * index / total,
                    inference_progress_message("DA3", index, total),
                ),
                cancel_check=context.check_cancelled,
            )
        context.check_cancelled()
        frame = adapter.surface_frame(prediction)
        color_paths = write_optional(
            context,
            parameters,
            images,
            progress_start=0.72,
            progress_span=0.08,
        )
        foreground_mask = None
        if color_paths:
            foreground_mask = foreground_mesh_mask(
                color_paths[0],
                frame.depth.shape,
            )
        context.progress(0.82, "Building surface")
        mesh = build_surface(frame, foreground_mask=foreground_mask)
        context.progress(0.92, "Writing surface")
        write_surface_archive(mesh, context.directory / "mesh.npz")
        if not color_paths:
            copy_color_images(images, context.directory)
        if adapter is moge2:
            write_normal_texture(
                prediction.normal[0],
                context.directory / "normal.png",
            )
    context.check_cancelled()
    context.progress(1.0, "Surface is ready")
    return {"surface": "mesh.npz"}