#
# Copyright (c) 2025, Daily
#
# SPDX-License-Identifier: BSD 2-Clause License
#

"""
Pipecat Voice AI Bot - Main entry point.

This module orchestrates the voice AI pipeline with:
- Real-time audio processing (STT → LLM → TTS)
- RTVI protocol for client communication
- Flow management for conversation state
- Transcript and analytics collection

Supports multiple transports:
- Daily.co (WebRTC/RTVI)
- Twilio (Media Streams WebSocket)
"""

import asyncio
import os
from typing import Any, Optional

import sentry_sdk
from dotenv import load_dotenv

# Langfuse tracing
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from pipecat.utils.tracing.setup import setup_tracing
from pipecat.pipeline.runner import PipelineRunner

# Local modules
from bot_schemas import BotData
from bot_config import BotConfig, TransportType
from services import create_llm_service, create_tts_service, create_stt_service, create_transport, create_livekit_transport
from utils.postgres import postgres
from core.pipeline import PipelineBuilder
from core.rtvi import setup_rtvi_processor, setup_client_ready_handler
from core.events import setup_event_handlers
from core.flow import end_conversation_handler, process_flow_config_handlers
from core.transcript import TranscriptCollector, TranscriptObserver
from local.recording import create_recording_processors
from utils.logging import log_info, log_warning, log_error
from fastapi import WebSocket
from metrics import (
    capture_event,
    get_default_provider,
    create_llm_observer,
    create_latency_observer,
)
from pipecat_flows import FlowConfig, FlowManager
from pipecatcloud.agent import DailySessionArguments, WebSocketSessionArguments

load_dotenv(override=True)


def create_transcript_handlers(log_key: str, enabled: bool = True):
    """Create transcript handlers for collecting conversation."""
    if not enabled:
        log_info(log_key, "Transcript collection disabled")
        return None, []

    transcript_collector = TranscriptCollector()
    transcript_observer = TranscriptObserver(transcript_collector)
    log_info(log_key, "Transcript collection enabled")

    return transcript_collector, [transcript_observer]


def setup_tracing_if_enabled(log_key: str) -> None:
    """Configure Langfuse tracing if enabled via environment."""
    if os.getenv("ENABLE_TRACING", "false").lower() != "true":
        return

    try:
        otlp_exporter = OTLPSpanExporter()
        setup_tracing(
            service_name="pipecat-bot",
            exporter=otlp_exporter,
            console_export=os.getenv("OTEL_CONSOLE_EXPORT", "false").lower() == "true",
        )
        log_info(log_key, "Langfuse tracing configured")
    except Exception as e:
        log_error(log_key, f"Failed to configure tracing: {e}")


def setup_sentry(environment: str) -> None:
    """Initialize Sentry error tracking (optional)."""
    sentry_dsn = os.environ.get("SENTRY_DSN")
    if not sentry_dsn:
        return  # Sentry is optional

    sentry_sdk.init(
        dsn=sentry_dsn,
        traces_sample_rate=1.0,
        environment=environment,
    )


async def setup_flow_manager(
    payload: BotData,
    rtvi_client: Any,
    task: Any,
    llm: Any,
    context_aggregator: Any,
    log_key: str,
) -> Optional[FlowManager]:
    """Set up FlowManager for conversation flow."""
    if not payload.flow:
        log_error(log_key, "No flow provided")
        return None

    flow_config = payload.flow.model_dump()
    flow_config_obj = process_flow_config_handlers(
        FlowConfig(**flow_config),
        rtvi_client,
    )

    flow_manager = FlowManager(
        task=task,
        llm=llm,
        context_aggregator=context_aggregator,
        flow_config=flow_config_obj,
    )
    flow_manager.log_key = log_key
    flow_manager.register_action("end_conversation", end_conversation_handler)

    # Handle legacy initial_node attribute
    if not getattr(flow_manager, "initial_node", None):
        legacy_initial = getattr(flow_manager, "_initial_node", None)
        if legacy_initial:
            setattr(flow_manager, "initial_node", legacy_initial)

    return flow_manager


async def enforce_max_session_duration(timeout_seconds: int, task: Any, log_key: str):
    """Cancel task after max session duration."""
    try:
        await asyncio.sleep(timeout_seconds)
        log_warning(log_key, f"Max session duration ({timeout_seconds}s) reached")
        await task.cancel()
    except asyncio.CancelledError:
        pass


async def run_pipeline(config: BotConfig) -> dict:
    """
    Unified pipeline for all transport types.

    Returns dict with session results (transcript, duration, etc.)
    """
    result = {
        "transcript_messages": [],
        "duration": 0,
        "status": "COMPLETED",
    }

    # Parse bot data
    if config.bot_data:
        data = BotData(**config.bot_data) if isinstance(config.bot_data, dict) else config.bot_data
    elif config.agent_id:
        # Load from database
        bot_data_dict = await postgres.get_bot_data(config.agent_id)
        if not bot_data_dict:
            raise ValueError(f"Failed to load bot data for agent: {config.agent_id}")
        data = BotData(**bot_data_dict)
    else:
        raise ValueError("Either bot_data or agent_id must be provided")

    # Extract configuration
    transcript_key = data.metadata.transcript_key if hasattr(data.metadata, 'transcript_key') else None
    environment = data.metadata.environment or "production"
    user_id = data.metadata.user_id or config.caller_number or f"user-{config.session_id}"
    max_duration = data.metadata.max_duration or 3600

    # Create log key
    if config.transport_type == TransportType.TWILIO:
        log_key = f"twilio_{config.call_sid or config.session_id or 'unknown'}"
    else:
        log_key = transcript_key.replace("/", "_") if transcript_key else config.session_id or "unknown-session"

    log_info(log_key, f"Pipeline started [transport={config.transport_type.value}, env={environment}]")

    # Setup infrastructure
    setup_tracing_if_enabled(log_key)
    setup_sentry(environment)

    # Create or use provided transport
    if config.transport_type == TransportType.TWILIO:
        if not config.transport:
            raise ValueError("Twilio transport must be provided via config.transport")
        transport = config.transport
        rtvi = None  # No RTVI for Twilio
    elif config.transport_type == TransportType.LIVEKIT:
        # LiveKit Cloud transport
        if not config.livekit_url or not config.livekit_token:
            raise ValueError("livekit_url and livekit_token required for LiveKit transport")
        transport = create_livekit_transport(
            room_url=config.livekit_url,
            token=config.livekit_token,
            room_name=config.livekit_room_name or "",
            participant_identity=config.participant_identity,
        )
        rtvi = None  # RTVI not supported with LiveKit yet
        log_info(log_key, f"LiveKit transport created [identity={config.participant_identity}]")
    else:
        # Daily.co transport
        if not config.room_url or not config.token:
            raise ValueError("room_url and token required for Daily transport")
        transport = create_transport(config.room_url, config.token, config.bot_name)
        rtvi = None
        if config.rtvi_enabled:
            rtvi = setup_rtvi_processor(
                pipeline_getter=lambda: pipeline,
                log_func=lambda msg: log_info(log_key, msg),
            )

    # Create services with provider selection
    llm = create_llm_service(
        provider=data.config.llm_provider.value,
        model=data.config.model,
    )
    tts = create_tts_service(
        provider=data.config.tts_provider.value,
        voice_id=data.config.voice_id,
    )
    stt = create_stt_service(
        provider=data.config.stt_provider.value,
    )

    log_info(
        log_key,
        f"Services initialized [llm={data.config.llm_provider.value}/{data.config.model}, "
        f"tts={data.config.tts_provider.value}, stt={data.config.stt_provider.value}]"
    )

    # Setup recording (local only for now)
    audio_processor = None
    if config.local_recording:
        audio_processor = create_recording_processors(config.session_id, config.recording_dir)
        log_info(log_key, f"Local recording mode [dir={config.recording_dir}]")

    # Setup transcript collection - always enabled for Twilio to capture messages
    enable_transcript = config.transport_type == TransportType.TWILIO or transcript_key is not None
    transcript_collector, transcript_observers = create_transcript_handlers(log_key, enable_transcript)

    # Build pipeline
    pipeline = None
    builder = (
        PipelineBuilder()
        .with_transport(transport)
        .with_rtvi(rtvi)
        .with_stt(stt)
        .with_llm(llm)
        .with_tts(tts)
        .with_recording(audio_processor)
        .with_transcript()
    )

    pipeline = builder.build()

    # Build observers
    llm_observer = None
    latency_observer = None
    if config.transport_type == TransportType.DAILY:
        llm_observer = create_llm_observer(data, llm.model_name)
        latency_observer = create_latency_observer(data, rtvi if config.rtvi_enabled else None)

    observers = builder.build_observers(transcript_observers, llm_observer, latency_observer)

    # Create task
    task = builder.create_task(pipeline, observers, log_key, user_id, environment)

    # Setup flow manager
    flow_manager = await setup_flow_manager(
        data, rtvi, task, llm, builder.context_aggregator, log_key
    )
    if not flow_manager:
        log_error(log_key, "Failed to initialize FlowManager")
        result["status"] = "FAILED"
        return result

    log_info(log_key, "FlowManager initialized")

    # Setup max session duration
    log_info(log_key, f"Max duration: {max_duration}s")
    asyncio.create_task(enforce_max_session_duration(max_duration, task, log_key))

    # Setup event handlers based on transport type
    if config.transport_type == TransportType.TWILIO:
        # For Twilio, setup simplified event handlers
        setup_twilio_event_handlers(
            transport=transport,
            task=task,
            flow_manager=flow_manager,
            transcript_processor=builder.transcript_processor,
            transcript_collector=transcript_collector,
            log_key=log_key,
        )
    elif config.transport_type == TransportType.LIVEKIT:
        # LiveKit event handlers
        setup_livekit_event_handlers(
            transport=transport,
            task=task,
            flow_manager=flow_manager,
            transcript_processor=builder.transcript_processor,
            transcript_collector=transcript_collector,
            log_key=log_key,
        )
    else:
        # Daily.co event handlers
        if config.rtvi_enabled and rtvi:
            setup_client_ready_handler(rtvi, lambda msg: log_info(log_key, msg))

        setup_event_handlers(
            transport=transport,
            task=task,
            flow_manager=flow_manager,
            rtvi=rtvi,
            transcript=builder.transcript_processor,
            transcript_collector=transcript_collector,
            audio_processor=audio_processor,
            config=config,
            log_key=log_key,
            capture_event_func=capture_event,
            bot_data=data,
        )

    # Run pipeline
    runner = PipelineRunner()

    try:
        await runner.run(task)
        log_info(log_key, "Pipeline completed successfully")
        result["status"] = "COMPLETED"
    except Exception as e:
        log_error(log_key, f"Pipeline failed: {e}")
        result["status"] = "FAILED"
        if config.transport_type == TransportType.DAILY:
            capture_event(data, "pipecat_error_encountered", {"message": str(e)})
        raise
    finally:
        # Collect transcript messages
        if transcript_collector:
            result["transcript_messages"] = transcript_collector.get_messages()

        if config.transport_type == TransportType.DAILY:
            get_default_provider().close()

        log_info(log_key, "Pipeline finished")

    return result


def setup_twilio_event_handlers(
    transport: Any,
    task: Any,
    flow_manager: Any,
    transcript_processor: Any,
    transcript_collector: Optional[TranscriptCollector],
    log_key: str,
) -> None:
    """
    Setup event handlers for Twilio transport.

    Twilio doesn't have participant join/leave events like Daily.co,
    so we handle the flow differently.
    """
    from datetime import datetime, timezone

    # Initialize flow immediately when transport is ready
    @transport.event_handler("on_client_connected")
    async def handle_client_connected(transport_instance, websocket):
        log_info(log_key, "Twilio client connected")

        # Start transcript session
        if transcript_collector:
            recording_start_time = datetime.now(timezone.utc).isoformat()
            transcript_collector.start_session(recording_start_time)

        # Initialize flow
        await flow_manager.initialize()

        # Set initial node
        if flow_manager.initial_node:
            from core.flow import send_state_sync
            await send_state_sync(flow_manager, flow_manager.initial_node)

    @transport.event_handler("on_client_disconnected")
    async def handle_client_disconnected(transport_instance, websocket):
        log_info(log_key, "Twilio client disconnected")

        # Finalize transcript
        if transcript_collector:
            transcript_collector.finalize()

        # Cancel task
        await task.cancel()

    # Handle transcript updates
    if transcript_processor:
        @transcript_processor.event_handler("on_transcript_update")
        async def handle_transcript_update(processor, frame):
            if transcript_collector:
                for message in frame.messages:
                    transcript_collector.process_utterance(message.role, message.content)


def setup_livekit_event_handlers(
    transport: Any,
    task: Any,
    flow_manager: Any,
    transcript_processor: Any,
    transcript_collector: Optional[TranscriptCollector],
    log_key: str,
) -> None:
    """
    Setup event handlers for LiveKit transport.

    LiveKit has similar event model to Daily.co but with different event names.
    """
    from datetime import datetime, timezone

    @transport.event_handler("on_participant_connected")
    async def handle_participant_connected(transport_instance, participant):
        log_info(log_key, f"LiveKit participant connected: {participant.identity if hasattr(participant, 'identity') else participant}")

        # Start transcript session
        if transcript_collector:
            recording_start_time = datetime.now(timezone.utc).isoformat()
            transcript_collector.start_session(recording_start_time)

        # Initialize flow
        await flow_manager.initialize()

        # Set initial node
        if flow_manager.initial_node:
            from core.flow import send_state_sync
            await send_state_sync(flow_manager, flow_manager.initial_node)

    @transport.event_handler("on_participant_disconnected")
    async def handle_participant_disconnected(transport_instance, participant):
        log_info(log_key, f"LiveKit participant disconnected: {participant.identity if hasattr(participant, 'identity') else participant}")

        # Finalize transcript
        if transcript_collector:
            transcript_collector.finalize()

        # Cancel task
        await task.cancel()

    # Handle transcript updates if processor is available
    if transcript_processor:
        @transcript_processor.event_handler("on_transcript_update")
        async def handle_transcript_update(processor, frame):
            if transcript_collector:
                for message in frame.messages:
                    transcript_collector.process_utterance(message.role, message.content)


async def run_livekit_pipeline(config: BotConfig) -> dict:
    """
    Entry point for LiveKit bot sessions.

    This is called by the LiveKit runner to start a bot with LiveKit transport.
    """
    log_info(config.session_id or "livekit", f"Starting LiveKit pipeline [url={config.livekit_url}]")
    return await run_pipeline(config)


# -----------------------------------------------------------------------------
# Entry Points
# -----------------------------------------------------------------------------

async def bot(args):
    """
    Main bot entry point for Pipecat Cloud.

    Handles both transport types:
    - DailySessionArguments: WebRTC sessions via Daily.co
    - WebSocketSessionArguments: Twilio Media Streams via WebSocket
    """
    # Check argument type to determine transport
    if isinstance(args, WebSocketSessionArguments):
        # Twilio WebSocket transport
        log_info("", f"Twilio WebSocket session started [session_id={args.session_id}]")
        try:
            from services.twilio_transport import handle_twilio_call
            # Pass session_id to enable Pipecat Cloud logs retrieval
            await handle_twilio_call(args.websocket, run_twilio_pipeline_wrapper, args.session_id)
            log_info("", "Twilio WebSocket session completed")
        except Exception as e:
            log_error("", f"Twilio WebSocket session error: {str(e)}")
            raise
    elif isinstance(args, DailySessionArguments):
        # Daily.co WebRTC transport
        log_info("", f"Bot initialized [room={args.room_url}]")

        if not args.body:
            raise ValueError("No data provided")

        try:
            session_id = args.session_id if hasattr(args, "session_id") else None

            config = BotConfig(
                transport_type=TransportType.DAILY,
                session_id=session_id,
                room_url=args.room_url,
                token=args.token,
                bot_data=args.body,
            )

            await run_pipeline(config)
            log_info(None, "Bot completed")
        except Exception as e:
            log_error(None, f"Bot error: {str(e)}")
            raise
    else:
        raise TypeError(f"Unsupported session arguments type: {type(args)}")


async def run_twilio_pipeline_wrapper(
    transport: Any,
    agent_config: dict,
    context: Any,  # TwilioCallContext
) -> None:
    """
    Wrapper to run unified pipeline for Twilio calls.

    This bridges the old TwilioCallContext API with the new unified pipeline.
    """
    log_key = f"twilio_{context.call_sid}"

    # Load botData from agent config
    bot_data_json = agent_config.get("botData")
    if bot_data_json:
        log_info(log_key, "Using botData from database")
        bot_data = bot_data_json
    else:
        # Fallback to legacy generation
        log_warning(log_key, "No botData in database, using legacy generation")
        bot_data = postgres._generate_bot_data_legacy(agent_config)

    # Create unified config
    config = BotConfig(
        transport_type=TransportType.TWILIO,
        session_id=context.pipecat_session_id,
        call_sid=context.call_sid,
        stream_sid=context.stream_sid,
        caller_number=context.caller_number,
        agent_id=context.agent_id,
        bot_data=bot_data,
        transport=transport,
        rtvi_enabled=False,  # No RTVI for Twilio
    )

    # Run unified pipeline
    result = await run_pipeline(config)

    # Update context with results for database save
    if result.get("transcript_messages"):
        context.transcript_messages = result["transcript_messages"]


async def twilio_websocket_handler(websocket: WebSocket) -> None:
    """
    FastAPI WebSocket endpoint handler for Twilio Media Streams.

    This is for local development. Mount this in your FastAPI app:
        @app.websocket("/twilio")
        async def twilio_ws(websocket: WebSocket):
            await twilio_websocket_handler(websocket)
    """
    from services.twilio_transport import handle_twilio_call

    await handle_twilio_call(websocket, run_twilio_pipeline_wrapper)


# Legacy alias for backwards compatibility
main = run_pipeline