"""Unit tests for tail-Softmax removal in the ONNX worker.

這些測試需要 ``onnx``（建小型手工 graph）但**不需要 ktc / toolchain**：
`remove_tail_softmax` 走的是 repo 內 libs/ONNX_Convertor/optimizer_scripts 的
``tools.other.remove_nodes``（純 python + onnx），可以在本機直接驗。

模型結構模擬 staging 實測的情境（.autoflow/06-testing/reports/
tflite-520-softmax-removal-verify-2026-07-06.md）：Gemm(logits) → Softmax(terminal)。
注意：remove_nodes 的 output 重接依賴 value_info 內有 logits 條目（production
flow 由 onnx2onnx_flow 的 shape inference 補齊），因此手工 graph 要自帶 value_info。
"""

from pathlib import Path
import sys

import pytest

onnx = pytest.importorskip("onnx")

from onnx import TensorProto, helper  # noqa: E402  (after importorskip)

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))

from services.workers.onnx.core import remove_tail_softmax  # noqa: E402


def _make_model_with_tail_softmax(softmax_name="tail_softmax"):
    """input -> Gemm(logits) -> Softmax -> output，Softmax 為 terminal。"""
    gemm = helper.make_node(
        "Gemm", inputs=["x", "w"], outputs=["logits"], name="gemm0"
    )
    softmax = helper.make_node(
        "Softmax", inputs=["logits"], outputs=["probs"], name=softmax_name
    )
    graph = helper.make_graph(
        nodes=[gemm, softmax],
        name="tail_softmax_graph",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]),
            helper.make_tensor_value_info("w", TensorProto.FLOAT, [4, 3]),
        ],
        outputs=[helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3])],
        value_info=[
            # logits 的 value_info：remove_nodes 重接 output 時需要（production
            # 由 shape inference 補、這裡手工補）。
            helper.make_tensor_value_info("logits", TensorProto.FLOAT, [1, 3]),
        ],
    )
    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])


def _make_model_without_softmax():
    """input -> Gemm -> Relu -> output，完全沒有 Softmax。"""
    gemm = helper.make_node("Gemm", inputs=["x", "w"], outputs=["logits"], name="gemm0")
    relu = helper.make_node("Relu", inputs=["logits"], outputs=["y"], name="relu0")
    graph = helper.make_graph(
        nodes=[gemm, relu],
        name="no_softmax_graph",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]),
            helper.make_tensor_value_info("w", TensorProto.FLOAT, [4, 3]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3])],
        value_info=[
            helper.make_tensor_value_info("logits", TensorProto.FLOAT, [1, 3]),
        ],
    )
    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])


def _make_model_with_tail_softmax_missing_logits_value_info():
    """同 _make_model_with_tail_softmax、但刻意不給 logits 的 value_info。

    remove_nodes 的 output 重接依賴 value_info；缺失時該 output 會被 Abandon
    → 單一 output 模型變成 0 個 output → guard 必須 raise。
    """
    gemm = helper.make_node("Gemm", inputs=["x", "w"], outputs=["logits"], name="gemm0")
    softmax = helper.make_node(
        "Softmax", inputs=["logits"], outputs=["probs"], name="tail_softmax"
    )
    graph = helper.make_graph(
        nodes=[gemm, softmax],
        name="missing_value_info_graph",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]),
            helper.make_tensor_value_info("w", TensorProto.FLOAT, [4, 3]),
        ],
        outputs=[helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3])],
        # 刻意不給 value_info
    )
    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])


def _make_multi_output_model_with_tail_softmax(with_logits_value_info=True):
    """雙 output 模型：
    branch 1: x -> Gemm(logits) -> Softmax -> probs（graph output 1）
    branch 2: x2 -> Relu -> y（graph output 2、與 Softmax 無關）
    """
    gemm = helper.make_node("Gemm", inputs=["x", "w"], outputs=["logits"], name="gemm0")
    softmax = helper.make_node(
        "Softmax", inputs=["logits"], outputs=["probs"], name="tail_softmax"
    )
    relu = helper.make_node("Relu", inputs=["x2"], outputs=["y"], name="relu0")
    value_info = (
        [helper.make_tensor_value_info("logits", TensorProto.FLOAT, [1, 3])]
        if with_logits_value_info
        else []
    )
    graph = helper.make_graph(
        nodes=[gemm, softmax, relu],
        name="multi_output_graph",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]),
            helper.make_tensor_value_info("w", TensorProto.FLOAT, [4, 3]),
            helper.make_tensor_value_info("x2", TensorProto.FLOAT, [1, 3]),
        ],
        outputs=[
            helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3]),
            helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3]),
        ],
        value_info=value_info,
    )
    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])


def _make_model_with_softmax_output_also_consumed():
    """Softmax 的 output 同時是 graph output、又被另一個節點消費：
    x -> Softmax -> probs（graph output 1）、probs -> Relu -> y（graph output 2）。
    有 children → 非 terminal → 不可移除。
    """
    softmax = helper.make_node("Softmax", inputs=["x"], outputs=["probs"], name="softmax0")
    relu = helper.make_node("Relu", inputs=["probs"], outputs=["y"], name="relu0")
    graph = helper.make_graph(
        nodes=[softmax, relu],
        name="output_also_consumed_graph",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])],
        outputs=[
            helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3]),
            helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3]),
        ],
        value_info=[
            helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3]),
        ],
    )
    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])


def _make_model_with_mid_graph_softmax():
    """input -> Softmax -> Relu -> output，Softmax 不是 terminal（有 children）。"""
    softmax = helper.make_node(
        "Softmax", inputs=["x"], outputs=["probs"], name="mid_softmax"
    )
    relu = helper.make_node("Relu", inputs=["probs"], outputs=["y"], name="relu0")
    graph = helper.make_graph(
        nodes=[softmax, relu],
        name="mid_softmax_graph",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3])],
        value_info=[
            helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3]),
        ],
    )
    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])


def test_remove_tail_softmax_removes_terminal_softmax_and_rewires_logits():
    model = _make_model_with_tail_softmax()

    removed = remove_tail_softmax(model)

    assert removed == ["tail_softmax"]
    # Softmax 已不在 graph
    assert all(n.op_type != "Softmax" for n in model.graph.node)
    # graph output 自動重接成 logits（Softmax 的輸入）
    assert [o.name for o in model.graph.output] == ["logits"]
    # 砍完的模型仍是合法 ONNX
    onnx.checker.check_model(model)


def test_remove_tail_softmax_is_noop_when_no_softmax():
    model = _make_model_without_softmax()
    nodes_before = [(n.name, n.op_type) for n in model.graph.node]
    outputs_before = [o.name for o in model.graph.output]

    removed = remove_tail_softmax(model)

    assert removed == []
    # graph 完全不變（絕不能壞既有 onnx/bie 輸入的 flow）
    assert [(n.name, n.op_type) for n in model.graph.node] == nodes_before
    assert [o.name for o in model.graph.output] == outputs_before


def test_remove_tail_softmax_keeps_mid_graph_softmax():
    # 非 terminal 的 Softmax 不移除：remove_nodes 是 cut-from-node 語意，
    # 對中間層 Softmax 使用會把下游一併切掉。
    model = _make_model_with_mid_graph_softmax()
    nodes_before = [(n.name, n.op_type) for n in model.graph.node]

    removed = remove_tail_softmax(model)

    assert removed == []
    assert [(n.name, n.op_type) for n in model.graph.node] == nodes_before
    assert any(n.op_type == "Softmax" for n in model.graph.node)


def test_remove_tail_softmax_handles_unnamed_softmax_node():
    # 節點沒有名字時要有 fallback 命名（cut_nodes 靠 node.name 比對）。
    model = _make_model_with_tail_softmax(softmax_name="")

    removed = remove_tail_softmax(model)

    assert len(removed) == 1
    assert removed[0]  # 有給 fallback 名稱
    assert all(n.op_type != "Softmax" for n in model.graph.node)
    assert [o.name for o in model.graph.output] == ["logits"]


def test_remove_tail_softmax_raises_when_all_outputs_dropped():
    # logits 缺 value_info → remove_nodes 把唯一的 output Abandon
    # → graph 剩 0 個 output → 必須 raise、不能存壞模型。
    model = _make_model_with_tail_softmax_missing_logits_value_info()

    with pytest.raises(RuntimeError, match="no outputs"):
        remove_tail_softmax(model)


def test_remove_tail_softmax_multi_output_rewires_and_keeps_unrelated_output():
    # 多 output 模型：Softmax 餵的 output 重接成 logits、無關的 output 保留。
    model = _make_multi_output_model_with_tail_softmax(with_logits_value_info=True)

    removed = remove_tail_softmax(model)

    assert removed == ["tail_softmax"]
    assert all(n.op_type != "Softmax" for n in model.graph.node)
    assert {o.name for o in model.graph.output} == {"logits", "y"}
    onnx.checker.check_model(model)


def test_remove_tail_softmax_multi_output_raises_on_silently_dropped_output():
    # 多 output 模型 + logits 缺 value_info：Softmax 餵的 output 被 Abandon、
    # 但另一個 output 還在（graph.output 非空、checker 也會過）——
    # 逐 output 驗證的 guard 必須抓到並 raise（Major-3）。
    model = _make_multi_output_model_with_tail_softmax(with_logits_value_info=False)

    with pytest.raises(RuntimeError, match="silently dropped"):
        remove_tail_softmax(model)


def test_remove_tail_softmax_raises_on_duplicate_node_name():
    # cut_nodes 靠名字比對：病態的同名節點（Softmax 與其他節點撞名）
    # 會導致 remove_nodes 誤砍無關節點 → 必須 raise 而非靜默砍。
    model = _make_model_with_tail_softmax(softmax_name="dup")
    model.graph.node[0].name = "dup"  # gemm0 改成與 Softmax 同名

    with pytest.raises(RuntimeError, match="matches 2 nodes"):
        remove_tail_softmax(model)


def test_remove_tail_softmax_keeps_softmax_whose_output_is_also_consumed():
    # Softmax output 同時是 graph output 又被其他節點消費 → 非 terminal、不移除。
    model = _make_model_with_softmax_output_also_consumed()
    nodes_before = [(n.name, n.op_type) for n in model.graph.node]
    outputs_before = [o.name for o in model.graph.output]

    removed = remove_tail_softmax(model)

    assert removed == []
    assert [(n.name, n.op_type) for n in model.graph.node] == nodes_before
    assert [o.name for o in model.graph.output] == outputs_before