"""
Regression tests for LLM Stream Parser Module.

These tests ensure the refactored parser maintains the exact same behavior as the
original inline implementation in agentic_edit.py.

Test categories:
1. Plain text handling (buffer retention behavior)
2. <thinking> block parsing (including cross-chunk)
3. Tool tag parsing (including cross-chunk and multiple tools)
4. Unknown tags handling
5. Unterminated blocks (RetryEvent generation)
6. reasoning_content metadata handling
"""

import pytest
from typing import Any, Generator, List, Optional, Tuple, Type
from unittest.mock import MagicMock
from dataclasses import dataclass

from autocoder.common.v2.agent.llm_stream_parser import (
    RegexToolXmlParser,
    StreamingTagParser,
    ReasoningMetadataEmitter,
    LLMResponseStreamParser,
    create_llm_stream_parser,
)
from autocoder.common.v2.agent.agentic_edit_types import (
    BaseTool,
    ErrorEvent,
    LLMOutputEvent,
    LLMReasoningEvent,
    LLMThinkingEvent,
    RetryEvent,
    TokenUsageEvent,
    ToolCallEvent,
)


# ============================================================================
# Test Fixtures and Helpers
# ============================================================================


class MockReadFileTool(BaseTool):
    """Mock tool for testing."""

    path: str
    start_line: Optional[str] = None
    end_line: Optional[str] = None


class MockExecuteCommandTool(BaseTool):
    """Mock tool for testing."""

    command: str
    requires_approval: bool = True


class MockListFilesTool(BaseTool):
    """Mock tool for testing."""

    path: str
    recursive: Optional[bool] = False


class MockAskFollowupQuestionTool(BaseTool):
    """Mock tool for testing."""

    question: str
    options: Optional[List[str]] = None


# Test tool model map
TEST_TOOL_MODEL_MAP = {
    "read_file": MockReadFileTool,
    "execute_command": MockExecuteCommandTool,
    "list_files": MockListFilesTool,
    "ask_followup_question": MockAskFollowupQuestionTool,
}


def mock_reconstruct_tool_xml(tool: BaseTool) -> str:
    """Mock XML reconstructor for testing."""
    tool_name = type(tool).__name__.replace("Mock", "").lower()
    # Map back to tag names
    tag_map = {
        "readfiletool": "read_file",
        "executecommandtool": "execute_command",
        "listfilestool": "list_files",
        "askfollowupquestiontool": "ask_followup_question",
    }
    tag = tag_map.get(type(tool).__name__.lower(), "unknown")
    params = tool.model_dump(exclude_none=True)
    param_xml = "\n".join(f"<{k}>{v}</{k}>" for k, v in params.items())
    return f"<{tag}>\n{param_xml}\n</{tag}>"


def make_generator(
    chunks: List[Tuple[str, Any]],
) -> Generator[Tuple[str, Any], None, None]:
    """Create a generator from a list of (content, metadata) tuples."""
    for chunk in chunks:
        yield chunk


@dataclass
class MockMetadata:
    """Mock metadata object for testing reasoning_content."""

    reasoning_content: Optional[str] = None


def collect_events(gen) -> List:
    """Collect all events from a generator."""
    return list(gen)


# ============================================================================
# RegexToolXmlParser Tests
# ============================================================================


class TestRegexToolXmlParser:
    """Tests for the RegexToolXmlParser class."""

    @pytest.fixture
    def parser(self):
        return RegexToolXmlParser(TEST_TOOL_MODEL_MAP)

    def test_parse_simple_tool(self, parser):
        """Test parsing a simple tool with one parameter."""
        xml = "<read_file>\n<path>/some/path.py</path>\n</read_file>"
        tool = parser.parse(xml, "read_file")

        assert tool is not None
        assert isinstance(tool, MockReadFileTool)
        assert tool.path == "/some/path.py"

    def test_parse_tool_with_multiple_params(self, parser):
        """Test parsing a tool with multiple parameters."""
        xml = """<read_file>
<path>/some/path.py</path>
<start_line>10</start_line>
<end_line>20</end_line>
</read_file>"""
        tool = parser.parse(xml, "read_file")

        assert tool is not None
        assert tool.path == "/some/path.py"
        assert tool.start_line == "10"
        assert tool.end_line == "20"

    def test_parse_boolean_requires_approval(self, parser):
        """Test boolean conversion for requires_approval field."""
        xml = """<execute_command>
<command>ls -la</command>
<requires_approval>true</requires_approval>
</execute_command>"""
        tool = parser.parse(xml, "execute_command")

        assert tool is not None
        assert tool.command == "ls -la"
        assert tool.requires_approval is True

        # Test false value
        xml_false = """<execute_command>
<command>ls -la</command>
<requires_approval>false</requires_approval>
</execute_command>"""
        tool_false = parser.parse(xml_false, "execute_command")
        assert tool_false.requires_approval is False

    def test_parse_boolean_recursive(self, parser):
        """Test boolean conversion for list_files recursive field."""
        xml = """<list_files>
<path>/some/dir</path>
<recursive>true</recursive>
</list_files>"""
        tool = parser.parse(xml, "list_files")

        assert tool is not None
        assert tool.path == "/some/dir"
        assert tool.recursive is True

    def test_parse_json_options(self, parser):
        """Test JSON parsing for options field in ask_followup_question."""
        xml = """<ask_followup_question>
<question>Which option?</question>
<options>["option1", "option2", "option3"]</options>
</ask_followup_question>"""
        tool = parser.parse(xml, "ask_followup_question")

        assert tool is not None
        assert tool.question == "Which option?"
        assert tool.options == ["option1", "option2", "option3"]

    def test_parse_xml_escaped_content(self, parser):
        """Test parsing XML with escaped characters."""
        xml = """<read_file>
<path>/path/with/&lt;special&gt;/chars.py</path>
</read_file>"""
        tool = parser.parse(xml, "read_file")

        assert tool is not None
        assert tool.path == "/path/with/<special>/chars.py"

    def test_parse_unknown_tool(self, parser):
        """Test parsing an unknown tool tag returns None."""
        xml = "<unknown_tool><param>value</param></unknown_tool>"
        tool = parser.parse(xml, "unknown_tool")
        assert tool is None

    def test_parse_malformed_xml(self, parser):
        """Test parsing malformed XML returns None."""
        xml = "<read_file><path>/some/path.py"  # Missing closing tags
        tool = parser.parse(xml, "read_file")
        assert tool is None

    def test_is_known_tool(self, parser):
        """Test is_known_tool method."""
        assert parser.is_known_tool("read_file") is True
        assert parser.is_known_tool("execute_command") is True
        assert parser.is_known_tool("unknown_tool") is False


# ============================================================================
# StreamingTagParser Tests
# ============================================================================


class TestStreamingTagParser:
    """Tests for the StreamingTagParser class."""

    @pytest.fixture
    def parser(self):
        xml_parser = RegexToolXmlParser(TEST_TOOL_MODEL_MAP)
        return StreamingTagParser(
            tool_xml_parser=xml_parser,
            tool_xml_reconstructor=mock_reconstruct_tool_xml,
        )

    def test_plain_text_short_retained(self, parser):
        """Test that short plain text (< 100 chars) is retained until flush."""
        events = collect_events(parser.feed("Hello, world!"))
        # Short text should not yield anything (retained in buffer)
        assert len(events) == 0

        # On flush, it should be yielded
        flush_events = collect_events(parser.flush())
        assert len(flush_events) == 1
        assert isinstance(flush_events[0], LLMOutputEvent)
        assert flush_events[0].text == "Hello, world!"

    def test_plain_text_long_partial_yield(self, parser):
        """Test that long plain text yields all but last 100 chars."""
        long_text = "x" * 150
        events = collect_events(parser.feed(long_text))

        # Should yield first 50 chars (150 - 100 = 50)
        assert len(events) == 1
        assert isinstance(events[0], LLMOutputEvent)
        assert len(events[0].text) == 50

        # Flush should yield remaining 100
        flush_events = collect_events(parser.flush())
        assert len(flush_events) == 1
        assert len(flush_events[0].text) == 100

    def test_thinking_block_single_chunk(self, parser):
        """Test parsing a complete thinking block in one chunk."""
        content = "Hello <thinking>I need to think about this</thinking> world"
        events = collect_events(parser.feed(content))

        # Find the events we care about
        output_events = [e for e in events if isinstance(e, LLMOutputEvent)]
        thinking_events = [e for e in events if isinstance(e, LLMThinkingEvent)]

        # Should have yielded "Hello " as output and thinking content
        assert any(e.text == "Hello " for e in output_events)
        assert len(thinking_events) == 1
        assert thinking_events[0].text == "I need to think about this"

        # Flush should yield remaining " world"
        flush_events = collect_events(parser.flush())
        assert len(flush_events) >= 1
        all_flush_text = "".join(
            e.text for e in flush_events if isinstance(e, LLMOutputEvent)
        )
        assert " world" in all_flush_text

    def test_thinking_block_cross_chunk(self, parser):
        """Test parsing a thinking block split across chunks."""
        # First chunk: start of thinking block - "Hello " may be yielded
        events1 = collect_events(parser.feed("Hello <thinking>I need to"))
        output_events1 = [e for e in events1 if isinstance(e, LLMOutputEvent)]
        # Hello might be yielded or retained depending on buffer size
        if output_events1:
            assert any("Hello" in e.text for e in output_events1)

        # Second chunk: end of thinking block
        events2 = collect_events(parser.feed(" finish thinking</thinking> done"))
        thinking_events = [e for e in events2 if isinstance(e, LLMThinkingEvent)]
        assert len(thinking_events) == 1
        assert thinking_events[0].text == "I need to finish thinking"

        # Flush should yield remaining
        flush_events = collect_events(parser.flush())
        all_flush_text = "".join(
            e.text for e in flush_events if isinstance(e, LLMOutputEvent)
        )
        assert "done" in all_flush_text

    def test_tool_block_single_chunk(self, parser):
        """Test parsing a complete tool block in one chunk."""
        content = "Some text <read_file>\n<path>/test.py</path>\n</read_file> more"
        events = collect_events(parser.feed(content))

        # Should yield: "Some text " as output, then ToolCallEvent
        assert len(events) == 2
        assert isinstance(events[0], LLMOutputEvent)
        assert events[0].text == "Some text "
        assert isinstance(events[1], ToolCallEvent)
        assert isinstance(events[1].tool, MockReadFileTool)
        assert events[1].tool.path == "/test.py"

    def test_tool_block_cross_chunk(self, parser):
        """Test parsing a tool block split across chunks."""
        # First chunk
        events1 = collect_events(parser.feed("Text <read_file>\n<path>/te"))

        # Second chunk
        events2 = collect_events(parser.feed("st.py</path>\n</read_file> done"))

        all_events = events1 + events2
        # Filter to meaningful events
        tool_events = [e for e in all_events if isinstance(e, ToolCallEvent)]
        assert len(tool_events) == 1
        assert tool_events[0].tool.path == "/test.py"

    def test_multiple_tool_blocks(self, parser):
        """Test parsing multiple tool blocks in sequence."""
        content = """<read_file>
<path>/file1.py</path>
</read_file>
<read_file>
<path>/file2.py</path>
</read_file>"""

        events = collect_events(parser.feed(content))
        tool_events = [e for e in events if isinstance(e, ToolCallEvent)]

        assert len(tool_events) == 2
        assert tool_events[0].tool.path == "/file1.py"
        assert tool_events[1].tool.path == "/file2.py"

    def test_unknown_tag_as_text(self, parser):
        """Test that unknown tags are treated as plain text."""
        content = "Hello <unknown_tag>some content</unknown_tag> world"
        events = collect_events(parser.feed(content))

        # Unknown tags should not trigger special handling
        # The content should eventually be yielded as text
        flush_events = collect_events(parser.flush())
        all_events = events + flush_events

        # All events should be LLMOutputEvent
        for event in all_events:
            assert isinstance(event, LLMOutputEvent)

    def test_unterminated_thinking_block(self, parser):
        """Test that unterminated thinking block produces RetryEvent."""
        # Consume the feed generator to update state
        list(parser.feed("Hello <thinking>started thinking but never finished"))
        flush_events = collect_events(parser.flush())

        # Should produce RetryEvent and remaining content as thinking
        event_types = [type(e) for e in flush_events]
        assert RetryEvent in event_types
        assert LLMThinkingEvent in event_types

        retry_event = [e for e in flush_events if isinstance(e, RetryEvent)][0]
        assert "unterminated" in retry_event.message.lower()
        assert "<thinking>" in retry_event.message

    def test_unterminated_tool_block(self, parser):
        """Test that unterminated tool block produces RetryEvent."""
        # Consume the feed generator to update state
        list(parser.feed("Hello <read_file><path>/test.py"))
        flush_events = collect_events(parser.flush())

        # Should produce RetryEvent and remaining content as text
        event_types = [type(e) for e in flush_events]
        assert RetryEvent in event_types
        assert LLMOutputEvent in event_types

        retry_event = [e for e in flush_events if isinstance(e, RetryEvent)][0]
        assert "unterminated" in retry_event.message.lower()
        assert "read_file" in retry_event.message

    def test_thinking_before_tool(self, parser):
        """Test thinking tag appearing before tool tag."""
        content = "<thinking>Let me think</thinking><read_file><path>/test.py</path></read_file>"
        events = collect_events(parser.feed(content))

        thinking_events = [e for e in events if isinstance(e, LLMThinkingEvent)]
        tool_events = [e for e in events if isinstance(e, ToolCallEvent)]

        assert len(thinking_events) == 1
        assert thinking_events[0].text == "Let me think"
        assert len(tool_events) == 1
        assert isinstance(tool_events[0].tool, MockReadFileTool)

    def test_tool_with_xml_escaped_content(self, parser):
        """Test tool with XML-escaped content in parameters."""
        content = """<read_file>
<path>/path/with/&lt;brackets&gt;.py</path>
</read_file>"""
        events = collect_events(parser.feed(content))

        tool_events = [e for e in events if isinstance(e, ToolCallEvent)]
        assert len(tool_events) == 1
        assert tool_events[0].tool.path == "/path/with/<brackets>.py"

    def test_reset_clears_state(self, parser):
        """Test that reset() clears all parser state."""
        # Partially parse something - consume the generator to update state
        list(parser.feed("Hello <thinking>partial"))
        assert parser.is_in_thinking_block is True

        # Reset
        parser.reset()

        # State should be clean
        assert parser.is_in_thinking_block is False
        assert parser.is_in_tool_block is False
        assert parser.remaining_buffer == ""


# ============================================================================
# ReasoningMetadataEmitter Tests
# ============================================================================


class TestReasoningMetadataEmitter:
    """Tests for the ReasoningMetadataEmitter class."""

    @pytest.fixture
    def emitter(self):
        return ReasoningMetadataEmitter()

    def test_reasoning_content_with_wrapper(self, emitter):
        """Test that reasoning content is wrapped with <inner_thinking> tags."""
        metadata1 = MockMetadata(reasoning_content="First reasoning")
        metadata2 = MockMetadata(reasoning_content="Second reasoning")

        events1 = collect_events(emitter.process_metadata(metadata1))
        events2 = collect_events(emitter.process_metadata(metadata2))

        # First call should emit opening tag + content
        assert len(events1) == 2
        assert events1[0].text == "<inner_thinking>"
        assert events1[1].text == "First reasoning"

        # Second call should only emit content (no opening tag)
        assert len(events2) == 1
        assert events2[0].text == "Second reasoning"

    def test_reasoning_closes_on_content_start(self, emitter):
        """Test that reasoning is closed when content starts."""
        metadata = MockMetadata(reasoning_content="Some reasoning")
        collect_events(emitter.process_metadata(metadata))

        # Simulate content starting
        close_events = collect_events(emitter.on_content_start())

        assert len(close_events) == 1
        assert close_events[0].text == "</inner_thinking>\n"

    def test_no_reasoning_content(self, emitter):
        """Test handling metadata without reasoning_content."""
        metadata = MockMetadata(reasoning_content=None)
        events = collect_events(emitter.process_metadata(metadata))
        assert len(events) == 0

    def test_reasoning_after_content_ignored(self, emitter):
        """Test that reasoning content after content start is ignored."""
        # Simulate content already started - consume the generator
        list(emitter.on_content_start())

        metadata = MockMetadata(reasoning_content="Late reasoning")
        events = collect_events(emitter.process_metadata(metadata))

        assert len(events) == 0

    def test_reset_clears_state(self, emitter):
        """Test that reset() clears emitter state."""
        metadata = MockMetadata(reasoning_content="Some reasoning")
        collect_events(emitter.process_metadata(metadata))
        # Consume the on_content_start generator
        list(emitter.on_content_start())

        assert emitter.is_reasoning_finished is True

        emitter.reset()

        assert emitter.is_reasoning_finished is False
        assert emitter.is_in_reasoning is False


# ============================================================================
# LLMResponseStreamParser Integration Tests
# ============================================================================


class TestLLMResponseStreamParser:
    """Integration tests for the complete LLMResponseStreamParser."""

    @pytest.fixture
    def parser(self):
        return create_llm_stream_parser(
            tool_model_map=TEST_TOOL_MODEL_MAP,
            tool_xml_reconstructor=mock_reconstruct_tool_xml,
        )

    def test_simple_text_stream(self, parser):
        """Test parsing a simple text stream."""
        chunks = [
            ("Hello, ", None),
            ("world!", None),
        ]
        events = collect_events(parser.parse_stream(make_generator(chunks)))

        # Should have text events and token usage at end
        text_events = [e for e in events if isinstance(e, LLMOutputEvent)]
        usage_events = [e for e in events if isinstance(e, TokenUsageEvent)]

        assert len(text_events) >= 1
        assert len(usage_events) == 1

    def test_reasoning_then_content(self, parser):
        """Test reasoning content followed by regular content."""
        chunks = [
            ("", MockMetadata(reasoning_content="I'm thinking")),
            ("", MockMetadata(reasoning_content=" about this")),
            ("Here is my response", None),
        ]
        events = collect_events(parser.parse_stream(make_generator(chunks)))

        reasoning_events = [e for e in events if isinstance(e, LLMReasoningEvent)]
        output_events = [e for e in events if isinstance(e, LLMOutputEvent)]

        # Should have: opening tag, two reasoning chunks, closing tag
        assert len(reasoning_events) == 4
        assert reasoning_events[0].text == "<inner_thinking>"
        assert reasoning_events[1].text == "I'm thinking"
        assert reasoning_events[2].text == " about this"
        assert reasoning_events[3].text == "</inner_thinking>\n"

    def test_mixed_content_with_tools(self, parser):
        """Test stream with text, thinking, and tool calls."""
        # Add enough content in one chunk to trigger processing
        chunks = [
            (
                "<thinking>Let me analyze</thinking><read_file>\n<path>/test.py</path>\n</read_file>",
                None,
            ),
        ]
        events = collect_events(parser.parse_stream(make_generator(chunks)))

        thinking_events = [e for e in events if isinstance(e, LLMThinkingEvent)]
        tool_events = [e for e in events if isinstance(e, ToolCallEvent)]

        assert len(thinking_events) == 1
        assert thinking_events[0].text == "Let me analyze"

        assert len(tool_events) == 1
        assert isinstance(tool_events[0].tool, MockReadFileTool)

    def test_token_usage_at_end(self, parser):
        """Test that TokenUsageEvent is always at the end."""
        chunks = [
            ("Hello", None),
            ("World", MagicMock(some_field="value")),
        ]
        events = collect_events(parser.parse_stream(make_generator(chunks)))

        assert isinstance(events[-1], TokenUsageEvent)
        assert events[-1].usage is not None

    def test_empty_chunks_handled(self, parser):
        """Test that empty content chunks are handled correctly."""
        chunks = [
            ("", None),
            ("", None),
            ("Hello", None),
            ("", None),
        ]
        events = collect_events(parser.parse_stream(make_generator(chunks)))

        # Should not crash and should produce output
        output_events = [e for e in events if isinstance(e, LLMOutputEvent)]
        assert len(output_events) >= 1

    def test_cross_chunk_tool_parsing(self, parser):
        """Test tool XML split across multiple chunks."""
        chunks = [
            ("<read", None),
            ("_file>\n<path>/", None),
            ("test.py</path>\n</read_file>", None),
        ]
        events = collect_events(parser.parse_stream(make_generator(chunks)))

        tool_events = [e for e in events if isinstance(e, ToolCallEvent)]
        assert len(tool_events) == 1
        assert tool_events[0].tool.path == "/test.py"

    def test_event_order_preserved(self, parser):
        """Test that event order matches the original implementation."""
        chunks = [
            ("", MockMetadata(reasoning_content="Reasoning")),
            ("Text before ", None),
            ("<thinking>Think</thinking>", None),
            ("<read_file><path>/test.py</path></read_file>", None),
            (" Text after", None),
        ]
        events = collect_events(parser.parse_stream(make_generator(chunks)))

        # Expected order: reasoning events, output, thinking, tool, output, token usage
        event_types = [type(e).__name__ for e in events]

        # Verify reasoning comes first
        first_non_reasoning = next(
            i for i, t in enumerate(event_types) if t != "LLMReasoningEvent"
        )
        for i in range(first_non_reasoning):
            assert event_types[i] == "LLMReasoningEvent"

        # Verify token usage is last
        assert event_types[-1] == "TokenUsageEvent"


# ============================================================================
# Extension Points Tests
# ============================================================================


class TestReasoningParserExtension:
    """Tests for the reasoning parser extension point."""

    @pytest.fixture
    def parser_without_reasoning_parsing(self):
        """Create parser with reasoning parsing disabled (default)."""
        return create_llm_stream_parser(
            tool_model_map=TEST_TOOL_MODEL_MAP,
            tool_xml_reconstructor=mock_reconstruct_tool_xml,
            enable_reasoning_parsing=False,
        )

    @pytest.fixture
    def parser_with_reasoning_parsing(self):
        """Create parser with reasoning parsing enabled."""
        return create_llm_stream_parser(
            tool_model_map=TEST_TOOL_MODEL_MAP,
            tool_xml_reconstructor=mock_reconstruct_tool_xml,
            enable_reasoning_parsing=True,
        )

    def test_reasoning_with_tool_xml_default_no_extraction(
        self, parser_without_reasoning_parsing
    ):
        """Test that reasoning_content with tool XML does NOT extract tools by default."""
        # Reasoning content contains a tool XML
        reasoning_with_tool = "<read_file><path>/test.py</path></read_file>"
        chunks = [
            ("", MockMetadata(reasoning_content=reasoning_with_tool)),
            ("Some content", None),
        ]
        events = collect_events(
            parser_without_reasoning_parsing.parse_stream(make_generator(chunks))
        )

        # Should have reasoning events but NO ToolCallEvent
        reasoning_events = [e for e in events if isinstance(e, LLMReasoningEvent)]
        tool_events = [e for e in events if isinstance(e, ToolCallEvent)]

        assert len(reasoning_events) >= 1
        assert len(tool_events) == 0

        # The tool XML should be in the reasoning text (as raw text)
        reasoning_text = "".join(e.text for e in reasoning_events)
        assert "<read_file>" in reasoning_text or "read_file" in reasoning_text

    def test_reasoning_with_tool_xml_extraction_enabled(
        self, parser_with_reasoning_parsing
    ):
        """Test that reasoning parsing extracts tool XML when enabled."""
        # Reasoning content contains a tool XML
        reasoning_with_tool = "<read_file><path>/test.py</path></read_file>"
        chunks = [
            ("", MockMetadata(reasoning_content=reasoning_with_tool)),
            ("Some content", None),
        ]
        events = collect_events(
            parser_with_reasoning_parsing.parse_stream(make_generator(chunks))
        )

        # Should have ToolCallEvent from reasoning
        tool_events = [e for e in events if isinstance(e, ToolCallEvent)]

        assert len(tool_events) == 1
        assert tool_events[0].tool.path == "/test.py"

    def test_reasoning_parser_property_exposed(self, parser_without_reasoning_parsing):
        """Test that reasoning_emitter is accessible for customization."""
        emitter = parser_without_reasoning_parsing.reasoning_emitter
        assert emitter is not None
        assert not emitter.has_reasoning_parser

    def test_attach_reasoning_parser_dynamically(
        self, parser_without_reasoning_parsing
    ):
        """Test dynamically attaching a reasoning parser."""
        from autocoder.common.v2.agent.llm_stream_parser import StreamingTagParser

        # Create a separate parser for reasoning
        xml_parser = parser_without_reasoning_parsing.tool_xml_parser
        reasoning_parser = StreamingTagParser(
            tool_xml_parser=xml_parser,
            tool_xml_reconstructor=mock_reconstruct_tool_xml,
        )

        # Attach it
        parser_without_reasoning_parsing.reasoning_emitter.attach_reasoning_parser(
            reasoning_parser
        )

        assert parser_without_reasoning_parsing.reasoning_emitter.has_reasoning_parser

        # Now it should extract tools from reasoning
        reasoning_with_tool = "<read_file><path>/dynamic.py</path></read_file>"
        chunks = [
            ("", MockMetadata(reasoning_content=reasoning_with_tool)),
            ("Content", None),
        ]
        events = collect_events(
            parser_without_reasoning_parsing.parse_stream(make_generator(chunks))
        )

        tool_events = [e for e in events if isinstance(e, ToolCallEvent)]
        assert len(tool_events) == 1
        assert tool_events[0].tool.path == "/dynamic.py"


class TestCustomBlockHandlers:
    """Tests for custom block handler extension point."""

    @pytest.fixture
    def parser(self):
        """Create a standard parser."""
        xml_parser = RegexToolXmlParser(TEST_TOOL_MODEL_MAP)
        return StreamingTagParser(
            tool_xml_parser=xml_parser,
            tool_xml_reconstructor=mock_reconstruct_tool_xml,
        )

    def test_register_custom_block_handler(self, parser):
        """Test registering a custom block handler."""
        from autocoder.common.v2.agent.llm_stream_parser import BlockHandler, BlockMatch

        # Define a simple custom block handler for <artifact>...</artifact>
        def find_artifact_start(buffer: str):
            pos = buffer.find("<artifact>")
            if pos != -1:
                return BlockMatch(
                    position=pos,
                    handler_id="artifact",
                    tag_length=len("<artifact>"),
                    metadata={"tag": "artifact"},
                )
            return None

        def find_artifact_end(buffer: str, metadata):
            pos = buffer.find("</artifact>")
            if pos != -1:
                return pos + len("</artifact>")
            return -1

        def create_artifact_event(content: str, metadata):
            # Just yield the content as LLMOutputEvent with a marker
            yield LLMOutputEvent(text=f"[ARTIFACT]{content}[/ARTIFACT]")

        handler = BlockHandler(
            handler_id="artifact",
            start_finder=find_artifact_start,
            end_finder=find_artifact_end,
            event_factory=create_artifact_event,
            priority=10,
        )

        parser.register_block_handler(handler)

        # Test parsing with the custom block
        content = "Before <artifact>Artifact content</artifact> After"
        events = collect_events(parser.feed(content))
        flush_events = collect_events(parser.flush())
        all_events = events + flush_events

        # Should have output events with the artifact content transformed
        output_texts = [e.text for e in all_events if isinstance(e, LLMOutputEvent)]
        full_text = "".join(output_texts)

        assert "Before" in full_text
        assert "[ARTIFACT]" in full_text
        assert "Artifact content" in full_text
        assert "[/ARTIFACT]" in full_text
        assert "After" in full_text

    def test_custom_handler_priority(self, parser):
        """Test that custom handlers are checked in priority order."""
        from autocoder.common.v2.agent.llm_stream_parser import BlockHandler, BlockMatch

        # Low priority handler
        def find_low(buffer: str):
            pos = buffer.find("<low>")
            if pos != -1:
                return BlockMatch(pos, "low", len("<low>"), None)
            return None

        # High priority handler that matches at same position
        def find_high(buffer: str):
            pos = buffer.find("<low>")  # Same position as low
            if pos != -1:
                return BlockMatch(pos, "high", len("<low>"), None)
            return None

        low_handler = BlockHandler(
            handler_id="low",
            start_finder=find_low,
            end_finder=lambda b, m: b.find("</low>") + 6 if "</low>" in b else -1,
            event_factory=lambda c, m: iter([LLMOutputEvent(text=f"LOW:{c}")]),
            priority=1,
        )

        high_handler = BlockHandler(
            handler_id="high",
            start_finder=find_high,
            end_finder=lambda b, m: b.find("</low>") + 6 if "</low>" in b else -1,
            event_factory=lambda c, m: iter([LLMOutputEvent(text=f"HIGH:{c}")]),
            priority=10,
        )

        # Register low first, then high - high should still win
        parser.register_block_handler(low_handler)
        parser.register_block_handler(high_handler)

        content = "<low>content</low>"
        events = collect_events(parser.feed(content))
        flush_events = collect_events(parser.flush())
        all_events = events + flush_events

        output_texts = [e.text for e in all_events if isinstance(e, LLMOutputEvent)]
        full_text = "".join(output_texts)

        # High priority handler should have processed it
        assert "HIGH:" in full_text

    def test_unregister_custom_handler(self, parser):
        """Test unregistering a custom block handler."""
        from autocoder.common.v2.agent.llm_stream_parser import BlockHandler, BlockMatch

        def find_test(buffer: str):
            pos = buffer.find("<test>")
            if pos != -1:
                return BlockMatch(pos, "test", len("<test>"), None)
            return None

        handler = BlockHandler(
            handler_id="test",
            start_finder=find_test,
            end_finder=lambda b, m: b.find("</test>") + 7 if "</test>" in b else -1,
            event_factory=lambda c, m: iter([LLMOutputEvent(text=f"TEST:{c}")]),
        )

        parser.register_block_handler(handler)
        removed = parser.unregister_block_handler("test")

        assert removed is not None
        assert removed.handler_id == "test"

        # Now the handler should not be active
        content = "<test>content</test>"
        events = collect_events(parser.feed(content))
        flush_events = collect_events(parser.flush())
        all_events = events + flush_events

        output_texts = [e.text for e in all_events if isinstance(e, LLMOutputEvent)]
        full_text = "".join(output_texts)

        # Should NOT have TEST: prefix since handler was unregistered
        assert "TEST:" not in full_text


class TestRoundProcessingPolicy:
    """Tests for the RoundProcessingPolicy class."""

    def test_default_policy_single_tool(self):
        """Test default policy allows only one tool per round."""
        from autocoder.common.v2.agent.agentic_edit_types import (
            RoundProcessingPolicy,
            ToolSource,
        )

        policy = RoundProcessingPolicy()

        assert policy.max_tools_per_round == 1
        assert policy.tools_executed_count == 0
        assert policy.can_execute_more_tools()
        assert not policy.is_draining

        # Execute one tool
        policy.record_tool_execution()
        assert policy.tools_executed_count == 1
        assert not policy.can_execute_more_tools()
        assert policy.should_enter_drain_mode()

    def test_policy_reset(self):
        """Test that reset clears policy state."""
        from autocoder.common.v2.agent.agentic_edit_types import RoundProcessingPolicy

        policy = RoundProcessingPolicy()
        policy.record_tool_execution()
        policy.enter_drain_mode()
        policy.set_completion_pending(True)

        assert policy.is_draining
        assert policy.has_completion_pending

        policy.reset()

        assert policy.tools_executed_count == 0
        assert not policy.is_draining
        assert not policy.has_completion_pending

    def test_policy_multi_tool(self):
        """Test policy with multiple tools allowed."""
        from autocoder.common.v2.agent.agentic_edit_types import RoundProcessingPolicy

        policy = RoundProcessingPolicy(max_tools_per_round=3)

        assert policy.can_execute_more_tools()
        policy.record_tool_execution()
        assert policy.can_execute_more_tools()
        policy.record_tool_execution()
        assert policy.can_execute_more_tools()
        policy.record_tool_execution()
        assert not policy.can_execute_more_tools()
        assert policy.should_enter_drain_mode()

    def test_policy_unlimited_tools(self):
        """Test policy with unlimited tools (max_tools_per_round <= 0)."""
        from autocoder.common.v2.agent.agentic_edit_types import RoundProcessingPolicy

        policy = RoundProcessingPolicy(max_tools_per_round=0)

        for _ in range(10):
            assert policy.can_execute_more_tools()
            policy.record_tool_execution()

        assert policy.can_execute_more_tools()  # Still can execute more

    def test_policy_tool_sources(self):
        """Test policy tool source configuration."""
        from autocoder.common.v2.agent.agentic_edit_types import (
            RoundProcessingPolicy,
            ToolSource,
        )

        # Default: only content
        policy = RoundProcessingPolicy()
        assert policy.accepts_tool_source(ToolSource.CONTENT)
        assert not policy.accepts_tool_source(ToolSource.REASONING)

        # With reasoning enabled
        policy_with_reasoning = RoundProcessingPolicy(
            tool_sources=[ToolSource.CONTENT, ToolSource.REASONING]
        )
        assert policy_with_reasoning.accepts_tool_source(ToolSource.CONTENT)
        assert policy_with_reasoning.accepts_tool_source(ToolSource.REASONING)


# ============================================================================
# Run tests
# ============================================================================

if __name__ == "__main__":
    pytest.main(["-xvs", __file__])
