"""
LLM Stream Parser Module

This module provides reusable, extensible components for parsing LLM response streams,
extracting tool calls (XML), thinking blocks, and reasoning content.

Key Components:
- RegexToolXmlParser: Parses tool XML strings into Pydantic tool models
- StreamingTagParser: State machine for parsing <thinking> and <tool_tag> blocks from content stream
- ReasoningMetadataEmitter: Handles reasoning_content from metadata with <inner_thinking> wrapper
- LLMResponseStreamParser: Orchestrates all parsers for complete LLM response parsing

Extension Points:
- To parse multiple XML formats: Register additional tag handlers in StreamingTagParser
- To parse reasoning_content for tools: Attach a StreamingTagParser instance to ReasoningMetadataEmitter
"""

import json
import re
import xml.sax.saxutils
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import (
    Any,
    Callable,
    Dict,
    Generator,
    List,
    Optional,
    Protocol,
    Set,
    Tuple,
    Type,
    Union,
    runtime_checkable,
)

from loguru import logger

from autocoder.common.v2.agent.agentic_edit_types import (
    BaseTool,
    ErrorEvent,
    LLMOutputEvent,
    LLMReasoningEvent,
    LLMThinkingEvent,
    RetryEvent,
    TokenUsageEvent,
    ToolCallEvent,
)


# Type aliases for clarity
ToolTagValidator = Callable[[str], bool]
ToolXmlReconstructor = Callable[[BaseTool], str]
ToolModelMap = Dict[str, Type[BaseTool]]

# Union of all events that can be yielded by the parser
ParsedEvent = Union[
    LLMOutputEvent,
    LLMThinkingEvent,
    LLMReasoningEvent,
    ToolCallEvent,
    ErrorEvent,
    RetryEvent,
    TokenUsageEvent,
]


@runtime_checkable
class ContentParserProtocol(Protocol):
    """
    Protocol for content parsers that can be attached to ReasoningMetadataEmitter.

    This allows any parser implementing feed() and flush() to be used for parsing
    reasoning_content, enabling tool XML extraction from reasoning phase.
    """

    def feed(self, content_chunk: str) -> Generator[ParsedEvent, None, None]:
        """Feed a content chunk and yield parsed events."""
        ...

    def flush(self) -> Generator[ParsedEvent, None, None]:
        """Flush remaining buffer and yield final events."""
        ...

    def reset(self) -> None:
        """Reset parser state for a new stream."""
        ...


class RegexToolXmlParser:
    """
    Parses tool XML strings into Pydantic tool models using regex.

    This class encapsulates the logic for:
    - Extracting parameters from XML
    - Type conversions (boolean, JSON arrays)
    - Tool-specific field handling

    Extensibility:
    - Add custom field handlers via register_field_handler()
    - Override tool model map via constructor
    """

    def __init__(self, tool_model_map: ToolModelMap):
        """
        Initialize the parser with a tool model map.

        Args:
            tool_model_map: Mapping from XML tag names to Pydantic tool model classes
        """
        self._tool_model_map = tool_model_map
        self._param_pattern = re.compile(r"<([a-zA-Z0-9_]+)>(.*?)</\1>", re.DOTALL)
        # Field handlers: (tool_tag, field_name) -> handler function
        self._field_handlers: Dict[Tuple[str, str], Callable[[str], Any]] = {}
        self._register_default_handlers()

    def _register_default_handlers(self) -> None:
        """Register default field handlers for known special cases."""
        # Boolean conversion for requires_approval (applies to all tools)
        self._field_handlers[("*", "requires_approval")] = lambda v: v.lower() == "true"
        # Boolean conversion for recursive in list_files
        self._field_handlers[("list_files", "recursive")] = (
            lambda v: v.lower() == "true"
        )
        # JSON parsing for options in ask_followup_question
        self._field_handlers[("ask_followup_question", "options")] = self._parse_json
        # JSON parsing for options in plan_mode_respond
        self._field_handlers[("plan_mode_respond", "options")] = self._parse_json

    def _parse_json(self, value: str) -> Any:
        """Parse JSON string, returning original string on failure."""
        try:
            return json.loads(value)
        except json.JSONDecodeError:
            logger.warning(f"Could not decode JSON: {value[:100]}...")
            return value

    def register_field_handler(
        self, tool_tag: str, field_name: str, handler: Callable[[str], Any]
    ) -> None:
        """
        Register a custom field handler for specific tool/field combinations.

        Args:
            tool_tag: The tool tag name, or "*" for all tools
            field_name: The field name to handle
            handler: Function that takes raw string value and returns processed value
        """
        self._field_handlers[(tool_tag, field_name)] = handler

    def get_tool_class(self, tool_tag: str) -> Optional[Type[BaseTool]]:
        """Get the tool class for a given tag name."""
        return self._tool_model_map.get(tool_tag)

    def is_known_tool(self, tag_name: str) -> bool:
        """Check if a tag name corresponds to a known tool."""
        return tag_name in self._tool_model_map

    def parse(self, tool_xml: str, tool_tag: str) -> Optional[BaseTool]:
        """
        Parse a tool XML string into a Pydantic tool model.

        Args:
            tool_xml: The complete XML string including opening and closing tags
            tool_tag: The tool tag name (e.g., "read_file", "execute_command")

        Returns:
            Parsed tool model instance, or None if parsing failed
        """
        params: Dict[str, Any] = {}
        try:
            # Find content between <tool_tag> and </tool_tag>
            inner_xml_match = re.search(
                rf"<{tool_tag}>(.*?)</{tool_tag}>", tool_xml, re.DOTALL
            )
            if not inner_xml_match:
                logger.error(
                    f"Could not find content within <{tool_tag}>...</{tool_tag}>"
                )
                return None
            inner_xml = inner_xml_match.group(1).strip()

            # Find <param>value</param> pairs within the inner content
            for m in self._param_pattern.finditer(inner_xml):
                key = m.group(1)
                # Basic unescaping
                val = xml.sax.saxutils.unescape(m.group(2))
                params[key] = val

            # Apply field handlers
            for key in list(params.keys()):
                # Check tool-specific handler first, then wildcard
                handler = self._field_handlers.get(
                    (tool_tag, key)
                ) or self._field_handlers.get(("*", key))
                if handler:
                    params[key] = handler(params[key])

            # Get tool class and instantiate
            tool_cls = self._tool_model_map.get(tool_tag)
            if tool_cls:
                return tool_cls(**params)
            else:
                logger.error(f"Tool class not found for tag: {tool_tag}")
                return None

        except Exception as e:
            logger.exception(
                f"Failed to parse tool XML for <{tool_tag}>: {e}\nXML:\n{tool_xml}"
            )
            return None


@dataclass
class StreamingTagParserState:
    """Internal state for the streaming tag parser."""

    buffer: str = ""
    in_tool_block: bool = False
    in_thinking_block: bool = False
    current_tool_tag: Optional[str] = None
    # For custom block handlers
    in_custom_block: bool = False
    current_custom_handler_id: Optional[str] = None


@dataclass
class BlockMatch:
    """
    Represents a matched block start in the buffer.

    Used by block handlers to report where a block starts and how to handle it.
    """

    position: int  # Start position in buffer
    handler_id: str  # Identifier for the handler that matched
    tag_length: int  # Length of the opening tag (to skip when entering block)
    metadata: Optional[Dict[str, Any]] = None  # Handler-specific data (e.g., tag name)


@dataclass
class BlockHandler:
    """
    Configuration for a custom block handler.

    Block handlers allow extending StreamingTagParser to recognize and process
    additional XML-like block structures beyond the built-in <thinking> and tool blocks.

    Extensibility:
    - Define start_finder to locate the opening of your block type
    - Define end_finder to locate the closing tag
    - Define event_factory to create the appropriate event from block content
    """

    handler_id: str  # Unique identifier for this handler
    start_finder: Callable[
        [str], Optional[BlockMatch]
    ]  # Find block start, return BlockMatch or None
    end_finder: Callable[
        [str, Optional[Dict[str, Any]]], int
    ]  # Find end position given buffer and metadata, return -1 if not found
    event_factory: Callable[
        [str, Optional[Dict[str, Any]]], Generator[ParsedEvent, None, None]
    ]  # Create event(s) from block content and metadata
    priority: int = 0  # Higher priority handlers are checked first (default: 0)


class StreamingTagParser:
    """
    State machine for parsing <thinking> and <tool_tag> blocks from a content stream.

    This parser maintains a buffer and tracks whether we're inside a thinking block,
    tool block, or plain text. It yields appropriate events as blocks are completed.

    Extensibility:
    - Inject custom tool tag validator via constructor
    - Inject custom tool XML parser and reconstructor via constructor
    - The parser can be instantiated multiple times for different sources
    - Register additional block handlers via extra_block_handlers parameter
      to parse custom XML structures (e.g., <artifact>, <code_review>, etc.)

    Block Handler Extension:
    To add custom block parsing, create a BlockHandler with:
    - start_finder: Function to find block start, returns BlockMatch or None
    - end_finder: Function to find block end given buffer and metadata
    - event_factory: Generator function to create events from block content
    """

    # Default buffer retention size for potential tag starts
    DEFAULT_BUFFER_RETAIN_SIZE = 100

    def __init__(
        self,
        tool_xml_parser: RegexToolXmlParser,
        tool_xml_reconstructor: ToolXmlReconstructor,
        tool_tag_validator: Optional[ToolTagValidator] = None,
        thinking_start_tag: str = "<thinking>",
        thinking_end_tag: str = "</thinking>",
        buffer_retain_size: int = DEFAULT_BUFFER_RETAIN_SIZE,
        extra_block_handlers: Optional[List[BlockHandler]] = None,
    ):
        """
        Initialize the streaming tag parser.

        Args:
            tool_xml_parser: Parser for converting tool XML to tool models
            tool_xml_reconstructor: Function to reconstruct XML from tool models
            tool_tag_validator: Optional function to validate if a tag is a known tool
                               (defaults to tool_xml_parser.is_known_tool)
            thinking_start_tag: Opening tag for thinking blocks
            thinking_end_tag: Closing tag for thinking blocks
            buffer_retain_size: Number of characters to retain in buffer for potential tags
            extra_block_handlers: Optional list of additional BlockHandler instances
                                 for parsing custom block structures. Handlers with
                                 higher priority are checked first. Default is None.
        """
        self._tool_xml_parser = tool_xml_parser
        self._tool_xml_reconstructor = tool_xml_reconstructor
        self._tool_tag_validator = tool_tag_validator or tool_xml_parser.is_known_tool
        self._thinking_start_tag = thinking_start_tag
        self._thinking_end_tag = thinking_end_tag
        self._buffer_retain_size = buffer_retain_size
        self._tool_start_pattern = re.compile(r"<([a-zA-Z0-9_]+)>")

        # Extra block handlers sorted by priority (descending)
        self._extra_block_handlers: List[BlockHandler] = sorted(
            extra_block_handlers or [], key=lambda h: h.priority, reverse=True
        )
        # Map handler_id -> handler for quick lookup
        self._handler_map: Dict[str, BlockHandler] = {
            h.handler_id: h for h in self._extra_block_handlers
        }
        # Current custom block metadata (if in a custom block)
        self._current_custom_metadata: Optional[Dict[str, Any]] = None

        # Parser state
        self._state = StreamingTagParserState()

    def reset(self) -> None:
        """Reset the parser state for a new stream."""
        self._state = StreamingTagParserState()
        self._current_custom_metadata = None

    def register_block_handler(self, handler: BlockHandler) -> None:
        """
        Register an additional block handler dynamically.

        The handler will be added to the list and the list will be re-sorted by priority.

        Args:
            handler: The BlockHandler to register
        """
        self._extra_block_handlers.append(handler)
        self._extra_block_handlers.sort(key=lambda h: h.priority, reverse=True)
        self._handler_map[handler.handler_id] = handler

    def unregister_block_handler(self, handler_id: str) -> Optional[BlockHandler]:
        """
        Unregister a block handler by its ID.

        Args:
            handler_id: The ID of the handler to remove

        Returns:
            The removed handler, or None if not found
        """
        handler = self._handler_map.pop(handler_id, None)
        if handler:
            self._extra_block_handlers = [
                h for h in self._extra_block_handlers if h.handler_id != handler_id
            ]
        return handler

    @property
    def is_in_thinking_block(self) -> bool:
        """Check if currently inside a thinking block."""
        return self._state.in_thinking_block

    @property
    def is_in_tool_block(self) -> bool:
        """Check if currently inside a tool block."""
        return self._state.in_tool_block

    @property
    def is_in_custom_block(self) -> bool:
        """Check if currently inside a custom block (from extra handlers)."""
        return self._state.in_custom_block

    @property
    def current_tool_tag(self) -> Optional[str]:
        """Get the current tool tag name if inside a tool block."""
        return self._state.current_tool_tag

    @property
    def current_custom_handler_id(self) -> Optional[str]:
        """Get the current custom handler ID if inside a custom block."""
        return self._state.current_custom_handler_id

    @property
    def remaining_buffer(self) -> str:
        """Get any remaining content in the buffer."""
        return self._state.buffer

    def _find_custom_block_start(self) -> Optional[BlockMatch]:
        """
        Search for custom block starts in the buffer.

        Checks all registered extra block handlers in priority order.

        Returns:
            The BlockMatch with the earliest position, or None if no match.
        """
        best_match: Optional[BlockMatch] = None

        for handler in self._extra_block_handlers:
            match = handler.start_finder(self._state.buffer)
            if match is not None:
                if best_match is None or match.position < best_match.position:
                    best_match = match

        return best_match

    def feed(self, content_chunk: str) -> Generator[ParsedEvent, None, None]:
        """
        Feed a content chunk into the parser and yield any completed events.

        Args:
            content_chunk: A chunk of content from the LLM stream

        Yields:
            Parsed events (LLMOutputEvent, LLMThinkingEvent, ToolCallEvent, ErrorEvent,
            or custom events from registered block handlers)
        """
        self._state.buffer += content_chunk

        while True:
            found_event = False

            # 1. Check for </thinking> if inside thinking block
            if self._state.in_thinking_block:
                end_think_pos = self._state.buffer.find(self._thinking_end_tag)
                if end_think_pos != -1:
                    thinking_content = self._state.buffer[:end_think_pos]
                    yield LLMThinkingEvent(text=thinking_content)
                    self._state.buffer = self._state.buffer[
                        end_think_pos + len(self._thinking_end_tag) :
                    ]
                    self._state.in_thinking_block = False
                    found_event = True
                    continue
                else:
                    # Need more data to close thinking block
                    break

            # 2. Check for </tool_tag> if inside tool block
            elif self._state.in_tool_block:
                end_tag = f"</{self._state.current_tool_tag}>"
                end_tool_pos = self._state.buffer.find(end_tag)
                if end_tool_pos != -1:
                    tool_block_end_index = end_tool_pos + len(end_tag)
                    tool_xml = self._state.buffer[:tool_block_end_index]

                    # Parse the tool XML
                    tool_obj = self._tool_xml_parser.parse(
                        tool_xml, self._state.current_tool_tag
                    )

                    if tool_obj:
                        # Reconstruct XML accurately after successful parsing
                        reconstructed_xml = self._tool_xml_reconstructor(tool_obj)
                        if reconstructed_xml.startswith("<error>"):
                            yield ErrorEvent(
                                message=f"Failed to reconstruct XML for tool {self._state.current_tool_tag}"
                            )
                        else:
                            yield ToolCallEvent(
                                tool=tool_obj, tool_xml=reconstructed_xml
                            )
                    else:
                        # Failed to parse - yield as plain text
                        yield LLMOutputEvent(
                            text=f"Failed to parse tool: <{self._state.current_tool_tag}> {tool_xml}"
                        )

                    self._state.buffer = self._state.buffer[tool_block_end_index:]
                    self._state.in_tool_block = False
                    self._state.current_tool_tag = None
                    found_event = True
                    continue
                else:
                    # Need more data to close tool block
                    break

            # 2.5. Check for custom block end if inside a custom block
            elif self._state.in_custom_block:
                handler_id = self._state.current_custom_handler_id
                handler = self._handler_map.get(handler_id) if handler_id else None
                if handler:
                    end_pos = handler.end_finder(
                        self._state.buffer, self._current_custom_metadata
                    )
                    if end_pos != -1:
                        # Extract block content and process through handler
                        block_content = self._state.buffer[:end_pos]
                        yield from handler.event_factory(
                            block_content, self._current_custom_metadata
                        )

                        self._state.buffer = self._state.buffer[end_pos:]
                        self._state.in_custom_block = False
                        self._state.current_custom_handler_id = None
                        self._current_custom_metadata = None
                        found_event = True
                        continue
                    else:
                        # Need more data to close custom block
                        break
                else:
                    # Handler not found - exit custom block state (shouldn't happen)
                    logger.warning(f"Custom block handler not found: {handler_id}")
                    self._state.in_custom_block = False
                    self._state.current_custom_handler_id = None
                    self._current_custom_metadata = None
                    break

            # 3. Check for <thinking>, <tool_tag>, or custom block if in plain text state
            else:
                start_think_pos = self._state.buffer.find(self._thinking_start_tag)
                tool_match = self._tool_start_pattern.search(self._state.buffer)
                start_tool_pos = tool_match.start() if tool_match else -1
                tool_name = tool_match.group(1) if tool_match else None

                # Also check for custom block starts
                custom_match = self._find_custom_block_start()

                # Determine which tag comes first (if any)
                first_tag_pos = -1
                is_thinking = False
                is_tool = False
                is_custom = False
                custom_block_match: Optional[BlockMatch] = None

                # Compare thinking, tool, and custom positions
                candidates: List[Tuple[int, str]] = []
                if start_think_pos != -1:
                    candidates.append((start_think_pos, "thinking"))
                if (
                    start_tool_pos != -1
                    and tool_name
                    and self._tool_tag_validator(tool_name)
                ):
                    candidates.append((start_tool_pos, "tool"))
                if custom_match is not None:
                    candidates.append((custom_match.position, "custom"))

                if candidates:
                    # Sort by position, thinking gets priority on ties
                    candidates.sort(
                        key=lambda x: (x[0], 0 if x[1] == "thinking" else 1)
                    )
                    first_pos, first_type = candidates[0]
                    first_tag_pos = first_pos
                    is_thinking = first_type == "thinking"
                    is_tool = first_type == "tool"
                    is_custom = first_type == "custom"
                    if is_custom:
                        custom_block_match = custom_match

                if first_tag_pos != -1:
                    # Found either <thinking>, a known <tool>, or custom block
                    # Yield preceding text if any
                    preceding_text = self._state.buffer[:first_tag_pos]
                    if preceding_text:
                        yield LLMOutputEvent(text=preceding_text)

                    # Transition state
                    if is_thinking:
                        self._state.buffer = self._state.buffer[
                            first_tag_pos + len(self._thinking_start_tag) :
                        ]
                        self._state.in_thinking_block = True
                    elif is_tool:
                        # Keep the starting tag in buffer
                        self._state.buffer = self._state.buffer[first_tag_pos:]
                        self._state.in_tool_block = True
                        self._state.current_tool_tag = tool_name
                    elif is_custom and custom_block_match is not None:
                        # Enter custom block state
                        self._state.buffer = self._state.buffer[
                            first_tag_pos + custom_block_match.tag_length :
                        ]
                        self._state.in_custom_block = True
                        self._state.current_custom_handler_id = (
                            custom_block_match.handler_id
                        )
                        self._current_custom_metadata = custom_block_match.metadata

                    found_event = True
                    continue

                else:
                    # No tags found, or only unknown tags found
                    # Yield text chunk but keep some buffer for potential tag start
                    split_point = max(
                        0, len(self._state.buffer) - self._buffer_retain_size
                    )
                    text_to_yield = self._state.buffer[:split_point]
                    if text_to_yield:
                        yield LLMOutputEvent(text=text_to_yield)
                        self._state.buffer = self._state.buffer[split_point:]
                    break

            # If no event was processed in this iteration, break inner loop
            if not found_event:
                break

    def flush(self) -> Generator[ParsedEvent, None, None]:
        """
        Flush any remaining content when the stream ends.

        Yields:
            Final events for any remaining buffer content, plus RetryEvent if unterminated blocks
        """
        if self._state.in_thinking_block:
            # Unterminated thinking block
            yield RetryEvent(message="Stream ended with unterminated <thinking> block.")
            if self._state.buffer:
                yield LLMThinkingEvent(text=self._state.buffer)
        elif self._state.in_tool_block:
            # Unterminated tool block
            yield RetryEvent(
                message=f"Stream ended with unterminated <{self._state.current_tool_tag}> block."
            )
            if self._state.buffer:
                yield LLMOutputEvent(text=self._state.buffer)
        elif self._state.in_custom_block:
            # Unterminated custom block
            handler_id = self._state.current_custom_handler_id or "unknown"
            yield RetryEvent(
                message=f"Stream ended with unterminated custom block (handler: {handler_id})."
            )
            if self._state.buffer:
                yield LLMOutputEvent(text=self._state.buffer)
        elif self._state.buffer:
            # Yield remaining plain text
            yield LLMOutputEvent(text=self._state.buffer)

        # Clear the buffer after flush
        self._state.buffer = ""


@dataclass
class ReasoningEmitterState:
    """Internal state for the reasoning metadata emitter."""

    reasoning_finished: bool = False
    in_reasoning: bool = False


class ReasoningMetadataEmitter:
    """
    Handles reasoning_content from LLM metadata, emitting it with <inner_thinking> wrapper.

    The reasoning content appears in metadata before the main content starts.
    This emitter wraps reasoning content with <inner_thinking> tags and closes
    them when the first content chunk appears.

    Extensibility:
    - To parse reasoning_content for tool XML, attach a ContentParserProtocol instance
      via attach_reasoning_parser(). When attached:
      - Text events from the parser are mapped to LLMReasoningEvent
      - ToolCallEvent and ErrorEvent pass through unchanged
      - This enables tool XML extraction from reasoning_content
    - Default behavior (no parser attached): emit raw text as LLMReasoningEvent
    """

    def __init__(
        self,
        opening_tag: str = "<inner_thinking>",
        closing_tag: str = "</inner_thinking>\n",
        reasoning_parser: Optional["ContentParserProtocol"] = None,
    ):
        """
        Initialize the reasoning emitter.

        Args:
            opening_tag: Tag to emit at start of reasoning content
            closing_tag: Tag to emit when reasoning ends (content starts)
            reasoning_parser: Optional parser for extracting tool XML from reasoning.
                             When provided, reasoning_content is parsed instead of
                             being emitted as raw text. Default is None (raw text mode).
        """
        self._opening_tag = opening_tag
        self._closing_tag = closing_tag
        self._reasoning_parser = reasoning_parser
        self._state = ReasoningEmitterState()

    def reset(self) -> None:
        """Reset the emitter state for a new stream."""
        self._state = ReasoningEmitterState()
        if self._reasoning_parser is not None:
            self._reasoning_parser.reset()

    @property
    def is_reasoning_finished(self) -> bool:
        """Check if reasoning phase has ended (content has started)."""
        return self._state.reasoning_finished

    @property
    def is_in_reasoning(self) -> bool:
        """Check if currently emitting reasoning content."""
        return self._state.in_reasoning

    @property
    def has_reasoning_parser(self) -> bool:
        """Check if a reasoning parser is attached."""
        return self._reasoning_parser is not None

    def attach_reasoning_parser(self, parser: "ContentParserProtocol") -> None:
        """
        Attach a parser for extracting tool XML from reasoning_content.

        When attached, reasoning_content will be fed through the parser.
        Text events are mapped to LLMReasoningEvent, while ToolCallEvent
        and ErrorEvent pass through unchanged.

        Args:
            parser: A parser implementing ContentParserProtocol
        """
        self._reasoning_parser = parser

    def detach_reasoning_parser(self) -> Optional["ContentParserProtocol"]:
        """
        Detach and return the current reasoning parser.

        Returns:
            The previously attached parser, or None if none was attached.
        """
        parser = self._reasoning_parser
        self._reasoning_parser = None
        return parser

    def _map_parser_event_to_reasoning(
        self, event: ParsedEvent
    ) -> Generator[ParsedEvent, None, None]:
        """
        Map events from the reasoning parser to appropriate output events.

        Text events (LLMOutputEvent, LLMThinkingEvent) are converted to LLMReasoningEvent.
        ToolCallEvent and ErrorEvent pass through unchanged.
        Other events are logged and skipped.

        Args:
            event: An event from the reasoning parser

        Yields:
            Mapped events appropriate for the reasoning phase
        """
        if isinstance(event, (LLMOutputEvent, LLMThinkingEvent)):
            # Map text output from parser to reasoning event
            yield LLMReasoningEvent(text=event.text)
        elif isinstance(event, (ToolCallEvent, ErrorEvent)):
            # Tool calls and errors pass through unchanged
            yield event
        elif isinstance(event, RetryEvent):
            # Log retry events but don't yield (parsing will complete in flush)
            logger.debug(f"Retry event in reasoning parser: {event.message}")
        else:
            # Skip other event types (TokenUsageEvent shouldn't appear here)
            logger.debug(f"Skipping unexpected event type in reasoning: {type(event)}")

    def process_metadata(self, metadata: Any) -> Generator[ParsedEvent, None, None]:
        """
        Process metadata for reasoning_content, yielding events if present.

        When a reasoning_parser is attached, the reasoning_content is fed through
        the parser and events are mapped appropriately. Otherwise, raw text is
        emitted as LLMReasoningEvent.

        Args:
            metadata: The metadata object from the LLM stream

        Yields:
            LLMReasoningEvent for reasoning text, or ToolCallEvent/ErrorEvent
            if a parser is attached and finds tool XML in reasoning_content
        """
        if self._state.reasoning_finished:
            return

        if metadata is not None:
            if hasattr(metadata, "reasoning_content") and metadata.reasoning_content:
                reasoning_content_text = metadata.reasoning_content

                # Add opening tag at the start of reasoning
                if not self._state.in_reasoning:
                    yield LLMReasoningEvent(text=self._opening_tag)
                    self._state.in_reasoning = True

                # Process through parser if attached, otherwise emit raw text
                if self._reasoning_parser is not None:
                    for event in self._reasoning_parser.feed(reasoning_content_text):
                        yield from self._map_parser_event_to_reasoning(event)
                else:
                    yield LLMReasoningEvent(text=reasoning_content_text)

    def on_content_start(self) -> Generator[ParsedEvent, None, None]:
        """
        Called when content starts, to close reasoning phase if active.

        If a reasoning parser is attached, it will be flushed to emit any
        remaining buffered content.

        Yields:
            LLMReasoningEvent with closing tag if reasoning was active,
            plus any events from flushing the parser
        """
        # Flush the reasoning parser if attached
        if self._reasoning_parser is not None and self._state.in_reasoning:
            for event in self._reasoning_parser.flush():
                yield from self._map_parser_event_to_reasoning(event)

        if self._state.in_reasoning and not self._state.reasoning_finished:
            yield LLMReasoningEvent(text=self._closing_tag)
            self._state.in_reasoning = False
        self._state.reasoning_finished = True


class LLMResponseStreamParser:
    """
    Orchestrates parsing of complete LLM response streams.

    This class combines:
    - ReasoningMetadataEmitter for reasoning_content handling
    - StreamingTagParser for content parsing (thinking/tool blocks)
    - TokenUsageEvent emission at stream end

    Extensibility:
    - Replace or extend individual parsers via constructor injection
    - Add additional parsers for new content sources
    - Enable reasoning parsing via enable_reasoning_parsing flag to extract
      tool XML from reasoning_content (disabled by default)
    """

    def __init__(
        self,
        tool_model_map: ToolModelMap,
        tool_xml_reconstructor: ToolXmlReconstructor,
        cancel_token_checker: Optional[Callable[[], None]] = None,
        enable_reasoning_parsing: bool = False,
    ):
        """
        Initialize the LLM response stream parser.

        Args:
            tool_model_map: Mapping from XML tag names to tool model classes
            tool_xml_reconstructor: Function to reconstruct XML from tool models
            cancel_token_checker: Optional function to check for cancellation
            enable_reasoning_parsing: If True, creates a separate StreamingTagParser
                                     for parsing tool XML from reasoning_content.
                                     Default is False (reasoning emitted as raw text).
        """
        self._tool_xml_parser = RegexToolXmlParser(tool_model_map)
        self._tag_parser = StreamingTagParser(
            tool_xml_parser=self._tool_xml_parser,
            tool_xml_reconstructor=tool_xml_reconstructor,
        )

        # Optionally create a reasoning parser for tool XML extraction
        reasoning_parser: Optional[StreamingTagParser] = None
        if enable_reasoning_parsing:
            reasoning_parser = StreamingTagParser(
                tool_xml_parser=self._tool_xml_parser,
                tool_xml_reconstructor=tool_xml_reconstructor,
            )

        self._reasoning_emitter = ReasoningMetadataEmitter(
            reasoning_parser=reasoning_parser
        )
        self._cancel_token_checker = cancel_token_checker
        self._enable_reasoning_parsing = enable_reasoning_parsing

    def reset(self) -> None:
        """Reset all parser states for a new stream."""
        self._tag_parser.reset()
        self._reasoning_emitter.reset()

    @property
    def tool_xml_parser(self) -> RegexToolXmlParser:
        """Access the tool XML parser for customization."""
        return self._tool_xml_parser

    @property
    def tag_parser(self) -> StreamingTagParser:
        """Access the streaming tag parser for customization."""
        return self._tag_parser

    @property
    def reasoning_emitter(self) -> ReasoningMetadataEmitter:
        """Access the reasoning emitter for customization."""
        return self._reasoning_emitter

    @property
    def is_reasoning_parsing_enabled(self) -> bool:
        """Check if reasoning parsing is enabled."""
        return self._enable_reasoning_parsing

    def parse_stream(
        self, generator: Generator[Tuple[str, Any], None, None]
    ) -> Generator[ParsedEvent, None, None]:
        """
        Parse an LLM response stream, yielding events for all content types.

        Args:
            generator: Iterator yielding (content_chunk, metadata) tuples

        Yields:
            All parsed events including reasoning, output, thinking, tool calls,
            errors, retries, and token usage
        """
        last_metadata = None

        try:
            for content_chunk, metadata in generator:
                # Check for cancellation if checker provided
                if self._cancel_token_checker:
                    self._cancel_token_checker()

                # Handle reasoning_content from metadata (before content starts)
                yield from self._reasoning_emitter.process_metadata(metadata)

                if not content_chunk:
                    last_metadata = metadata
                    continue

                # Once content appears, close reasoning phase
                yield from self._reasoning_emitter.on_content_start()

                last_metadata = metadata

                # Feed content to tag parser
                yield from self._tag_parser.feed(content_chunk)

        except Exception as e:
            # Re-raise cancellation exceptions
            from autocoder.common.global_cancel import CancelRequestedException

            if isinstance(e, CancelRequestedException):
                raise e

            # Format error details
            import traceback

            full_traceback = traceback.format_exc()
            error_detail = (
                full_traceback[-300:] if len(full_traceback) > 300 else full_traceback
            )

            logger.error(f"LLM connection failed: {error_detail}")
            yield ErrorEvent(message=f"Connection failed. Error: {error_detail}")

        # Flush remaining content
        yield from self._tag_parser.flush()

        # Emit token usage event last
        yield TokenUsageEvent(usage=last_metadata)


# Factory function for creating a parser with standard configuration
def create_llm_stream_parser(
    tool_model_map: ToolModelMap,
    tool_xml_reconstructor: ToolXmlReconstructor,
    cancel_token_checker: Optional[Callable[[], None]] = None,
    enable_reasoning_parsing: bool = False,
) -> LLMResponseStreamParser:
    """
    Create a configured LLM response stream parser.

    Args:
        tool_model_map: Mapping from XML tag names to tool model classes
        tool_xml_reconstructor: Function to reconstruct XML from tool models
        cancel_token_checker: Optional function to check for cancellation
        enable_reasoning_parsing: If True, enables parsing of tool XML from
                                 reasoning_content. Default is False.

    Returns:
        Configured LLMResponseStreamParser instance
    """
    return LLMResponseStreamParser(
        tool_model_map=tool_model_map,
        tool_xml_reconstructor=tool_xml_reconstructor,
        cancel_token_checker=cancel_token_checker,
        enable_reasoning_parsing=enable_reasoning_parsing,
    )
