import asyncio

from loguru import logger

from api.constants import ENABLE_CALL_RECORDING_UPLOAD
from api.db import db_client
from api.enums import PostHogEvent, WorkflowRunState
from api.services.campaign.circuit_breaker import circuit_breaker
from api.services.integrations import IntegrationRuntimeSession
from api.services.pipecat.audio_config import AudioConfig
from api.services.pipecat.audio_playback import play_audio_loop
from api.services.pipecat.in_memory_buffers import (
    InMemoryLogsBuffer,
    InMemoryRecordingBuffers,
)
from api.services.pipecat.pipeline_metrics_aggregator import PipelineMetricsAggregator
from api.services.pipecat.termination_funnel_processor import (
    TerminationFunnelProcessor,
)
from api.services.pipecat.tracing_config import get_trace_url
from api.services.pipecat.transcript_log_coordinator import TranscriptLogCoordinator
from api.services.posthog_client import capture_event
from api.services.workflow.initial_context import merge_external_initial_context
from api.services.workflow.pipecat_engine import PipecatEngine
from api.services.workflow_run_artifacts import upload_workflow_run_artifacts
from api.tasks.arq import enqueue_job
from api.tasks.function_names import FunctionNames
from pipecat.frames.frames import (
    ErrorFrame,
    Frame,
)
from pipecat.pipeline.worker import PipelineWorker
from pipecat.processors.audio.audio_buffer_processor import AudioBufferProcessor
from pipecat.utils.enums import EndTaskReason


async def _capture_call_event(
    workflow_run_id: int,
    user_provider_id: str | None,
    event: str,
    extra_properties: dict | None = None,
) -> None:
    """Look up workflow_run for call metadata and fire a PostHog event.
    Meant to be run via asyncio.create_task() so it never blocks the pipeline."""
    try:
        workflow_run = await db_client.get_workflow_run_by_id(workflow_run_id)
        properties = {
            "workflow_run_id": workflow_run_id,
            "workflow_id": workflow_run.workflow_id if workflow_run else None,
            "call_type": workflow_run.mode if workflow_run else None,
            "call_direction": (
                (workflow_run.initial_context or {}).get("direction", "outbound")
                if workflow_run
                else None
            ),
        }
        if extra_properties:
            properties.update(extra_properties)
        capture_event(
            distinct_id=user_provider_id,
            event=event,
            properties=properties,
        )
    except Exception:
        logger.exception(f"Background PostHog capture failed for '{event}'")


def register_event_handlers(
    task: PipelineWorker,
    transport,
    workflow_run_id: int,
    engine: PipecatEngine,
    audio_buffer: AudioBufferProcessor,
    in_memory_logs_buffer: InMemoryLogsBuffer,
    transcript_log_coordinator: TranscriptLogCoordinator,
    pipeline_metrics_aggregator: PipelineMetricsAggregator,
    termination_funnel: TerminationFunnelProcessor,
    audio_config=AudioConfig,
    pre_call_fetch_task: asyncio.Task | None = None,
    user_provider_id: str | None = None,
    integration_runtime_sessions: list[IntegrationRuntimeSession] | None = None,
    include_transcript_end_timestamps: bool = False,
    answer_supervisor=None,
):
    """Register all event handlers for transport and task events.

    Returns:
        In-memory recording buffers for use by other handlers.
    """
    # Initialize in-memory buffers with proper audio configuration
    sample_rate = audio_config.pipeline_sample_rate if audio_config else 16000
    num_channels = 1  # Pipeline audio is always mono

    logger.debug(
        f"Initializing audio buffer for workflow {workflow_run_id} "
        f"with sample_rate={sample_rate}Hz, channels={num_channels}"
    )

    in_memory_audio_buffers = InMemoryRecordingBuffers(
        workflow_run_id=workflow_run_id,
        sample_rate=sample_rate,
        num_channels=num_channels,
    )
    # Track both events to ensure the initial response is only triggered after both occur
    ready_state = {
        "pipeline_started": False,
        "client_connected": False,
        "initial_response_triggered": False,
    }

    async def maybe_trigger_initial_response():
        """Start the conversation after both pipeline_started and client_connected events.

        If a pre-call fetch is in progress, plays a ringer while waiting for the
        response, then merges the result into the call context before proceeding.
        """
        if (
            ready_state["pipeline_started"]
            and ready_state["client_connected"]
            and not ready_state["initial_response_triggered"]
        ):
            ready_state["initial_response_triggered"] = True

            if answer_supervisor is not None:
                answer_supervisor.arm()

            asyncio.create_task(
                _capture_call_event(
                    workflow_run_id, user_provider_id, PostHogEvent.CALL_STARTED
                )
            )

            # Wait for pre-call fetch if in progress, playing ringer meanwhile
            if pre_call_fetch_task is not None:
                if not pre_call_fetch_task.done():
                    logger.info(
                        "Pre-call fetch still in progress, playing ringer while waiting"
                    )
                    stop_ringer = asyncio.Event()
                    sample_rate = audio_config.pipeline_sample_rate or 16000
                    ringer_task = asyncio.create_task(
                        play_audio_loop(
                            stop_event=stop_ringer,
                            sample_rate=sample_rate,
                            queue_frame=transport.output().queue_frame,
                        )
                    )
                    try:
                        fetch_result = await pre_call_fetch_task
                    finally:
                        stop_ringer.set()
                        await ringer_task
                else:
                    fetch_result = pre_call_fetch_task.result()

                if fetch_result:
                    engine._call_context_vars = merge_external_initial_context(
                        engine._call_context_vars, fetch_result
                    )
                    try:
                        await db_client.update_workflow_run(
                            workflow_run_id,
                            initial_context={**engine._call_context_vars},
                        )
                    except Exception as e:
                        logger.error(f"Failed to persist pre-call fetch context: {e}")
                    logger.info(
                        f"Pre-call fetch complete, merged keys: "
                        f"{list(fetch_result.keys())}"
                    )

            # Set the start node now (after pre-call fetch data is merged)
            # so that render_template() has the complete _call_context_vars.
            await engine.set_node(engine.workflow.start_node_id)
            if answer_supervisor is not None:
                await engine.handle_answer_supervision()
                return
            await engine.queue_node_opening(
                node_id=engine.workflow.start_node_id,
                previous_node_id=None,
                generate_if_no_greeting=True,
            )

    @transport.event_handler("on_client_connected")
    async def on_client_connected(_transport, _participant):
        logger.debug("In on_client_connected callback handler")
        await audio_buffer.start_recording()
        ready_state["client_connected"] = True
        await maybe_trigger_initial_response()

    @transport.event_handler("on_client_disconnected")
    async def on_client_disconnected(_transport, _participant):
        call_disposed = engine.is_call_disposed()

        logger.debug(
            f"In on_client_disconnected callback handler. Call disposed: {call_disposed}"
        )

        # Stop recordings
        await audio_buffer.stop_recording()

        # A disconnect right after a handoff is the external PBX taking the
        # customer leg, not the caller hanging up. The reason stays distinct
        # from EndTaskReason.TRANSFER_CALL on purpose: that value drives the
        # serializer's ARI bridge-transfer strategy, whereas an external-PBX
        # transfer still wants the ordinary hangup strategy to tear down our
        # own legs.
        gathered_context = await engine.get_gathered_context()
        reason = (
            EndTaskReason.CALL_TRANSFERRED.value
            if gathered_context.get("external_pbx_transferred")
            else EndTaskReason.USER_HANGUP.value
        )

        await engine.end_call_with_reason(reason, abort_immediately=True)

    @task.event_handler("on_pipeline_started")
    async def on_pipeline_started(_task: PipelineWorker, _frame: Frame):
        logger.debug("In on_pipeline_started callback handler")
        ready_state["pipeline_started"] = True
        await maybe_trigger_initial_response()

    async def _record_pipeline_error() -> None:
        try:
            workflow_run = await db_client.get_workflow_run_by_id(workflow_run_id)
            if workflow_run and workflow_run.campaign_id:
                await circuit_breaker.record_and_evaluate(
                    campaign_id=workflow_run.campaign_id,
                    is_failure=True,
                    workflow_run_id=workflow_run_id,
                    reason="pipeline_error",
                )
            asyncio.create_task(
                _capture_call_event(
                    workflow_run_id,
                    user_provider_id,
                    PostHogEvent.CALL_FAILED,
                    extra_properties={"error_reason": "pipeline_error"},
                )
            )
        except Exception as e:
            logger.error(f"Error recording circuit breaker failure: {e}", exc_info=True)

    async def dispose_call(reason: str, error: ErrorFrame | None = None) -> None:
        """Dispose of a call whose pipeline asked to stop.

        The engine queues the terminal frame itself once it is done, so this is
        what actually ends the pipeline. Everything that must be recorded
        against the run happens before that, in order, rather than in a
        separate task racing the shutdown.
        """
        if error is not None:
            logger.error(f"Pipeline error for workflow run {workflow_run_id}: {error}")
            await _record_pipeline_error()
        await engine.end_call_with_reason(reason, abort_immediately=True)

    termination_funnel.set_termination_handler(dispose_call)

    @task.event_handler("on_pipeline_error")
    async def on_pipeline_error(_task: PipelineWorker, frame: Frame):
        # Pipecat emits recoverable ErrorFrames for reconnect/retry paths. The
        # observer classifies them, but only fatal frames should dispose the call.
        if isinstance(frame, ErrorFrame) and not frame.fatal:
            return
        # Only errors raised by the input transport get this far: the funnel
        # intercepts every other fatal ErrorFrame on its way up the pipeline
        # and disposes of the call through the same path. Reaching here means
        # the worker is about to cancel the pipeline on its own, so this is a
        # race the funnel cannot close -- see TerminationFunnelProcessor.
        await dispose_call(
            EndTaskReason.PIPELINE_ERROR.value,
            frame if isinstance(frame, ErrorFrame) else None,
        )

    @task.event_handler("on_pipeline_finished")
    async def on_pipeline_finished(
        task: PipelineWorker,
        _frame: Frame,
    ):
        logger.debug("In on_pipeline_finished callback handler")

        # Turn and feedback observers run on independent queues. Drain them
        # before finalizing immutable transcripts and taking the DB snapshot.
        await task.wait_for_observers()
        await transcript_log_coordinator.flush()

        workflow_run = await db_client.get_workflow_run_by_id(workflow_run_id)

        # Stop recordings
        await audio_buffer.stop_recording()

        # Add trace URL if available (must be done before conversation tracing ends)
        if task.turn_trace_observer:
            trace_id = task.turn_trace_observer.get_trace_id()
            if trace_id:
                trace_url = get_trace_url(trace_id)
                if trace_url:
                    engine.record_context({"trace_url": trace_url})
                    logger.debug(f"Added trace URL to gathered_context: {trace_url}")

        try:
            has_user_speech = in_memory_logs_buffer.contains_user_speech()
        except Exception:
            has_user_speech = False

        engine.record_call_tags(["user_speech"] if has_user_speech else [])

        # One read, after every writer has had its say. Keys other processes put
        # on this run -- AMD results, ARI transfer state, campaign retry tags --
        # are deliberately not merged in here: `update_workflow_run` reconciles
        # them under a row lock at write time, so re-merging a staler unlocked
        # copy would only be a second, worse answer.
        gathered_context = await engine.get_gathered_context()

        # Store disposition code in workflow for dynamic filtering
        disposition_code = gathered_context.get("mapped_call_disposition")
        if disposition_code and workflow_run:
            try:
                await db_client.add_call_disposition_code(
                    workflow_run.workflow_id, disposition_code
                )
            except Exception as e:
                logger.error(
                    f"Error storing disposition code in workflow: {e}",
                    exc_info=True,
                )

        # Clean up engine resources (including voicemail detector)
        integration_logs: dict[str, object] = {}
        for runtime_session in integration_runtime_sessions or []:
            try:
                session_logs = await runtime_session.on_call_finished(
                    gathered_context=gathered_context
                )
                if session_logs:
                    integration_logs.update(session_logs)
            except Exception as e:
                logger.error(
                    f"Error finalizing integration runtime session '{runtime_session.name}': {e}",
                    exc_info=True,
                )

        await engine.cleanup()

        # ------------------------------------------------------------------
        # Close Smart-Turn WebSocket if the transport's analyzer supports it
        # ------------------------------------------------------------------
        try:
            turn_analyzer = None

            # Most transports store their params (with turn_analyzer) directly.
            if hasattr(transport, "_params") and transport._params:
                turn_analyzer = getattr(transport._params, "turn_analyzer", None)

            # Fallback: some transports expose params through input() instance.
            if turn_analyzer is None and hasattr(transport, "input"):
                try:
                    input_transport = transport.input()
                    if input_transport and hasattr(input_transport, "_params"):
                        turn_analyzer = getattr(
                            input_transport._params, "turn_analyzer", None
                        )
                except Exception:
                    pass

            if turn_analyzer and hasattr(turn_analyzer, "close"):
                await turn_analyzer.close()
                logger.debug("Closed turn analyzer websocket")
        except Exception as exc:
            logger.warning(f"Failed to close Smart-Turn analyzer gracefully: {exc}")

        usage_info = pipeline_metrics_aggregator.get_all_usage_metrics_serialized()

        logger.debug(
            f"Usage metrics: {usage_info}, Gathered context: {gathered_context}"
        )

        await db_client.update_workflow_run(
            run_id=workflow_run_id,
            usage_info=usage_info,
            gathered_context=gathered_context,
            is_completed=True,
            state=WorkflowRunState.COMPLETED.value,
        )

        asyncio.create_task(
            _capture_call_event(
                workflow_run_id, user_provider_id, PostHogEvent.CALL_COMPLETED
            )
        )

        logs_update: dict[str, object] = {}
        if not in_memory_logs_buffer.is_empty:
            try:
                feedback_events = in_memory_logs_buffer.get_events()
                logs_update["realtime_feedback_events"] = feedback_events
                logger.debug(
                    f"Saved {len(feedback_events)} feedback events to workflow run logs"
                )
            except Exception as e:
                logger.error(f"Error saving realtime feedback logs: {e}", exc_info=True)
        else:
            logger.debug("Logs buffer is empty, skipping save")

        logs_update.update(integration_logs)

        if logs_update:
            try:
                await db_client.update_workflow_run(
                    run_id=workflow_run_id,
                    logs=logs_update,
                )
            except Exception as e:
                logger.error(f"Error saving workflow run logs: {e}", exc_info=True)

        # Upload artifacts straight from the in-memory buffers so nothing has
        # to cross a process/host boundary via temp files. Must complete
        # before the completion job is enqueued so QA and webhooks see the
        # artifacts in storage.
        try:
            mixed_audio_wav = None
            user_audio_wav = None
            bot_audio_wav = None

            if not ENABLE_CALL_RECORDING_UPLOAD:
                logger.info(
                    "Call recording upload is disabled "
                    "(ENABLE_CALL_RECORDING_UPLOAD=false), skipping audio upload"
                )
            else:
                if not in_memory_audio_buffers.mixed.is_empty:
                    mixed_audio_wav = await in_memory_audio_buffers.mixed.to_wav_bytes()
                else:
                    logger.debug("Audio buffer is empty, skipping upload")

                if not in_memory_audio_buffers.user.is_empty:
                    user_audio_wav = await in_memory_audio_buffers.user.to_wav_bytes()
                else:
                    logger.debug("User audio buffer is empty, skipping upload")

                if not in_memory_audio_buffers.bot.is_empty:
                    bot_audio_wav = await in_memory_audio_buffers.bot.to_wav_bytes()
                else:
                    logger.debug("Bot audio buffer is empty, skipping upload")

            transcript_text = in_memory_logs_buffer.generate_transcript_text(
                include_end_timestamps=include_transcript_end_timestamps
            )
            if not transcript_text:
                logger.debug("No transcript events in logs buffer, skipping upload")

            await upload_workflow_run_artifacts(
                workflow_run_id,
                mixed_audio_wav=mixed_audio_wav,
                user_audio_wav=user_audio_wav,
                bot_audio_wav=bot_audio_wav,
                transcript_text=transcript_text,
            )
        except Exception as e:
            logger.error(f"Error uploading call artifacts: {e}", exc_info=True)

        # Combined task: runs integrations (including QA), then calculates
        # cost (so QA token usage is captured in usage_info)
        await enqueue_job(
            FunctionNames.PROCESS_WORKFLOW_COMPLETION,
            workflow_run_id,
        )

    # Return the buffer so it can be passed to other handlers
    return in_memory_audio_buffers


def register_audio_data_handler(
    audio_buffer: AudioBufferProcessor,
    workflow_run_id,
    in_memory_buffers: InMemoryRecordingBuffers,
):
    """Register event handler for audio data"""
    logger.info(f"Registering audio data handler for workflow run {workflow_run_id}")

    @audio_buffer.event_handler("on_audio_data")
    async def on_audio_data(buffer, audio, sample_rate, num_channels):
        if not audio:
            return

        try:
            await in_memory_buffers.mixed.append(audio)
        except MemoryError as e:
            logger.error(f"Mixed audio buffer full: {e}")

    @audio_buffer.event_handler("on_track_audio_data")
    async def on_track_audio_data(
        buffer, user_audio, bot_audio, sample_rate, num_channels
    ):
        try:
            if user_audio:
                await in_memory_buffers.user.append(user_audio)
            if bot_audio:
                await in_memory_buffers.bot.append(bot_audio)
        except MemoryError as e:
            logger.error(f"Track audio buffer full: {e}")
