import mlflow
from mlflow.entities.span import SpanType
from mlflow.exceptions import MlflowException
from mlflow.genai.judges.builtin import CategoricalRating
from mlflow.genai.scorers import ToolCallCorrectness
from mlflow.tracing.constant import SpanAttributeKey

AVAILABLE_TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "search",
            "description": "Search for information",
            "parameters": {"type": "object", "properties": {"query": {"type": "string"}}},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "summarize",
            "description": "Summarize text",
            "parameters": {"type": "object", "properties": {"max_length": {"type": "integer"}}},
        },
    },
]


def create_trace_with_tool_calls(question: str, tool_calls: list[dict]) -> mlflow.entities.Trace:
    with mlflow.start_span(name="agent") as span:
        span.set_inputs({"question": question})
        with mlflow.start_span(name="llm_call", span_type=SpanType.LLM) as llm_span:
            llm_span.set_attribute(SpanAttributeKey.CHAT_TOOLS, AVAILABLE_TOOLS)
            llm_span.set_inputs({"messages": [{"role": "user", "content": question}]})
            llm_span.set_outputs({"content": "I'll help you with that."})
        for tool in tool_calls:
            with mlflow.start_span(name=tool["name"], span_type=SpanType.TOOL) as tool_span:
                tool_span.set_inputs(tool.get("arguments", {}))
                tool_span.set_outputs(tool.get("output", "result"))
        span.set_outputs("Agent response")
    return mlflow.get_trace(span.trace_id)


trace_search_summarize = create_trace_with_tool_calls(
    "Search for MLflow and summarize",
    [
        {"name": "search", "arguments": {"query": "MLflow"}, "output": "MLflow is an ML platform"},
        {"name": "summarize", "arguments": {"max_length": 100}, "output": "Summary"},
    ],
)

trace_single_search = create_trace_with_tool_calls(
    "Search for MLflow",
    [{"name": "search", "arguments": {"query": "MLflow"}, "output": "MLflow is an ML platform"}],
)

# TEST 1: Exact match, ordered, full expectations - MATCH
scorer = ToolCallCorrectness(should_exact_match=True, should_consider_ordering=True)
result = scorer(
    trace=trace_search_summarize,
    expectations={
        "expected_tool_calls": [
            {"name": "search", "arguments": {"query": "MLflow"}},
            {"name": "summarize", "arguments": {"max_length": 100}},
        ]
    },
)
print(f"TEST 1 - Value: {result.value}, Rationale: {result.rationale}")
assert result.value == CategoricalRating.YES
assert "match" in result.rationale.lower()

# TEST 2: Exact match, ordered, full expectations - WRONG ORDER
scorer = ToolCallCorrectness(should_exact_match=True, should_consider_ordering=True)
result = scorer(
    trace=trace_search_summarize,
    expectations={
        "expected_tool_calls": [
            {"name": "summarize", "arguments": {"max_length": 100}},
            {"name": "search", "arguments": {"query": "MLflow"}},
        ]
    },
)
print(f"TEST 2 - Value: {result.value}, Rationale: {result.rationale}")
assert result.value == CategoricalRating.NO
assert "Position 1" in result.rationale

# TEST 3: Exact match, unordered, full expectations - DIFFERENT ORDER OK
scorer = ToolCallCorrectness(should_exact_match=True, should_consider_ordering=False)
result = scorer(
    trace=trace_search_summarize,
    expectations={
        "expected_tool_calls": [
            {"name": "summarize", "arguments": {"max_length": 100}},
            {"name": "search", "arguments": {"query": "MLflow"}},
        ]
    },
)
print(f"TEST 3 - Value: {result.value}, Rationale: {result.rationale}")
assert result.value == CategoricalRating.YES

# TEST 4: Exact match, partial expectations (names only) - MATCH
scorer = ToolCallCorrectness(should_exact_match=True, should_consider_ordering=False)
result = scorer(
    trace=trace_search_summarize,
    expectations={"expected_tool_calls": [{"name": "search"}, {"name": "summarize"}]},
)
print(f"TEST 4 - Value: {result.value}, Rationale: {result.rationale}")
assert result.value == CategoricalRating.YES

# TEST 5: Exact match, partial expectations - WRONG NAME
scorer = ToolCallCorrectness(should_exact_match=True, should_consider_ordering=False)
result = scorer(
    trace=trace_search_summarize,
    expectations={"expected_tool_calls": [{"name": "search"}, {"name": "translate"}]},
)
print(f"TEST 5 - Value: {result.value}, Rationale: {result.rationale}")
assert result.value == CategoricalRating.NO
assert "translate" in result.rationale

# TEST 6: Exact match, count mismatch
scorer = ToolCallCorrectness(should_exact_match=True, should_consider_ordering=False)
result = scorer(
    trace=trace_single_search,
    expectations={"expected_tool_calls": [{"name": "search"}, {"name": "summarize"}]},
)
print(f"TEST 6 - Value: {result.value}, Rationale: {result.rationale}")
assert result.value == CategoricalRating.NO
assert "Expected 2" in result.rationale

# TEST 7: Exact match without expectations - should raise error
scorer = ToolCallCorrectness(should_exact_match=True)
try:
    scorer(trace=trace_single_search)
    assert False, "Should have raised MlflowException"
except MlflowException as e:
    print(f"TEST 7 - Correctly raised: {e}")
    assert "should_exact_match=True requires expectations" in str(e)

# TEST 8: Fuzzy match with full expectations
scorer = ToolCallCorrectness(should_exact_match=False)
result = scorer(
    trace=trace_search_summarize,
    expectations={
        "expected_tool_calls": [
            {"name": "search", "arguments": {"query": "MLflow"}},
            {"name": "summarize", "arguments": {"max_length": 100}},
        ]
    },
)
print(f"TEST 8 - Value: {result.value}, Rationale: {result.rationale}")
assert result.value == CategoricalRating.YES

# TEST 9: Fuzzy match, ground-truth free
scorer = ToolCallCorrectness(should_exact_match=False)
result = scorer(trace=trace_search_summarize)
print(f"TEST 9 - Value: {result.value}, Rationale: {result.rationale}")

# TEST 10: Fuzzy match with partial expectations
scorer = ToolCallCorrectness(should_exact_match=False)
result = scorer(
    trace=trace_search_summarize,
    expectations={"expected_tool_calls": [{"name": "search"}, {"name": "summarize"}]},
)
print(f"TEST 10 - Value: {result.value}, Rationale: {result.rationale}")
assert result.value == CategoricalRating.YES

print("All 10 tests passed!")