"""Tests for message types in the formalization workflow.

Generated by our friend.
"""

import pytest

from ax_prover.models.files import Location
from ax_prover.models.messages import (
    FEEDBACK_CONSTRUCTORS,
    BuildFailedFeedback,
    BuildSuccessFeedback,
    FeedbackMessage,
    MaxIterationsFeedback,
    MissingTargetTheoremFeedback,
    ProposalMessage,
    ReviewApprovedFeedback,
    ReviewRejectedFeedback,
    SorriesGoalStateFeedback,
    StructuredOutputParsingFailedFeedback,
)


@pytest.fixture
def sample_location():
    """A reusable Location for tests."""
    return Location(name="my_thm", module_path="Test.Module", is_external=False)


class TestProposalMessage:
    """Tests for ProposalMessage content formatting."""

    def test_content_includes_reasoning(self, sample_location):
        """Content starts with reasoning."""
        msg = ProposalMessage(reasoning="My reasoning", code="sorry", location=sample_location)
        assert "Reasoning: My reasoning" in msg.content

    def test_content_includes_code_block(self, sample_location):
        """Code is wrapped in a lean code block."""
        msg = ProposalMessage(
            reasoning="r", code="theorem foo := by sorry", location=sample_location
        )
        assert "```lean" in msg.content
        assert "theorem foo := by sorry" in msg.content

    def test_content_includes_location(self, sample_location):
        """Location's formatted_context appears in content."""
        msg = ProposalMessage(reasoning="r", code="c", location=sample_location)
        assert "Test.Module:my_thm" in msg.content

    def test_content_includes_imports_when_present(self, sample_location):
        """Imports section appears when imports are given."""
        msg = ProposalMessage(
            reasoning="r", code="c", location=sample_location, imports=["Mathlib.Data.Nat"]
        )
        assert "Imports:" in msg.content
        assert "Mathlib.Data.Nat" in msg.content

    def test_content_omits_imports_when_empty(self, sample_location):
        """No Imports section when imports list is empty."""
        msg = ProposalMessage(reasoning="r", code="c", location=sample_location, imports=[])
        assert "Imports:" not in msg.content

    def test_content_includes_opens_when_present(self, sample_location):
        """Opens section appears when opens are given."""
        msg = ProposalMessage(reasoning="r", code="c", location=sample_location, opens=["Nat"])
        assert "Opens:" in msg.content

    def test_has_changes_true_with_code(self, sample_location):
        """has_changes is True when code is non-empty."""
        msg = ProposalMessage(reasoning="r", code="theorem foo := sorry", location=sample_location)
        assert msg.has_changes is True

    def test_has_changes_false_when_empty(self, sample_location):
        """has_changes is False when code, imports, opens are all empty."""
        msg = ProposalMessage(reasoning="r", code="", location=sample_location)
        assert msg.has_changes is False

    def test_no_location(self):
        """ProposalMessage works with location=None."""
        msg = ProposalMessage(reasoning="r", code="c", location=None)
        assert "No location" in msg.content


class TestBuildSuccessFeedback:
    """Tests for BuildSuccessFeedback."""

    def test_content_is_empty(self):
        """Success feedback has empty content (not shown to LLM)."""
        msg = BuildSuccessFeedback()
        assert msg.content == ""

    def test_properties(self):
        """is_success=True, is_terminal=False."""
        msg = BuildSuccessFeedback()
        assert msg.is_success is True
        assert msg.is_terminal is False


class TestBuildFailedFeedback:
    """Tests for BuildFailedFeedback."""

    def test_content_contains_error(self):
        """Content includes 'BUILD FAILED' and the error output."""
        msg = BuildFailedFeedback(error_output="type mismatch at line 5")
        assert "BUILD FAILED" in msg.content
        assert "type mismatch at line 5" in msg.content

    def test_properties(self):
        """is_success=False, is_terminal=False."""
        msg = BuildFailedFeedback(error_output="err")
        assert msg.is_success is False
        assert msg.is_terminal is False


class TestReviewApprovedFeedback:
    """Tests for ReviewApprovedFeedback."""

    def test_properties(self):
        """is_success=True, is_terminal=False."""
        msg = ReviewApprovedFeedback()
        assert msg.is_success is True
        assert msg.is_terminal is False


class TestReviewRejectedFeedback:
    """Tests for ReviewRejectedFeedback."""

    def test_content_contains_feedback(self):
        """Content includes 'REVIEW NEEDS REVISION' and the feedback."""
        msg = ReviewRejectedFeedback(feedback="Statement changed")
        assert "REVIEW NEEDS REVISION" in msg.content
        assert "Statement changed" in msg.content

    def test_properties(self):
        """is_success=False."""
        msg = ReviewRejectedFeedback(feedback="x")
        assert msg.is_success is False


class TestMaxIterationsFeedback:
    """Tests for MaxIterationsFeedback."""

    def test_content_contains_limit(self):
        """Content includes the max iterations number."""
        msg = MaxIterationsFeedback(max_iterations=5)
        assert "5" in msg.content

    def test_is_terminal(self):
        """MaxIterations is a terminal feedback."""
        msg = MaxIterationsFeedback(max_iterations=5)
        assert msg.is_terminal is True
        assert msg.is_success is False


class TestMissingTargetTheoremFeedback:
    """Tests for MissingTargetTheoremFeedback."""

    def test_content_contains_theorem_name(self):
        """Content includes the missing theorem name."""
        msg = MissingTargetTheoremFeedback(theorem_name="my_theorem")
        assert "my_theorem" in msg.content
        assert "MISSING TARGET THEOREM" in msg.content


class TestSorriesGoalStateFeedback:
    """Tests for SorriesGoalStateFeedback."""

    def test_content_contains_count_and_goals(self):
        """Content includes sorry count and goal state."""
        msg = SorriesGoalStateFeedback(sorry_count=2, goal_state_at_sorries="⊢ P")
        assert "2" in msg.content
        assert "⊢ P" in msg.content


class TestStructuredOutputParsingFailedFeedback:
    """Tests for StructuredOutputParsingFailedFeedback."""

    def test_content_contains_error(self):
        """Content includes the error message."""
        msg = StructuredOutputParsingFailedFeedback(error_message="invalid JSON")
        assert "invalid JSON" in msg.content
        assert "STRUCTURED OUTPUT PARSING FAILED" in msg.content


class TestFeedbackConstructors:
    """Tests for the FEEDBACK_CONSTRUCTORS registry."""

    def test_all_subclasses_registered(self):
        """Every FeedbackMessage subclass is in FEEDBACK_CONSTRUCTORS."""
        for cls in FeedbackMessage.__subclasses__():
            feedback_type = cls.model_fields["feedback_type"].default
            assert feedback_type in FEEDBACK_CONSTRUCTORS
            assert FEEDBACK_CONSTRUCTORS[feedback_type] is cls

    def test_lookup_by_feedback_type(self):
        """Can look up concrete class by feedback_type string."""
        assert FEEDBACK_CONSTRUCTORS["build_failed"] is BuildFailedFeedback
        assert FEEDBACK_CONSTRUCTORS["review_approved"] is ReviewApprovedFeedback

    def test_construct_from_type_string(self):
        """Can construct a feedback message from its feedback_type."""
        cls = FEEDBACK_CONSTRUCTORS["build_failed"]
        msg = cls(error_output="some error")
        assert isinstance(msg, BuildFailedFeedback)
        assert "some error" in msg.content
