import asyncio
import logging
from dotenv import load_dotenv
from livekit import agents
from livekit.agents import AgentSession, JobContext, WorkerOptions, cli, llm, RoomOutputOptions
from livekit.plugins import openai, silero, tavus

load_dotenv()

# Configure logging
logging.basicConfig(level=<http://logging.INFO|logging.INFO>)
logger = logging.getLogger("official-tavus-agent")

# Tavus Configuration
TAVUS_API_KEY = "**"
PERSONA_ID = "**"
REPLICA_ID = "**"

async def entrypoint(ctx: JobContext):
    """Official Tavus integration following the documentation exactly."""
    
    # Connect to the room first
    await ctx.connect()

    # Create initial chat context
    initial_ctx = llm.ChatContext()
    initial_ctx.messages.append(llm.ChatMessage(
        role="system",
        content="You are a helpful AI assistant. Keep responses brief and conversational.",
    ))

    # Create agent session with STT, LLM, TTS
    session = AgentSession(
        stt=openai.STT(),
        llm=openai.LLM(),
        tts=openai.TTS(),
        vad=silero.VAD.load(),
        chat_ctx=initial_ctx,
    )

    # Create avatar session with official parameters
    avatar = tavus.AvatarSession(
        replica_id=REPLICA_ID,
        persona_id=PERSONA_ID,
        tavus_api_key=TAVUS_API_KEY,
        avatar_participant_name="Tavus-avatar-agent",
    )

    # Start the avatar and wait for it to join
    await avatar.start(session, room=ctx.room)

    # Start agent session with audio disabled (avatar handles audio)
    await session.start(
        room=ctx.room,
        room_output_options=RoomOutputOptions(
            # Disable audio output to the room. The avatar plugin publishes audio separately.
            audio_enabled=False,
        ),
    )

    # Keep the session alive
    await asyncio.sleep(1)


if __name__ == "__main__":
    cli.run_app(WorkerOptions(entrypoint_fnc=entrypoint))