"""
agent.py — LiveKit voice agent (livekit-agents 1.7+)

Requires an AgentServer at module level named 'agent'.
"""

import os, logging
from dotenv import load_dotenv

load_dotenv()

from livekit.agents import (
    Agent, AgentSession, AgentServer, JobContext, WorkerOptions, AutoSubscribe,
    EndpointingOptions,
)
from livekit.plugins import openai as lk_openai, silero

logger = logging.getLogger("agent")
logging.basicConfig(level=logging.INFO)


def make_tts():
    backend = os.getenv("TTS_BACKEND", "kokoro")
    if backend == "elevenlabs":
        from livekit.plugins import elevenlabs
        return elevenlabs.TTS(
            api_key=os.environ["ELEVENLABS_API_KEY"],
            voice_id=os.getenv("ELEVENLABS_VOICE_ID", "21m00Tcm4TlvDq8ikWAM"),
        )
    if backend == "piper":
        return _make_piper_tts(os.getenv("PIPER_MODEL", "en_US-libritts-high"))
    base_url = (
        os.getenv("KOKORO_BASE_URL", "http://localhost:8880/v1")
        if backend == "kokoro"
        else os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")
    )
    voice = os.getenv("KOKORO_VOICE", "af_heart") if backend == "kokoro" else "alloy"
    model = "kokoro" if backend == "kokoro" else "tts-1"
    return lk_openai.TTS(model=model, voice=voice, base_url=base_url, api_key="x")


def _make_piper_tts(model_name: str):
    from piper.voice import PiperVoice
    import wave, io, asyncio, pathlib, urllib.request
    from livekit.agents import tts as lk_tts
    from livekit import rtc

    class PiperTTS(lk_tts.TTS):
        def __init__(self):
            super().__init__(streaming=False)
            self._voice = None

        def _load(self):
            if self._voice: return
            d = pathlib.Path.home() / ".local/share/piper"
            d.mkdir(parents=True, exist_ok=True)
            mp = d / f"{model_name}.onnx"
            cp = d / f"{model_name}.onnx.json"
            base = ("https://huggingface.co/rhasspy/piper-voices/resolve/main/"
                    + "/".join(model_name.split("-")[:2]) + f"/medium/{model_name}")
            for p, u in [(mp, base+".onnx"), (cp, base+".onnx.json")]:
                if not p.exists():
                    urllib.request.urlretrieve(u, p)
            self._voice = PiperVoice.load(str(mp), config_path=str(cp))

        def synthesize(self, text):
            return _PiperStream(self, text)

    class _PiperStream(lk_tts.ChunkedStream):
        def __init__(self, tts, text):
            super().__init__(tts=tts, input_text=text)

        async def _run(self):
            loop = asyncio.get_event_loop()
            def _synth():
                self._tts._load()
                buf = io.BytesIO()
                with wave.open(buf, "wb") as wf:
                    self._tts._voice.synthesize(self._input_text, wf)
                return buf.getvalue()
            wav = await loop.run_in_executor(None, _synth)
            buf = io.BytesIO(wav)
            with wave.open(buf, "rb") as wf:
                sr, nc, data = wf.getframerate(), wf.getnchannels(), wf.readframes(wf.getnframes())
            self._event_ch.send_nowait(lk_tts.SynthesizedAudio(
                request_id=self._request_id,
                frame=rtc.AudioFrame(data=data, sample_rate=sr, num_channels=nc,
                                     samples_per_channel=len(data)//(2*nc)),
            ))

    return PiperTTS()


async def entrypoint(ctx: JobContext):
    claims     = ctx.token_claims()
    model      = (claims.attributes or {}).get("model") or os.getenv("DEFAULT_MODEL", "hf.co/mradermacher/Heretic-Dolphin3.0-Qwen2.5-3b-i1-GGUF:Q4_K_M")
    prompt     = (claims.attributes or {}).get("prompt") or os.getenv("SYSTEM_PROMPT", "You are a helpful voice assistant. Keep your answers concise and conversational.")
    ollama_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")

    logger.info(f"Starting agent: model={model} tts={os.getenv('TTS_BACKEND','kokoro')}")

    await ctx.connect(auto_subscribe=AutoSubscribe.AUDIO_ONLY)

    import openai as _openai
    whisper_client = _openai.AsyncOpenAI(
        base_url=os.getenv("WHISPER_BASE_URL", "http://localhost:8778/v1"),
        api_key="x",
        timeout=120.0,
    )

    session = AgentSession(
        vad=silero.VAD.load(),
        stt=lk_openai.STT(
            model="Systran/faster-whisper-base",
            client=whisper_client,
        ),
        llm=lk_openai.LLM(model=model, base_url=ollama_url, api_key="ollama"),
        tts=make_tts(),
        # Give whisper more time — it's slower than cloud STT
        min_endpointing_delay=1.2,
        max_endpointing_delay=6.0,
    )

    await session.start(
        agent=Agent(instructions=prompt),
        room=ctx.room,
    )

    await session.generate_reply(
        instructions="Greet your creator briefly and casually"
    )


# AgentServer instance — required by livekit-agents 1.7+ module discovery
agent = AgentServer(
    ws_url=os.getenv("LIVEKIT_URL", "ws://localhost:7880"),
    api_key=os.getenv("LIVEKIT_API_KEY", "devkey"),
    api_secret=os.getenv("LIVEKIT_API_SECRET", "secret"),
)

@agent.rtc_session
async def run_session(ctx: JobContext):
    await entrypoint(ctx)

if __name__ == "__main__":
    import sys
    from livekit.agents.__main__ import main
    sys.exit(main())