"""Tests for PromptOptimizer - BootstrapPromptOptimizer, LLMPromptOptimizer, factory"""

import pytest

from agentkit.evolution.prompt_optimizer import (
    BootstrapPromptOptimizer,
    LLMPromptOptimizer,
    Module,
    PromptOptimizer,
    Signature,
    create_prompt_optimizer,
)


def _make_module(instruction: str = "Find the best result.") -> Module:
    return Module(
        name="test_module",
        signature=Signature(
            input_fields={"query": "search query"},
            output_fields={"result": "search result"},
            instruction=instruction,
        ),
    )


# ── BootstrapPromptOptimizer ───────────────────────────────────


class TestBootstrapPromptOptimizer:
    """测试 BootstrapPromptOptimizer"""

    def test_is_alias_for_prompt_optimizer(self):
        """PromptOptimizer 是 BootstrapPromptOptimizer 的别名"""
        assert PromptOptimizer is BootstrapPromptOptimizer

    @pytest.mark.asyncio
    async def test_not_enough_examples_returns_unchanged(self):
        """样本不足时返回未修改的模块"""
        optimizer = BootstrapPromptOptimizer(min_examples_for_optimization=3)
        optimizer.add_example({"q": "1"}, {"a": "1"}, 0.9)

        result = await optimizer.optimize(_make_module())
        assert result.name == "test_module"  # Unchanged

    @pytest.mark.asyncio
    async def test_enough_examples_produces_optimized_module(self):
        """足够样本时产生优化模块"""
        optimizer = BootstrapPromptOptimizer(max_demos=3, min_examples_for_optimization=2)
        for i in range(3):
            optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)

        result = await optimizer.optimize(_make_module())
        assert result.name == "test_module_optimized"
        assert len(result.demos) == 3

    @pytest.mark.asyncio
    async def test_failure_examples_add_avoid_patterns(self):
        """失败样本添加避免模式到指令中"""
        optimizer = BootstrapPromptOptimizer(min_examples_for_optimization=1)
        optimizer.add_example({"q": "good"}, {"a": "good"}, 0.9)
        optimizer.add_example({"bad_input": "bad"}, {"a": "bad"}, 0.3)

        result = await optimizer.optimize(_make_module())
        assert "Avoid these patterns" in result.signature.instruction

    def test_example_count(self):
        """example_count 返回正确的成功/失败数"""
        optimizer = BootstrapPromptOptimizer()
        optimizer.add_example({"q": "1"}, {"a": "1"}, 0.9)
        optimizer.add_example({"q": "2"}, {"a": "2"}, 0.3)
        optimizer.add_example({"q": "3"}, {"a": "3"}, 0.8)

        success, failure = optimizer.example_count
        assert success == 2
        assert failure == 1


# ── LLMPromptOptimizer ─────────────────────────────────────────


class MockLLMResponse:
    """Mock LLM response"""
    def __init__(self, content: str):
        self.content = content


class MockLLMGateway:
    """Mock LLM Gateway"""
    def __init__(self, response_content: str = "Improved instruction for better results."):
        self._response = response_content
        self.chat_called = False

    async def chat(self, messages, model="default", agent_name="", task_type=""):
        self.chat_called = True
        return MockLLMResponse(self._response)


class FailingLLMGateway:
    """LLM Gateway that always fails"""
    async def chat(self, messages, **kwargs):
        raise RuntimeError("LLM unavailable")


@pytest.mark.asyncio
async def test_llm_optimizer_generates_improved_instruction():
    """LLMPromptOptimizer 生成改进的指令"""
    gateway = MockLLMGateway()
    optimizer = LLMPromptOptimizer(llm_gateway=gateway)

    # Add enough examples for bootstrap post-processing
    for i in range(3):
        optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)

    module = _make_module()
    result = await optimizer.optimize(module)

    assert result.name == "test_module_optimized"
    assert result.signature.instruction == "Improved instruction for better results."
    assert gateway.chat_called is True


@pytest.mark.asyncio
async def test_llm_optimizer_falls_back_to_bootstrap_on_failure():
    """LLM 调用失败时回退到 BootstrapPromptOptimizer"""
    gateway = FailingLLMGateway()
    optimizer = LLMPromptOptimizer(llm_gateway=gateway)

    # Add enough examples for bootstrap fallback
    for i in range(3):
        optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)

    module = _make_module()
    result = await optimizer.optimize(module)

    # Should fall back to bootstrap optimization
    assert result.name == "test_module_optimized"
    assert len(result.demos) == 3


@pytest.mark.asyncio
async def test_llm_optimizer_with_reflection_context():
    """LLMPromptOptimizer 传递反思上下文"""
    from agentkit.evolution.reflector import Reflection

    gateway = MockLLMGateway()
    optimizer = LLMPromptOptimizer(llm_gateway=gateway)

    for i in range(3):
        optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)

    reflection = Reflection(
        task_id="test-001",
        agent_name="test_agent",
        outcome="failure",
        quality_score=0.3,
        patterns=["slow_execution"],
        insights=["Low quality score"],
        suggestions=["Optimize prompt"],
    )

    module = _make_module()
    result = await optimizer.optimize(module, trace=None, reflection=reflection)

    assert result.name == "test_module_optimized"
    assert gateway.chat_called is True


@pytest.mark.asyncio
async def test_llm_optimizer_empty_response_falls_back():
    """LLM 返回空响应时回退到 bootstrap"""
    gateway = MockLLMGateway(response_content="  ")
    optimizer = LLMPromptOptimizer(llm_gateway=gateway)

    for i in range(3):
        optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)

    module = _make_module()
    result = await optimizer.optimize(module)

    # Should fall back to bootstrap
    assert result.name == "test_module_optimized"


def test_llm_optimizer_example_count():
    """LLMPromptOptimizer 的 example_count 委托给 bootstrap"""
    optimizer = LLMPromptOptimizer(llm_gateway=MockLLMGateway())
    optimizer.add_example({"q": "1"}, {"a": "1"}, 0.9)
    optimizer.add_example({"q": "2"}, {"a": "2"}, 0.3)

    success, failure = optimizer.example_count
    assert success == 1
    assert failure == 1


# ── Factory function ───────────────────────────────────────────


class TestCreatePromptOptimizer:
    """测试 create_prompt_optimizer 工厂函数"""

    def test_bootstrap_type(self):
        """bootstrap 类型返回 BootstrapPromptOptimizer"""
        optimizer = create_prompt_optimizer("bootstrap")
        assert isinstance(optimizer, BootstrapPromptOptimizer)

    def test_llm_type_with_gateway(self):
        """llm 类型有 gateway 时返回 LLMPromptOptimizer"""
        gateway = MockLLMGateway()
        optimizer = create_prompt_optimizer("llm", llm_gateway=gateway)
        assert isinstance(optimizer, LLMPromptOptimizer)

    def test_llm_type_without_gateway_falls_back(self):
        """llm 类型无 gateway 时回退到 BootstrapPromptOptimizer"""
        optimizer = create_prompt_optimizer("llm", llm_gateway=None)
        assert isinstance(optimizer, BootstrapPromptOptimizer)

    def test_auto_type_with_gateway(self):
        """auto 类型有 gateway 时返回 LLMPromptOptimizer"""
        gateway = MockLLMGateway()
        optimizer = create_prompt_optimizer("auto", llm_gateway=gateway)
        assert isinstance(optimizer, LLMPromptOptimizer)

    def test_auto_type_without_gateway(self):
        """auto 类型无 gateway 时返回 BootstrapPromptOptimizer"""
        optimizer = create_prompt_optimizer("auto", llm_gateway=None)
        assert isinstance(optimizer, BootstrapPromptOptimizer)

    def test_kwargs_passed_through(self):
        """额外参数传递给优化器"""
        optimizer = create_prompt_optimizer("bootstrap", max_demos=3, min_examples_for_optimization=2)
        assert optimizer._max_demos == 3
        assert optimizer._min_examples == 2