from __future__ import annotations

import asyncio
import re
import string
import uuid
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, List, Optional

import structlog

from rasa.core.channels.voice_stream.asr.asr_event import InterimTranscriptReady
from rasa.shared.nlu.constants import (
    COMMANDS,
    ENTITIES,
    INTENT,
    INTENT_NAME_KEY,
    PREDICTED_CONFIDENCE_KEY,
    TEXT,
)

if TYPE_CHECKING:
    from rasa.core.processor import MessageProcessor

logger = structlog.get_logger()

# Words that ASR frequently adds/removes but carry no semantic meaning.
_FILLER_WORDS = frozenset({
    "uh", "uhh", "um", "umm", "hmm", "hm", "ah", "ahh",
    "like", "please", "okay", "ok", "so", "well", "right",
})

_WHITESPACE_RE = re.compile(r"\s+")


def _normalize_for_comparison(text: str) -> str:
    """Normalize transcript text for fuzzy comparison.

    Strips punctuation, lowercases, and removes common filler/backchannel
    words that ASR engines frequently add or remove in final transcripts.
    """
    text = text.lower()
    text = text.translate(str.maketrans("", "", string.punctuation))
    words = _WHITESPACE_RE.split(text.strip())
    words = [w for w in words if w and w not in _FILLER_WORDS]
    return " ".join(words)


@dataclass
class SpeculativeResult:
    """Cached result from speculative processing."""

    transcript: str
    commands: List[Dict[str, Any]]
    # Agent speculation results (populated when speculate_agent is enabled)
    agent_invoked: bool = False
    filler_sent: bool = False
    pending_tool_calls: Optional[List[Dict[str, Any]]] = None
    llm_response_data: Optional[Dict[str, Any]] = None
    bot_uttered_data: Optional[Dict[str, Any]] = None
    agent_llm_messages: Optional[List[Dict[str, Any]]] = None


@dataclass
class SpeculationConfig:
    """Configuration for speculative processing."""

    enabled: bool = False
    min_new_tokens: int = 3
    speculate_agent: bool = False
    allowed_speculative_commands: List[str] = field(
        default_factory=lambda: ["continue agent", "restart agent"]
    )

    @classmethod
    def from_dict(cls, data: Optional[Dict[str, Any]]) -> SpeculationConfig:
        if not data:
            return cls()
        return cls(
            enabled=data.get("enabled", False),
            min_new_tokens=data.get("min_new_tokens", 3),
            speculate_agent=data.get("speculate_agent", False),
            allowed_speculative_commands=data.get(
                "allowed_speculative_commands",
                ["continue agent", "restart agent"],
            ),
        )


def _build_full_command_string(cmd: Dict[str, Any]) -> str:
    """Build a full command string from a command dict for matching.

    Command dicts have a "command" key with the base type (e.g. "start flow")
    and additional fields for parameters (e.g. "flow": "bill_inquiry").
    This reconstructs the full string like "start flow bill_inquiry".
    """
    base = cmd.get("command", "")
    # Parameterized commands: append the relevant parameter
    if base == "start flow" and "flow" in cmd:
        return f"start flow {cmd['flow']}"
    if base == "restart agent" and "agent_id" in cmd:
        return f"restart agent {cmd['agent_id']}"
    return base


def _are_commands_agent_speculative(
    commands: List[Dict[str, Any]],
    allowed_commands: List[str],
) -> bool:
    """Check if all commands in the list are eligible for agent speculation."""
    if not commands:
        return False
    return all(
        _build_full_command_string(cmd) in allowed_commands
        for cmd in commands
    )


class SpeculativeProcessor:
    """Manages speculative command generation and agent execution on interim transcripts.

    Runs the NLU graph on partial ASR transcripts to pre-compute commands.
    When commands are agent-compatible (continue/restart agent), also runs the
    agent's first LLM turn speculatively (with filler streaming to user).
    When the final transcript arrives, if it matches the speculative transcript,
    the pre-computed results are reused.
    """

    def __init__(self, config: SpeculationConfig) -> None:
        self.config = config
        self._speculative_cmd_task: Optional[asyncio.Task] = None
        self._speculative_agent_task: Optional[asyncio.Task] = None
        self._speculative_result: Optional[SpeculativeResult] = None
        self._latest_speculative_transcript: Optional[str] = None
        self._soft_miss_result: Optional[SpeculativeResult] = None
        # Late agent result: stored when a speculative agent task completes
        # after on_final_transcript() has already cleared the transcript
        # (making the staleness check fail). Instead of discarding the valid
        # result, we preserve it here so the voice channel can await and
        # reuse it when has_new_information is SAME.
        self._late_agent_result: Optional[SpeculativeResult] = None
        self._pending_pipeline_task: Optional[asyncio.Task] = None

    async def on_interim_transcript(
        self,
        event: InterimTranscriptReady,
        processor: MessageProcessor,
        sender_id: str,
    ) -> None:
        """Handle an interim transcript ready event by starting speculative cmd gen.

        Uses two-layer matching against any existing speculative result:
        1. Normalized comparison — if the new interim is cosmetically identical
           to the cached result, skip command gen entirely.
        2. If normalized miss, run command gen with the previous speculative
           transcript so the LLM can assess has_new_information. If SAME,
           the cached result (including any agent data) is preserved.

        Args:
            event: The interim transcript event with accumulated text.
            processor: The message processor for running the NLU graph.
            sender_id: The conversation sender ID for fetching tracker.
        """
        new_text = event.text
        trace_id = uuid.uuid4().hex[:8]

        # Layer 1: If we have an existing result, check normalized match.
        # If the new interim is only cosmetically different, skip command gen.
        if self._speculative_result and _normalize_for_comparison(
            self._speculative_result.transcript
        ) == _normalize_for_comparison(new_text):
            logger.info(
                "speculative.interim_normalized_hit",
                trace_id=trace_id,
                new_text=new_text,
                cached_text=self._speculative_result.transcript,
            )
            self._latest_speculative_transcript = new_text
            return

        # Layer 2: Normalized miss — need to run command gen.
        # Capture previous transcript for semantic diff in the command gen.
        # Prefer completed result's transcript; fall back to the latest
        # speculative transcript (even if its task hasn't completed yet)
        # so that interim-to-interim semantic diff always fires.
        previous_transcript = (
            self._speculative_result.transcript
            if self._speculative_result
            else self._latest_speculative_transcript
        )

        logger.info(
            "speculative.starting_pipeline",
            trace_id=trace_id,
            text=new_text,
            previous_transcript=previous_transcript,
            type="interim",
        )

        # Mark previous results as stale. We do NOT cancel in-flight tasks
        # because litellm's internal asyncio tasks don't handle cancellation
        # cleanly (orphaned task exceptions). Instead, each task checks
        # _latest_speculative_transcript before storing its result — stale
        # tasks simply discard their output.
        self._latest_speculative_transcript = new_text

        # Start new speculative command generation with previous transcript
        # context for semantic diff.
        self._speculative_cmd_task = asyncio.create_task(
            self._run_speculative_pipeline(
                new_text, processor, sender_id, previous_transcript, trace_id
            )
        )

    def on_final_transcript(self, final_text: str) -> Optional[SpeculativeResult]:
        """Check if speculative result can be reused for the final transcript.

        Uses two-layer matching:
        1. Normalized comparison (strips punctuation, case, filler words)
        2. If normalized miss but speculative processing was attempted, stores
           a "soft miss" context so the voice channel can pass it to the command
           generator for LLM-based semantic diff.

        Args:
            final_text: The final ASR transcript.

        Returns:
            SpeculativeResult if cache hit (normalized transcripts match), else None.
        """
        # Capture the speculative transcript before clearing — we need it
        # for soft miss even if the speculative task hasn't completed yet.
        speculative_transcript = self._latest_speculative_transcript

        # Preserve the in-flight pipeline task so the voice channel can
        # await late-arriving agent results when has_new_information=SAME.
        # Do NOT clear _late_agent_result here — if the speculative agent
        # already completed and stored its result before this call, we
        # need to keep it for await_speculative_agent_result() to return.
        self._pending_pipeline_task = self._speculative_cmd_task

        # Mark transcript as final so any in-flight tasks will store
        # their results in _late_agent_result instead of _speculative_result
        # (see the staleness check in _run_speculative_agent).
        self._latest_speculative_transcript = None
        self._soft_miss_result = None

        result = self._speculative_result
        self._speculative_result = None

        if result:
            # Layer 1: Normalized comparison
            if _normalize_for_comparison(
                result.transcript
            ) == _normalize_for_comparison(final_text):
                # Update transcript to the final text for downstream consistency
                result.transcript = final_text
                logger.info(
                    "speculative.cache_hit",
                    match_type="normalized",
                    transcript=final_text,
                    agent_invoked=result.agent_invoked,
                    filler_sent=result.filler_sent,
                )
                return result

            # Layer 2: Normalized miss with completed result — store full
            # soft miss with cached commands and agent data.
            logger.info(
                "speculative.soft_miss",
                final_text=final_text,
                speculative_text=result.transcript,
                agent_invoked=result.agent_invoked,
            )
            self._soft_miss_result = result
            return None

        # No completed result, but if speculative processing was attempted
        # (task in-flight or completed with no result), still store the
        # transcript for semantic diff in the command generator.
        if speculative_transcript:
            logger.info(
                "speculative.soft_miss_in_flight",
                final_text=final_text,
                speculative_text=speculative_transcript,
            )
            self._soft_miss_result = SpeculativeResult(
                transcript=speculative_transcript,
                commands=[],
            )
        else:
            logger.debug(
                "speculative.no_result",
                final_text=final_text,
            )
        return None

    def get_soft_miss_result(self) -> Optional[SpeculativeResult]:
        """Get the speculative result from a soft miss, if any.

        A soft miss occurs when the normalized transcripts don't match but
        a speculative result exists. The voice channel should pass the
        speculative transcript and cached agent data through metadata so the
        command generator can perform LLM-based semantic comparison.

        Returns:
            The speculative result if a soft miss occurred, else None.
            Calling this clears the soft miss state.
        """
        result = self._soft_miss_result
        self._soft_miss_result = None
        return result

    async def await_speculative_agent_result(
        self, timeout: float = 3.0
    ) -> Optional[SpeculativeResult]:
        """Await a late-arriving speculative agent result.

        When on_final_transcript() is called while a speculative pipeline
        task is still in-flight, the agent result will arrive late and be
        stored in _late_agent_result. This method awaits the pending task
        (with timeout) and returns the late result if available.

        Called by the voice channel when the command generator returns
        has_new_information=SAME, indicating the final transcript is
        semantically identical and the cached agent result should be reused.

        Args:
            timeout: Maximum seconds to wait for the pending task.

        Returns:
            The late SpeculativeResult if the agent completed, else None.
        """
        task = self._pending_pipeline_task
        if task is None or task.done():
            result = self._late_agent_result
            self._late_agent_result = None
            self._pending_pipeline_task = None
            return result

        try:
            await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
        except asyncio.TimeoutError:
            logger.debug("speculative.await_agent_timeout", timeout=timeout)
        except Exception:
            logger.debug("speculative.await_agent_error", exc_info=True)

        result = self._late_agent_result
        self._late_agent_result = None
        self._pending_pipeline_task = None
        return result

    def cancel_all(self) -> None:
        """Discard all speculative results and invalidate in-flight tasks.

        In-flight tasks are not forcefully cancelled (to avoid litellm's
        orphaned-task issues). Instead, setting _latest_speculative_transcript
        to None ensures any completing task will see its result as stale and
        discard it.
        """
        self._speculative_result = None
        self._latest_speculative_transcript = None
        self._soft_miss_result = None
        self._late_agent_result = None
        self._pending_pipeline_task = None

    async def _run_speculative_pipeline(
        self,
        text: str,
        processor: MessageProcessor,
        sender_id: str,
        previous_transcript: Optional[str] = None,
        trace_id: Optional[str] = None,
    ) -> None:
        """Run the full speculative pipeline: command gen, then optionally agent.

        After command generation completes, checks if the commands are eligible
        for agent speculation. If so, starts the speculative agent first-turn.

        If previous_transcript is provided, the command generator will also
        assess whether the new transcript has semantically new information.
        If not, the existing cached result is preserved instead of being
        overwritten.
        """
        _tid = trace_id or uuid.uuid4().hex[:8]
        try:
            # Step 1: Run speculative command generation
            commands, has_new_info = await self._run_speculative_command_gen(
                text, processor, sender_id, previous_transcript, _tid
            )
            if commands is None:
                return  # Failed or stale

            # If the command generator determined there's no new semantic
            # information and we have a cached result, preserve it — just
            # update the transcript text for downstream matching.
            if not has_new_info and self._speculative_result:
                self._speculative_result.transcript = text
                logger.info(
                    "speculative.semantic_same_reusing_cached",
                    trace_id=_tid,
                    transcript=text,
                    previous_transcript=previous_transcript,
                    agent_invoked=self._speculative_result.agent_invoked,
                )
                return

            # Store command-only result immediately
            self._speculative_result = SpeculativeResult(
                transcript=text,
                commands=commands,
            )

            logger.info(
                "speculative.command_gen_complete",
                trace_id=_tid,
                transcript=text,
                num_commands=len(commands),
                commands=[c.get("command", "unknown") for c in commands],
                has_new_info=has_new_info,
            )

            # Step 2: Check if we should speculatively invoke the agent.
            # For any command in the allowed list (continue agent, restart agent,
            # start flow X, etc.), we attempt to find the agent that would be
            # invoked and run its first LLM turn speculatively.
            if (
                self.config.speculate_agent
                and _are_commands_agent_speculative(
                    commands, self.config.allowed_speculative_commands
                )
            ):
                logger.info(
                    "speculative.starting_agent",
                    trace_id=_tid,
                    transcript=text,
                )
                await self._run_speculative_agent(
                    text, commands, processor, sender_id
                )

        except Exception:
            logger.exception("speculative.pipeline_error", trace_id=_tid, text=text)

    async def _run_speculative_command_gen(
        self,
        text: str,
        processor: MessageProcessor,
        sender_id: str,
        previous_transcript: Optional[str] = None,
        trace_id: Optional[str] = None,
    ) -> tuple[Optional[List[Dict[str, Any]]], bool]:
        """Run speculative command generation using the NLU graph.

        Args:
            text: The current interim transcript text.
            processor: The message processor for running the NLU graph.
            sender_id: The conversation sender ID.
            previous_transcript: If provided, the command generator will also
                assess whether the new transcript has new semantic information
                compared to this previous transcript.
            trace_id: Correlation ID for matching request/response logs.

        Returns:
            Tuple of (commands, has_new_information).
            commands is None if failed/stale.
            has_new_information defaults to True when no previous transcript
            is provided or when the flag is missing from the LLM output.
        """
        from rasa.agents.constants import (
            HAS_NEW_INFORMATION_KEY,
            SPECULATIVE_PREVIOUS_TRANSCRIPT_KEY,
            SPECULATIVE_TRACE_ID_KEY,
        )

        _tid = trace_id or uuid.uuid4().hex[:8]

        # Get a read-only tracker snapshot
        tracker = await processor.tracker_store.get_or_create_tracker(sender_id)
        if tracker is None:
            logger.debug(
                "speculative.no_tracker", trace_id=_tid, sender_id=sender_id
            )
            return None, True

        # Create a lightweight UserMessage for parsing.
        # Include trace_id in metadata for log correlation in command gen.
        # If we have a previous transcript, include it so the command
        # generator can assess has_new_information.
        from rasa.core.channels.channel import UserMessage

        metadata: Dict[str, Any] = {SPECULATIVE_TRACE_ID_KEY: _tid}
        if previous_transcript:
            metadata[SPECULATIVE_PREVIOUS_TRANSCRIPT_KEY] = previous_transcript

        speculative_message = UserMessage(
            text=text, sender_id=sender_id, metadata=metadata
        )

        logger.debug(
            "speculative.command_gen_start",
            trace_id=_tid,
            text=text,
            previous_transcript=previous_transcript,
        )

        # Run the NLU graph (command generator) with the interim text
        parse_data = await processor.parse_message(speculative_message, tracker)
        commands = parse_data.get(COMMANDS, [])
        has_new_info = parse_data.get(HAS_NEW_INFORMATION_KEY, True)

        # Check staleness
        if self._latest_speculative_transcript != text:
            logger.debug(
                "speculative.stale_result",
                trace_id=_tid,
                text=text,
                latest=self._latest_speculative_transcript,
            )
            return None, True

        logger.debug(
            "speculative.command_gen_done",
            trace_id=_tid,
            text=text,
            commands=[c.get("command", "unknown") for c in commands],
            has_new_info=has_new_info,
        )

        return commands, has_new_info

    async def _run_speculative_agent(
        self,
        text: str,
        commands: List[Dict[str, Any]],
        processor: MessageProcessor,
        sender_id: str,
    ) -> None:
        """Run the agent's first LLM turn speculatively.

        This method handles both cases:
        - Agent already on the stack (continue agent / restart agent)
        - Agent reachable via flow advancement (start flow X → agent call)

        It applies the speculative commands to a tracker copy, runs the
        command processor to update the dialogue stack, then finds the
        agent that would be invoked. A CollectingOutputChannel is used so
        nothing reaches the user during speculation. Tool calls are NEVER
        executed.
        """
        try:
            from rasa.agents.agent_manager import AgentManager
            from rasa.core.channels.channel import UserMessage
            from rasa.core.config.configuration import Configuration
            from rasa.core.policies.flows.agent_executor import _prepare_agent_input
            from rasa.core.policies.flows.flow_executor import select_next_step
            from rasa.dialogue_understanding.stack.frames.flow_stack_frame import (
                AgentStackFrame,
            )
            from rasa.shared.agents.utils import get_protocol_type
            from rasa.shared.core.events import UserUttered
            from rasa.shared.core.flows.steps import CallFlowStep
            from rasa.dialogue_understanding.stack.frames import BaseFlowStackFrame

            # Get a fresh tracker copy for speculative processing
            tracker = await processor.tracker_store.get_or_create_tracker(sender_id)
            if tracker is None:
                return

            # Apply the speculative user message + commands to the tracker copy
            tracker_copy = tracker.copy()
            parse_data = self.build_parse_data_from_result(
                SpeculativeResult(transcript=text, commands=commands)
            )
            tracker_copy.update(
                UserUttered(
                    text,
                    parse_data.get(INTENT, {}),
                    parse_data.get(ENTITIES, []),
                    parse_data,
                    message_id="speculative",
                ),
                processor.domain,
            )

            # Run the command processor to update the dialogue stack
            tracker_copy = await processor.run_command_processor(tracker_copy)

            # Now find the agent that would be invoked
            agent_step, agent_stack_frame = await self._find_next_agent_step(
                tracker_copy, processor
            )

            if not agent_step:
                logger.debug(
                    "speculative.no_agent_step_found",
                    text=text,
                )
                return

            # Build AgentInput
            agent_input = _prepare_agent_input(
                agent_stack_frame,
                agent_step,
                tracker_copy,
                list(tracker_copy.slots.values()),
            )

            # Get the agent instance
            protocol_type = get_protocol_type(
                agent_step,
                Configuration.get_instance().available_agents.get_agent_config(
                    agent_step.call
                ),
            )
            agent = AgentManager().get_agent(agent_step.call, protocol_type)

            # Process input
            processed_input = await agent.process_input(agent_input)

            # Use a collecting channel so nothing is sent to the user
            # during speculation. The response is cached and only sent
            # on cache hit through the real pipeline.
            from rasa.core.channels.channel import CollectingOutputChannel

            collecting_channel = CollectingOutputChannel()

            # Run first LLM turn only (no tool execution)
            llm_response, bot_uttered, messages = await agent.run_first_turn_only(
                processed_input, collecting_channel
            )

            # Check staleness — if the transcript has moved on (new interim
            # or final arrived), store the result as a late agent result
            # instead of discarding. The voice channel can await and reuse
            # it when the command generator returns has_new_information=SAME.
            is_stale = self._latest_speculative_transcript != text

            if llm_response is None:
                logger.debug("speculative.agent_no_response", text=text)
                return

            # Build the agent result
            if llm_response.tool_calls:
                # Case B: LLM returned tool calls (+ possibly filler text)
                agent_result = SpeculativeResult(
                    transcript=text,
                    commands=commands,
                    agent_invoked=True,
                    filler_sent=bot_uttered is not None,
                    pending_tool_calls=[
                        tc.model_dump() for tc in llm_response.tool_calls
                    ],
                    llm_response_data=llm_response.to_dict(),
                    agent_llm_messages=messages,
                )
            else:
                # Case A: Text-only response (final answer, no tools)
                agent_result = SpeculativeResult(
                    transcript=text,
                    commands=commands,
                    agent_invoked=True,
                    filler_sent=False,
                    llm_response_data=llm_response.to_dict(),
                    bot_uttered_data=(
                        bot_uttered.as_dict() if bot_uttered else None
                    ),
                    agent_llm_messages=messages,
                )

            if is_stale:
                # Store as late result for the voice channel to await
                logger.info(
                    "speculative.agent_late_result_stored",
                    transcript=text,
                    agent_id=agent_step.call,
                    has_tool_calls=bool(llm_response.tool_calls),
                )
                self._late_agent_result = agent_result
            else:
                self._speculative_result = agent_result

            logger.info(
                "speculative.agent_first_turn_complete",
                transcript=text,
                agent_id=agent_step.call,
                has_tool_calls=bool(llm_response.tool_calls),
                filler_sent=bot_uttered is not None,
                is_stale=is_stale,
            )

        except Exception:
            logger.exception("speculative.agent_error", text=text)

    @staticmethod
    async def _find_next_agent_step(
        tracker: Any,
        processor: MessageProcessor,
    ) -> tuple:
        """Walk the dialogue stack and flow steps to find the next agent call.

        After commands have been applied to the tracker, this method walks
        through the flow steps (simulating what advance_flows_until_next_action
        would do) to find if the next actionable step is an agent call.

        Only walks through "pass-through" steps (set_slots, noop, link, call
        to sub-flows). Stops at collect steps, action steps (utter), or
        agent calls.

        Returns:
            (CallFlowStep, AgentStackFrame or None) if an agent call is found,
            (None, None) otherwise.
        """
        from rasa.core.policies.flows.flow_executor import select_next_step
        from rasa.dialogue_understanding.stack.frames import BaseFlowStackFrame
        from rasa.dialogue_understanding.stack.frames.flow_stack_frame import (
            AgentStackFrame,
        )
        from rasa.shared.core.flows.steps import CallFlowStep
        from rasa.shared.core.flows.steps.collect import CollectInformationFlowStep
        from rasa.shared.core.flows.steps.action import ActionFlowStep
        from rasa.shared.core.flows.steps.set_slots import SetSlotsFlowStep
        from rasa.shared.core.flows.steps.no_operation import NoOperationFlowStep

        all_flows = await processor.get_flows()

        # First, check if there's already an agent stack frame on top
        # (for continue agent / restart agent commands)
        top_frame = tracker.stack.top()
        if isinstance(top_frame, AgentStackFrame):
            flow = all_flows.flow_by_id(top_frame.flow_id)
            if flow:
                for s in flow.steps:
                    if (
                        isinstance(s, CallFlowStep)
                        and s.id == top_frame.step_id
                        and s.is_calling_agent()
                    ):
                        return s, top_frame
            return None, None

        # Walk through flow steps to find the next agent call
        if not isinstance(top_frame, BaseFlowStackFrame):
            return None, None

        flow = top_frame.flow(all_flows)
        if not flow:
            return None, None

        # Walk up to MAX_WALK_STEPS to find an agent call
        MAX_WALK_STEPS = 20
        current_step = top_frame.step(all_flows)

        for _ in range(MAX_WALK_STEPS):
            next_step = select_next_step(
                current_step, flow, tracker.stack, tracker
            )
            if next_step is None:
                break

            if isinstance(next_step, CallFlowStep) and next_step.is_calling_agent():
                # Found the agent call step
                agent_stack_frame = tracker.stack.find_agent_stack_frame_by_agent(
                    next_step.call
                )
                return next_step, agent_stack_frame

            # Only walk through "transparent" steps that don't produce
            # user-visible side effects
            if isinstance(next_step, (SetSlotsFlowStep, NoOperationFlowStep)):
                current_step = next_step
                continue

            # Stop at steps that produce side effects (collect, action/utter,
            # non-agent calls, etc.)
            break

        return None, None

    @staticmethod
    def build_parse_data_from_result(result: SpeculativeResult) -> Dict[str, Any]:
        """Build parse_data dict from a speculative result for UserMessage."""
        return {
            TEXT: result.transcript,
            COMMANDS: result.commands,
            INTENT: {INTENT_NAME_KEY: None, PREDICTED_CONFIDENCE_KEY: 0.0},
            ENTITIES: [],
        }
