"""Tests for ProverAgent routing logic.

Generated by our friend.
"""

from unittest.mock import MagicMock, patch

import pytest

from ax_prover.models.files import Location
from ax_prover.models.messages import (
    BuildFailedFeedback,
    BuildSuccessFeedback,
    MaxIterationsFeedback,
    ProposalMessage,
    ReviewApprovedFeedback,
    ReviewRejectedFeedback,
    SorriesGoalStateFeedback,
)
from ax_prover.models.proving import ProverAgentState, TargetItem
from ax_prover.prover.agent import ProverAgent


@pytest.fixture
def sample_item():
    loc = Location(name="thm", module_path="Test.Module", is_external=False)
    return TargetItem(title="Test", location=loc)


@pytest.fixture
def agent():
    """Create a ProverAgent with mocked internals for routing tests."""
    with patch.object(ProverAgent, "__init__", lambda self, *a, **kw: None):
        a = ProverAgent.__new__(ProverAgent)
        a.logger = MagicMock()
        return a


@pytest.fixture
def proposal():
    loc = Location(name="thm", module_path="Test.Module", is_external=False)
    return ProposalMessage(reasoning="r", code="sorry", location=loc)


class TestRouteProposer:
    """Tests for route_proposer routing decisions."""

    def test_returns_end_on_terminal_feedback(self, agent, sample_item, proposal):
        """Terminal feedback (e.g., MaxIterations) routes to 'end'."""
        terminal = MaxIterationsFeedback(max_iterations=5)
        state = ProverAgentState(item=sample_item, messages=[proposal, terminal])
        assert agent.route_proposer(state) == "end"

    def test_returns_retry_when_last_not_proposal(self, agent, sample_item, proposal):
        """When last message is not a ProposalMessage, routes to 'retry'."""
        feedback = BuildFailedFeedback(error_output="err")
        state = ProverAgentState(item=sample_item, messages=[proposal, feedback])
        assert agent.route_proposer(state) == "retry"

    def test_returns_continue_on_valid_proposal(self, agent, sample_item, proposal):
        """When last message is a ProposalMessage, routes to 'continue'."""
        state = ProverAgentState(item=sample_item, messages=[proposal])
        assert agent.route_proposer(state) == "continue"


class TestRouteBuilder:
    """Tests for route_builder routing decisions."""

    def test_returns_continue_on_success(self, agent, sample_item, proposal):
        """BuildSuccessFeedback routes to 'continue'."""
        state = ProverAgentState(item=sample_item, messages=[proposal, BuildSuccessFeedback()])
        assert agent.route_builder(state) == "continue"

    def test_returns_retry_on_failure(self, agent, sample_item, proposal):
        """BuildFailedFeedback routes to 'retry'."""
        state = ProverAgentState(
            item=sample_item, messages=[proposal, BuildFailedFeedback(error_output="err")]
        )
        assert agent.route_builder(state) == "retry"

    def test_returns_retry_on_sorries(self, agent, sample_item, proposal):
        """SorriesGoalStateFeedback routes to 'retry'."""
        state = ProverAgentState(
            item=sample_item,
            messages=[
                proposal,
                SorriesGoalStateFeedback(sorry_count=1, goal_state_at_sorries="⊢ P"),
            ],
        )
        assert agent.route_builder(state) == "retry"


class TestRouteReviewer:
    """Tests for route_reviewer routing decisions."""

    def test_returns_continue_on_approval(self, agent, sample_item, proposal):
        """ReviewApprovedFeedback routes to 'continue'."""
        state = ProverAgentState(item=sample_item, messages=[proposal, ReviewApprovedFeedback()])
        assert agent.route_reviewer(state) == "continue"

    def test_returns_retry_on_rejection(self, agent, sample_item, proposal):
        """ReviewRejectedFeedback routes to 'retry'."""
        state = ProverAgentState(
            item=sample_item,
            messages=[proposal, ReviewRejectedFeedback(feedback="bad")],
        )
        assert agent.route_reviewer(state) == "retry"

    def test_returns_retry_on_non_review_feedback(self, agent, sample_item, proposal):
        """Non-review feedback (e.g., BuildSuccess) routes to 'retry'."""
        state = ProverAgentState(item=sample_item, messages=[proposal, BuildSuccessFeedback()])
        assert agent.route_reviewer(state) == "retry"
