"""
LLM Output Processor

Intercepts LLM output to:
1. Extract tool calls and send them via text_output_queue
2. Forward clean text to TTS pipeline
3. Sanitize ApiPilot internal markers before sending output to clients or TTS
"""

from __future__ import annotations

import logging
import re
from collections.abc import Iterator
from queue import Queue

from speech_to_speech.baseHandler import BaseHandler
from speech_to_speech.pipeline.events import AssistantTextEvent, ResponseFailedEvent, TokenUsageEvent

# Sanitize ApiPilot CALL/RESULT markers before emitting AssistantTextEvent
# to WebSocket clients or TTSInput to Qwen3-TTS, which would speak the markers.
# Observed formats from upstream Vocero start_pipeline_rag:
#   1) <<<APIPILOT_CALL:NAME=tool,ARGUMENTS=json>>>   standard triple brackets
#   2) <APIPILOT_CALL:NAME=get_user_briefing,ARGUMENTS={"format":"voice"}
#                                 single brackets and model-generated arguments
#   3) <<<APIPILOT_CALL:...}>  \n>>                      split across lines
#   4) >>>APIPILOT_CALL:...<<<                            inverted brackets
#   5) <<<APIPILOT_RESULT:PASS2>>> followed by a result block
_APIPILOT_CALL_NORM = re.compile(
    r"[<]{1,3}\s*APIPILOT_CALL:NAME=([A-Za-z0-9_.-]+),ARGUMENTS=(.*?)[>]{1,3}",
    re.IGNORECASE | re.DOTALL,
)
_APIPILOT_CALL_INV = re.compile(
    r"[>]{1,3}\s*APIPILOT_CALL:NAME=([A-Za-z0-9_.-]+),ARGUMENTS=(.*?)[<]{1,3}",
    re.IGNORECASE | re.DOTALL,
)
_APIPILOT_RESULT = re.compile(
    r"[<]{1,3}\s*APIPILOT_RESULT(?::PASS2)?[^<>\n]{0,4000}[>]{1,3}",
    re.IGNORECASE | re.DOTALL,
)
_APIPILOT_SPANISH_PLACEHOLDER = ("Voy a consultar tus datos, un momento por favor.")


def _sanitize_apipilot_before_client_and_tts(text: str) -> tuple[str, bool, str]:
    if not isinstance(text, str) or not text.strip():
        return text, False, ""
    found_call = False
    found_call_tool = ""
    def _repl_call(m: re.Match) -> str:
        nonlocal found_call, found_call_tool
        found_call = True
        try:
            nm = m.group(1)
            if isinstance(nm, str) and nm.strip():
                found_call_tool = nm.strip()
        except Exception:
            pass
        return f" {_APIPILOT_SPANISH_PLACEHOLDER} "
    new_text = _APIPILOT_CALL_NORM.sub(_repl_call, text)
    if not found_call:
        new_text = _APIPILOT_CALL_INV.sub(_repl_call, new_text)
    found_result = False
    if _APIPILOT_RESULT.search(new_text):
        new_text = _APIPILOT_RESULT.sub(" ", new_text)
        found_result = True
    # Catch marker fragments without a wrapper, including mid-stream truncation.
    low = new_text.lower()
    if "apipilot" in low and ("call" in low or "result" in low or "name=" in low or "arguments=" in low):
        new_text = re.sub(r"(?i)\bAPIPILOT(?:_CALL|_RESULT)?\b", " ", new_text)
        new_text = re.sub(r"(?i)\b(?:NAME|ARGUMENTS|STATUS|PASS2)\s*[:=]", " ", new_text)
        found_call = True
    if found_call or found_result:
        new_text = re.sub(r"\s{2,}", " ", new_text).strip()
    tag = ""
    if found_call:
        tag = f"CALL({found_call_tool or '?'})"
    if found_result:
        tag = (tag + "+" if tag else "") + "RESULT"
    return new_text, (found_call or found_result), tag
from speech_to_speech.pipeline.handler_types import LLMOut, TTSIn
from speech_to_speech.pipeline.messages import EndOfResponse, LLMResponseChunk, TokenUsage, TTSInput
from speech_to_speech.pipeline.queue_types import TextEventItem
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
from speech_to_speech.utils.utils import response_wants_audio

logger = logging.getLogger(__name__)


class LMOutputProcessor(BaseHandler[LLMOut, TTSIn]):
    """
    Processes LLM output to extract tool calls and forward clean text to TTS.

    Input: :class:`LLMResponseChunk`, :class:`TokenUsage`, or :class:`EndOfResponse` from LLM
    Output: :class:`TTSInput` or :class:`EndOfResponse` to TTS
    Side effect: Sends :class:`AssistantTextEvent` / :class:`TokenUsageEvent` to text_output_queue
    """

    def setup(
        self,
        text_output_queue: Queue[TextEventItem] | None = None,
        speculative_turns: SpeculativeTurnTracker | None = None,
    ) -> None:
        """
        Initialize the processor.

        Args:
            text_output_queue: Queue to send text messages and tool calls
        """
        self.text_output_queue = text_output_queue
        self.speculative_turns = speculative_turns

    def _turn_output_allowed(self, turn_id: str | None, turn_revision: int | None) -> bool:
        if self.speculative_turns is None:
            return True
        return self.speculative_turns.is_latest_after_reopen_grace(turn_id, turn_revision)

    def process(self, lm_output: LLMOut) -> Iterator[TTSIn]:
        """
        Process LLM output: send text/tools to WebSocket, forward clean text to TTS.

        Yields:
            :class:`TTSInput` or :class:`EndOfResponse` for TTS
        """
        if isinstance(lm_output, TokenUsage):
            if not self._turn_output_allowed(
                lm_output.turn_id,
                lm_output.turn_revision,
            ):
                logger.debug(
                    "Dropping stale token usage for turn=%s rev=%s", lm_output.turn_id, lm_output.turn_revision
                )
                return
            if self.text_output_queue is not None:
                self.text_output_queue.put(
                    TokenUsageEvent(
                        input_tokens=lm_output.input_tokens or 0,
                        output_tokens=lm_output.output_tokens or 0,
                        turn_id=lm_output.turn_id,
                        turn_revision=lm_output.turn_revision,
                    )
                )
            return

        if isinstance(lm_output, EndOfResponse):
            if not self._turn_output_allowed(
                lm_output.turn_id,
                lm_output.turn_revision,
            ):
                logger.debug(
                    "Dropping stale end-of-response for turn=%s rev=%s",
                    lm_output.turn_id,
                    lm_output.turn_revision,
                )
                return
            # A failed generation (e.g. invalid out-of-band input) closes the response as
            # "failed" via the text side-channel, then falls through to emit the normal
            # EndOfResponse so the audio path still re-enables listening / releases the slot.
            if lm_output.error and self.text_output_queue is not None:
                self.text_output_queue.put(
                    ResponseFailedEvent(
                        message=lm_output.error,
                        turn_id=lm_output.turn_id,
                        turn_revision=lm_output.turn_revision,
                    )
                )
            yield EndOfResponse(
                turn_id=lm_output.turn_id,
                turn_revision=lm_output.turn_revision,
                cancel_generation=lm_output.cancel_generation,
            )
            return

        if not isinstance(lm_output, LLMResponseChunk):
            logger.warning("LMOutputProcessor received unexpected type: %s", type(lm_output))
            return

        if not self._turn_output_allowed(
            lm_output.turn_id,
            lm_output.turn_revision,
        ):
            logger.debug("Dropping stale LLM chunk for turn=%s rev=%s", lm_output.turn_id, lm_output.turn_revision)
            return

        logger.debug(f"LM processor: text='{lm_output.text}', tools={lm_output.tools}")

        # Sanitize centrally before AssistantTextEvent and TTSInput so neither
        # clients nor Qwen3-TTS receive internal ApiPilot markers.
        clean_text = lm_output.text
        any_sanitized = False
        sanitize_tag = ""
        if clean_text and clean_text.strip():
            clean_text, any_sanitized, sanitize_tag = _sanitize_apipilot_before_client_and_tts(clean_text)
            if any_sanitized:
                logger.warning(
                    "[LM-OUTPUT-APIPILOT-SANITIZE] %s — tool=%s rawLen=%d cleanLen=%d (removed or replaced with the Spanish placeholder before client/TTS output).",
                    sanitize_tag or "-",
                    (sanitize_tag if sanitize_tag else "marker"),
                    len(lm_output.text or ""),
                    len(clean_text or ""),
                )
                # Use the placeholder when sanitization removes the entire chunk,
                # avoiding empty text that would trigger the TTS greeting fallback.
                if not clean_text.strip():
                    clean_text = _APIPILOT_SPANISH_PLACEHOLDER

        if self.text_output_queue is not None:
            event = AssistantTextEvent(
                text=clean_text,
                turn_id=lm_output.turn_id,
                turn_revision=lm_output.turn_revision,
                cancel_generation=lm_output.cancel_generation,
            )
            if lm_output.tools:
                event.tools = lm_output.tools
                logger.info(f"Sending to clients: text='{clean_text}', tools={[t.name for t in lm_output.tools]}")
            else:
                logger.debug(f"Sending to clients: text='{clean_text}' (no tools)")
            self.text_output_queue.put(event)

        if clean_text and response_wants_audio(lm_output.response):
            logger.debug(f"Forwarding to TTS: '{clean_text}' (original marker was: {sanitize_tag if any_sanitized else 'none'})")
            yield TTSInput(
                text=clean_text,
                language_code=lm_output.language_code,
                runtime_config=lm_output.runtime_config,
                response=lm_output.response,
                turn_id=lm_output.turn_id,
                turn_revision=lm_output.turn_revision,
                speech_stopped_at_s=lm_output.speech_stopped_at_s,
                cancel_generation=lm_output.cancel_generation,
            )