"""Simulation runner - orchestrates the simulator pipeline."""

import asyncio
import time
from typing import Any

from loguru import logger
from openai.types.chat import ChatCompletionMessageParam
from pipecat.audio.vad.silero import SileroVADAnalyzer
from pipecat.audio.vad.vad_analyzer import VADParams
from pipecat.frames.frames import EndFrame, LLMRunFrame
from pipecat.observers.turn_tracking_observer import TurnTrackingObserver
from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.runner import PipelineRunner
from pipecat.pipeline.task import PipelineParams, PipelineTask
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.aggregators.llm_response_universal import LLMContextAggregatorPair

from voiceground.observer import VoicegroundObserver
from voiceground.simulation.bridge import (
    VoicegroundBridgeTransport,
    VoicegroundBridgeTransportParams,
)
from voiceground.simulation.config import VoicegroundSimulatorConfig
from voiceground.simulation.reporter import SimulationReporter
from voiceground.simulation.results import (
    VoicegroundSimulationResults,
    VoicegroundTerminationReason,
)


def build_simulator_pipeline(
    simulator_transport: VoicegroundBridgeTransport,
    config: VoicegroundSimulatorConfig,
) -> Pipeline:
    """Build the simulator pipeline that acts as the simulated user."""
    messages: ChatCompletionMessageParam = {"role": "system", "content": config.system_prompt}
    context = LLMContext(messages=[messages])
    context_aggregator = LLMContextAggregatorPair(context)

    return Pipeline(
        [
            simulator_transport.input(),
            config.stt,
            context_aggregator.user(),
            config.llm,
            config.tts,
            simulator_transport.output(),
            context_aggregator.assistant(),
        ]
    )


class VoicegroundSimulation:
    """A simulation environment for testing Pipecat bots.

    Use as an async context manager to prepare a simulation, then run your
    bot pipeline normally against it. The simulation provides a transport
    that is a drop-in replacement for LocalAudioTransport or WebsocketTransport.

    Example:
        async with VoicegroundSimulation(config) as simulation:
            # Use simulation.transport just like LocalAudioTransport
            pipeline = Pipeline([
                simulation.transport.input(),   # Standard transport.input()
                stt, llm, tts,
                simulation.transport.output(),  # Standard transport.output()
            ])

            # Run your bot normally
            task = PipelineTask(pipeline, observers=[...])
            runner = PipelineRunner()
            await runner.run(task)

        # Results available after context exits
        print(simulation.results.transcript)
    """

    def __init__(self, config: VoicegroundSimulatorConfig):
        """Initialize the simulation.

        Args:
            config: Configuration for the simulated user.
        """
        self._config = config

        # Create transport params with VAD enabled
        vad_params = VADParams(stop_secs=2.5)
        bot_params = VoicegroundBridgeTransportParams(
            audio_in_enabled=True,
            audio_out_enabled=True,
            vad_analyzer=SileroVADAnalyzer(params=vad_params),
        )
        sim_params = VoicegroundBridgeTransportParams(
            audio_in_enabled=True,
            audio_out_enabled=True,
            vad_analyzer=SileroVADAnalyzer(params=vad_params),
        )

        # Bot transport - this is what the user's bot will use
        self._bot_transport = VoicegroundBridgeTransport(params=bot_params)

        # Simulator transport - paired with bot transport for audio routing
        self._simulator_transport = self._bot_transport.create_paired_transport(params=sim_params)

        # Register handler to inject first message when bot "connects"
        self._simulator_transport.add_event_handler("on_client_connected", self._on_bot_connected)

        # Simulation state
        self._simulator_task: PipelineTask | None = None
        self._simulator_runner: PipelineRunner | None = None
        self._simulator_future: asyncio.Task[None] | None = None
        self._timeout_task: asyncio.Task[None] | None = None
        self._stop_monitor_task: asyncio.Task[None] | None = None
        self._simulation_observer: VoicegroundObserver | None = None
        self._start_time: float = 0.0
        self._results: VoicegroundSimulationResults | None = None
        self._termination_reason: VoicegroundTerminationReason = (
            VoicegroundTerminationReason.UNKNOWN
        )
        self._stop_event = asyncio.Event()

        self._turn_observer = TurnTrackingObserver()
        self._turn_observer.add_event_handler("on_turn_ended", self._on_turn_ended)
        self._simulation_reporter = SimulationReporter()

    @property
    def transport(self) -> VoicegroundBridgeTransport:
        """Get the transport to use in your bot pipeline.

        This is a drop-in replacement for LocalAudioTransport or any other
        Pipecat transport. Use it exactly as you would use LocalAudioTransport:

            pipeline = Pipeline([
                simulation.transport.input(),   # Receives simulated user audio
                stt, llm, tts,
                simulation.transport.output(),  # Sends bot audio to simulator
            ])
        """
        return self._bot_transport

    @property
    def results(self) -> VoicegroundSimulationResults:
        """Get the simulation results.

        Only available after the simulation has completed (context exited).
        """
        if self._results is None:
            raise RuntimeError("Simulation has not completed yet. Use 'async with' context.")
        return self._results

    async def _on_bot_connected(self, transport: VoicegroundBridgeTransport) -> None:
        """Handle bot connection - trigger the simulator to start the conversation."""
        if self._config.initiate_conversation and self._simulator_task:
            logger.debug("Bot connected, triggering simulator to start conversation")
            await self._simulator_task.queue_frame(LLMRunFrame())

    async def _on_turn_ended(
        self,
        observer: TurnTrackingObserver,
        turn_count: int,
        duration: float,
        was_interrupted: bool,
    ) -> None:
        """Handle turn end events to check for max turns."""
        logger.debug(
            f"Turn {turn_count} ended (duration={duration:.2f}s, interrupted={was_interrupted})"
        )

        if turn_count >= self._config.max_turns:
            logger.info(f"Turn limit ({self._config.max_turns}) reached, stopping simulation")
            self._termination_reason = VoicegroundTerminationReason.MAX_TURNS
            self._stop_event.set()

    async def __aenter__(self) -> "VoicegroundSimulation":
        """Start the simulation."""
        self._start_time = time.time()
        self._stop_event.clear()

        simulator_pipeline = build_simulator_pipeline(self._simulator_transport, self._config)
        self._simulation_observer = VoicegroundObserver(reporters=[self._simulation_reporter])

        self._simulator_task = PipelineTask(
            simulator_pipeline,
            params=PipelineParams(allow_interruptions=True, enable_metrics=True),
            observers=[self._turn_observer, self._simulation_observer],
        )
        self._simulator_runner = PipelineRunner()

        self._simulator_future = asyncio.create_task(self._run_simulator())
        self._timeout_task = asyncio.create_task(self._monitor_timeout())
        self._stop_monitor_task = asyncio.create_task(self._monitor_stop_event())

        return self

    async def _monitor_timeout(self) -> None:
        """Monitor for timeout and trigger stop if exceeded."""
        try:
            await asyncio.sleep(self._config.timeout_seconds)
            if not self._stop_event.is_set():
                logger.info(
                    f"Timeout ({self._config.timeout_seconds}s) reached, stopping simulation"
                )
                self._termination_reason = VoicegroundTerminationReason.TIMEOUT
                self._stop_event.set()
        except asyncio.CancelledError:
            pass

    async def _monitor_stop_event(self) -> None:
        """Monitor stop event and send EndFrame when triggered."""
        try:
            await self._stop_event.wait()
            # Stop both pipelines
            if self._simulator_task:
                await self._simulator_task.queue_frame(EndFrame())
        except asyncio.CancelledError:
            pass

    async def _run_simulator(self) -> None:
        """Run the simulator pipeline."""
        if self._simulator_task and self._simulator_runner:
            await self._simulator_runner.run(self._simulator_task)

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: Any,  # traceback object type varies by Python version
    ) -> None:
        """Stop the simulation and collect results."""
        await self._cancel_task(self._timeout_task)
        await self._cancel_task(self._stop_monitor_task)

        if self._simulator_task and not self._stop_event.is_set():
            await self._simulator_task.queue_frame(EndFrame())

        if self._simulator_future:
            try:
                await asyncio.wait_for(self._simulator_future, timeout=5.0)
            except asyncio.TimeoutError:
                await self._cancel_task(self._simulator_future)

        await self._bot_transport.close()

        self._results = VoicegroundSimulationResults(
            transcript=self._simulation_reporter.transcript_collector.entries,
            events=self._simulation_reporter.events,
            turn_count=self._turn_observer._turn_count,
            duration_seconds=time.time() - self._start_time,
            termination_reason=self._termination_reason,
        )

    async def _cancel_task(self, task: asyncio.Task | None) -> None:
        """Cancel a task and wait for it to finish."""
        if task:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
