"""Tests for unified workflow service helpers."""

import argparse
from unittest.mock import patch

from model_osmosis.services.unified_workflow import (
    build_build_args_from_protect,
    build_launch_action_namespace,
    build_pipeline_args_from_make,
    evaluate_check_workflow,
)


def test_build_pipeline_args_from_make_cross_merge_sets_force_cross():
    args = argparse.Namespace(
        mode="cross-merge",
        output_type="osmoimg",
        models=["a", "b"],
        source=None,
        target=None,
        output="out.osmoimg",
        name="n",
        sample_data=None,
        auto_grow=True,
        merge_strategy="auto",
        strategy_engine="auto",
        cross_opt_stage="auto",
        weights=None,
        coverage_floor=0.70,
        top_k=2,
        train_gate=True,
        hf_base_model=None,
        onnx_input_shape=None,
        trace_input_shape=None,
        epochs=3,
        seed=42,
        allow_partial_load=False,
        allow_lossy_cast=False,
        dry_run=False,
        report_file=None,
        self_check=True,
        self_check_level="smoke",
    )
    ns = build_pipeline_args_from_make(args)
    assert ns.mode == "merge"
    assert ns.force_cross is True
    assert ns.output_type == "osmoimg"


def test_build_build_args_from_protect_maps_fields():
    args = argparse.Namespace(
        model="foo/bar",
        output="protected.osmoimg",
        name="protected",
        compression_profile="lossless-max",
        trust_sign=True,
        provenance="private",
        repro_mode="strict",
        signer_id="default",
        signer_backend="local-keyring",
    )
    ns = build_build_args_from_protect(args)
    assert ns.command == "build"
    assert ns.model == "foo/bar"
    assert ns.output == "protected.osmoimg"
    assert ns.compression_profile == "lossless-max"


def test_build_launch_action_namespace_run_and_export():
    run_args = argparse.Namespace(
        action="run",
        image="m.osmoimg",
        prompt="hi",
        input=None,
        output_file=None,
        max_tokens=12,
        temperature=0.0,
        include_prompt=False,
        runtime_profile="default",
        device="auto",
        runtime_quantization="none",
        runtime_cache=True,
        require_trust=False,
        trust_profile="standard",
        policy_file=None,
        host="0.0.0.0",
        port=8080,
        warmup=False,
        format="pytorch",
        output=None,
    )
    action, ns = build_launch_action_namespace(run_args)
    assert action == "run"
    assert ns.command == "run"
    assert ns.prompt == "hi"

    export_args = argparse.Namespace(**{**run_args.__dict__, "action": "export", "format": "onnx"})
    action, ns = build_launch_action_namespace(export_args)
    assert action == "export"
    assert ns.command == "export"
    assert ns.format == "onnx"


def test_evaluate_check_workflow_payload_contract():
    with patch("model_osmosis.format.image.OsmoImage.inspect", return_value={
        "runtime": {"payload_codec": "adaptive-zstd-v1", "compression_profile": "lossless-max"}
    }):
        payload = evaluate_check_workflow(
            image="model.osmoimg",
            self_check_level="smoke",
            trust_profile="standard",
            require_trust=False,
            policy_file=None,
            benchmark_runtime=False,
            benchmark_profile="ci_standard",
            runtime_profile="default",
            device="auto",
            runtime_quantization="none",
            run_release_check=False,
            trust_checker=lambda *_args, **_kwargs: {"trust_status": "trusted"},
            self_checker=lambda **_kwargs: {
                "self_check_status": "passed",
                "self_check_level": "smoke",
                "self_check_checks": ["artifact_exists"],
                "self_check_skips": [],
            },
        )
    assert payload["status"] == "PASS"
    assert payload["trust_status"] == "trusted"
    assert payload["payload_codec"] == "adaptive-zstd-v1"
    assert payload["self_check_status"] == "passed"

