"""Executor 测试。"""

from unittest.mock import patch

from promptgrad.executor import Executor
from promptgrad.models import TestCase


class TestRenderPrompt:
    def test_simple_render(self):
        executor = Executor(
            prompt_config={"system": "You help.", "template": "Query: {{query}}", "few_shot_examples": []},
            model_config={"name": "test-model", "temperature": 0.3, "max_tokens": 100},
        )
        result = executor.render_prompt("Input: {{query}}, Context: {{context}}", {"query": "hello", "context": "test"})
        assert result == "Input: hello, Context: test"

    def test_render_missing_var_kept(self):
        executor = Executor(
            prompt_config={"system": "You help.", "template": "{{query}}", "few_shot_examples": []},
            model_config={"name": "test-model", "temperature": 0.3, "max_tokens": 100},
        )
        result = executor.render_prompt("{{query}} and {{missing}}", {"query": "hello"})
        assert "hello" in result
        assert "{{missing}}" in result


class TestExecutorRun:
    @patch("promptgrad.executor.call_llm")
    def test_run_single_case(self, mock_call_llm):
        mock_call_llm.return_value = {
            "output": "rewritten query",
            "token_usage": {"input": 50, "output": 20},
            "latency_ms": 100.0,
        }
        executor = Executor(
            prompt_config={"system": "Rewrite.", "template": "Query: {{query}}", "few_shot_examples": []},
            model_config={"name": "test-model", "temperature": 0.3, "max_tokens": 100},
        )
        cases = [TestCase(id="c1", input={"query": "python list"}, tags=[])]
        results = executor.run(cases)
        assert len(results) == 1
        assert results[0].case_id == "c1"
        assert results[0].output == "rewritten query"

    @patch("promptgrad.executor.call_llm")
    def test_run_concurrent(self, mock_call_llm):
        mock_call_llm.return_value = {
            "output": "output",
            "token_usage": {"input": 10, "output": 5},
            "latency_ms": 50.0,
        }
        executor = Executor(
            prompt_config={"system": "Help.", "template": "{{q}}", "few_shot_examples": []},
            model_config={"name": "test-model", "temperature": 0.3, "max_tokens": 100},
            concurrency=3,
        )
        cases = [TestCase(id=f"c{i}", input={"q": f"q{i}"}, tags=[]) for i in range(5)]
        results = executor.run(cases)
        assert len(results) == 5
        assert mock_call_llm.call_count == 5


class TestBuildMessages:
    def test_messages_with_system_and_user(self):
        executor = Executor(
            prompt_config={
                "system": "You are a helper.",
                "template": "Q: {{query}}",
                "few_shot_examples": [],
            },
            model_config={"name": "m", "temperature": 0, "max_tokens": 100},
        )
        msgs = executor._build_messages({"query": "test"})
        assert msgs[0]["role"] == "system"
        assert msgs[0]["content"] == "You are a helper."
        assert msgs[1]["role"] == "user"
        assert msgs[1]["content"] == "Q: test"

    def test_messages_with_few_shot(self):
        executor = Executor(
            prompt_config={
                "system": "Help.",
                "template": "Q: {{query}}",
                "few_shot_examples": [
                    {"user": "Q: example", "assistant": "A: answer"},
                ],
            },
            model_config={"name": "m", "temperature": 0, "max_tokens": 100},
        )
        msgs = executor._build_messages({"query": "real"})
        assert len(msgs) == 4  # system + few_shot_user + few_shot_assistant + user
        assert msgs[1]["role"] == "user"
        assert msgs[2]["role"] == "assistant"
