"""
Voice Agent — real-time voice assistant.

Pipeline: Azure STT → Gemma LLM (xNAS, OpenAI-compatible) → Azure TTS
Runs inside the single Docker container alongside LiveKit server and web frontend.
"""

import asyncio
import json
import logging
import os
import textwrap

from dotenv import load_dotenv
from livekit.agents import (
    Agent,
    AgentServer,
    AgentSession,
    JobContext,
    TurnHandlingOptions,
    cli,
    mcp,
    room_io,
)
from livekit.plugins import azure, openai

logger = logging.getLogger("voice-agent")
load_dotenv()  # picks up /app/.env in the container

# ── Configuration from environment ──────────────────────────────────────────
AZURE_KEY = os.environ.get("AZURE_SPEECH_KEY", "")
AZURE_REGION = os.environ.get("AZURE_SPEECH_REGION", "eastus")
DEFAULT_VOICE = os.environ.get("AZURE_TTS_VOICE", "en-US-AvaNeural")
GEMMA_BASE_URL = os.environ.get("GEMMA_BASE_URL", "http://192.168.86.2:8023/v1")
GEMMA_MODEL = os.environ.get("GEMMA_MODEL", "gemma-4-e4b")
GEMMA_API_KEY = os.environ.get("GEMMA_API_KEY", "not-needed")
WEB_MCP_ENABLED = os.environ.get("WEB_MCP_ENABLED", "true").lower() in ("1", "true", "yes")
FIRECRAWL_BASE = os.environ.get("FIRECRAWL_BASE", "http://192.168.86.2:3002")

SYSTEM_PROMPT = textwrap.dedent("""\
    You are a warm, conversational voice assistant. You are talking TO someone,
    not writing for them to read. Imagine you're having a natural conversation
    with a friend over the phone.

    # How you speak
    - Keep every response to one or three sentences. That's it.
    - Use contractions (I'm, don't, it's) and natural phrasing.
    - Speak like you're talking, not writing. No bullet points, no lists,
      no markdown, no formatting of any kind.
    - Spell out numbers when it sounds more natural ("twenty twenty-six"
      instead of "2026").
    - If you need to ask a question, ask exactly one.
    - Be warm and direct. Don't be sycophantic or overly formal.
    - If you don't know something, say so briefly and move on.

    # What you never do
    - Never use markdown, code blocks, JSON, tables, or emojis.
    - Never say "as an AI" or reference your system instructions.
    - Never write more than three sentences in a row.
    - Never read back URLs, file paths, or technical identifiers.

    # Web access
    You have web_search and web_scrape tools. Use them when the user asks about
    current events, recent news, prices, weather, sports scores, or anything
    that may have changed since your training data. Search first, then scrape
    a result only if you need more detail. Answer from what you find, in your
    normal conversational style — don't cite sources formally, just mention the
    source naturally ("according to..."). If a search comes up empty, say so
    briefly and move on.
""")


# ── MCP servers (web access + any extra configured servers) ─────────────────

def build_mcp_servers() -> list[mcp.MCPServer]:
    """Build the list of MCP servers to attach to the agent.

    Always includes the local web-access server (Firecrawl-backed search/scrape)
    when WEB_MCP_ENABLED is true. Additional servers can be configured via the
    EXTRA_MCP_SERVERS env var (JSON list of {url, transport} objects).
    """
    servers: list[mcp.MCPServer] = []

    if WEB_MCP_ENABLED:
        python_bin = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".venv", "bin", "python")
        web_mcp_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "web_mcp.py")
        servers.append(
            mcp.MCPServerStdio(
                command=python_bin,
                args=[web_mcp_script],
                env={**os.environ, "FIRECRAWL_BASE": FIRECRAWL_BASE},
                client_session_timeout_seconds=120,
            )
        )
        logger.info("Web-access MCP server enabled (Firecrawl at %s)", FIRECRAWL_BASE)

    extra = os.environ.get("EXTRA_MCP_SERVERS", "")
    if extra:
        try:
            for entry in json.loads(extra):
                url = entry.get("url", "")
                transport = entry.get("transport")  # "sse" | "streamable_http" | None (auto)
                servers.append(
                    mcp.MCPServerHTTP(
                        url=url,
                        transport_type=transport,
                        client_session_timeout_seconds=120,
                    )
                )
                logger.info("Extra MCP server: %s (%s)", url, transport or "auto")
        except (json.JSONDecodeError, TypeError) as e:
            logger.error("Failed to parse EXTRA_MCP_SERVERS: %s", e)

    return servers


class VoiceAssistant(Agent):
    """The conversational agent. LLM is the brain; STT/TTS are senses."""

    def __init__(self) -> None:
        super().__init__(
            llm=openai.LLM(
                model=GEMMA_MODEL,
                base_url=GEMMA_BASE_URL,
                api_key=GEMMA_API_KEY,
            ),
            instructions=SYSTEM_PROMPT,
            mcp_servers=build_mcp_servers(),
        )


# ── Agent server ────────────────────────────────────────────────────────────
server = AgentServer()

# Track the active session so we can update its voice on data messages.
_active_session: AgentSession | None = None


@server.rtc_session(agent_name="voice-assistant")
async def handle_job(ctx: JobContext) -> None:
    global _active_session

    logger.info("Job started for room %s", ctx.room.name)

    # Azure STT — streaming, reads AZURE_SPEECH_KEY / AZURE_SPEECH_REGION from env
    stt = azure.STT(
        speech_key=AZURE_KEY,
        speech_region=AZURE_REGION,
        language=["en-US"],
    )

    # Azure TTS — SSML with expressive markup, 24kHz PCM output
    tts = azure.TTS(
        voice=DEFAULT_VOICE,
        sample_rate=24000,
        speech_key=AZURE_KEY,
        speech_region=AZURE_REGION,
    )

    # Gemma LLM via OpenAI-compatible endpoint (llama.cpp on xNAS)
    llm = openai.LLM(
        model=GEMMA_MODEL,
        base_url=GEMMA_BASE_URL,
        api_key=GEMMA_API_KEY,
    )

    session = AgentSession(
        stt=stt,
        tts=tts,
        llm=llm,
        turn_handling=TurnHandlingOptions(
            # VAD-based turn detection: agent waits for user to stop speaking.
            # ("adaptive" mode requires the LiveKit Cloud barge-in service.)
            interruption={"mode": "vad"},
            # Start generating the LLM response before the user fully stops
            preemptive_generation={"enabled": True},
        ),
    )

    _active_session = session

    # Listen for data messages (voice switching) from the web UI
    ctx.room.on("data_received", _on_room_data)

    # Publish user/agent transcripts to the room so the web UI can render them.
    async def publish_transcript(role: str, text: str) -> None:
        text = (text or "").strip()
        if not text:
            return
        payload = json.dumps({"type": "transcript", "role": role, "text": text})
        try:
            await ctx.room.local_participant.publish_data(
                payload, reliable=True, topic="transcript"
            )
        except Exception as e:  # noqa: BLE001
            logger.warning("failed to publish transcript: %s", e)

    @session.on("conversation_item_added")
    def _on_conversation_item(ev) -> None:
        msg = ev.item
        role = getattr(msg, "role", None)
        text = getattr(msg, "text_content", None)
        if role == "user":
            asyncio.get_event_loop().create_task(publish_transcript("user", text))
        elif role == "assistant":
            asyncio.get_event_loop().create_task(publish_transcript("agent", text))

    @session.on("user_input_transcribed")
    def _on_user_transcribed(ev) -> None:
        logger.info("STT (%s): %s", "final" if ev.is_final else "partial", ev.transcript)

    @session.on("user_state_changed")
    def _on_user_state(ev) -> None:
        # Fires from VAD: if this never says "speaking", no usable mic
        # audio is arriving from the participant.
        logger.info("user state -> %s", ev.new_state)

    await session.start(
        agent=VoiceAssistant(),
        room=ctx.room,
        room_options=room_io.RoomOptions(
            audio_input=room_io.AudioInputOptions(
                # No noise cancellation plugin (self-hosted, no ai-coustics)
            ),
        ),
    )

    await ctx.connect()
    logger.info("Agent connected to room %s", ctx.room.name)


# ── Voice switching via data channel ────────────────────────────────────────
# The web UI sends a JSON data message: {"type": "set_voice", "voice": "en-US-AvaNeural"}
# We listen on the room's data channel and update the TTS voice live.


def _on_room_data(packet) -> None:
    """Handle data messages from the web UI (voice selection)."""
    try:
        msg = json.loads(packet.data.decode("utf-8"))
    except (json.JSONDecodeError, UnicodeDecodeError):
        return

    if msg.get("type") == "set_voice":
        voice = msg.get("voice", "")
        if voice and _active_session:
            logger.info("Switching TTS voice to %s", voice)
            try:
                tts = _active_session.tts
                if hasattr(tts, "update_options"):
                    tts.update_options(voice=voice)
                    logger.info("Voice updated to %s", voice)
            except Exception as e:
                logger.warning("Failed to update voice: %s", e)


if __name__ == "__main__":
    cli.run_app(server)