"""
Custom LLM plugin that proxies to IT0 agent-service.

Instead of calling Claude directly, this plugin:
1. POSTs to agent-service /api/v1/agent/tasks (engineType configurable: claude_agent_sdk or claude_api)
2. Subscribes to the agent-service WebSocket /ws/agent for streaming text events
3. Emits ChatChunk objects into the LiveKit pipeline

Voice Session Mode (long-lived agent)
--------------------------------------
When engine_type == "claude_agent_sdk", the plugin uses a long-lived voice session:

  Voice call starts → start_voice_session()
    → POST /api/v1/agent/sessions/voice/start
    → agent-service spawns a persistent run loop for this session

  User speaks (each turn) → AgentServiceLLMStream._run()
    → POST /api/v1/agent/sessions/{sessionId}/voice/inject
    → run loop receives the message and executes one SDK turn
    → streams events back via WebSocket → ChatChunks → TTS

  User hangs up → terminate_voice_session()
    → DELETE /api/v1/agent/sessions/{sessionId}/voice
    → agent-service kills the run loop and marks session completed

This preserves all agent-service capabilities: Tool Use, conversation history,
tenant isolation, and session management.
"""

import asyncio
import json
import logging
import uuid
from typing import Any

import httpx
import websockets

from livekit.agents import llm
from livekit.agents.types import (
    DEFAULT_API_CONNECT_OPTIONS,
    NOT_GIVEN,
    APIConnectOptions,
    NotGivenOr,
)

logger = logging.getLogger(__name__)


class AgentServiceLLM(llm.LLM):
    """LLM that proxies to IT0 agent-service for Claude + Tool Use."""

    def __init__(
        self,
        *,
        agent_service_url: str = "http://agent-service:3002",
        auth_header: str = "",
        engine_type: str = "claude_agent_sdk",
    ):
        super().__init__()
        self._agent_service_url = agent_service_url
        self._auth_header = auth_header
        self._engine_type = engine_type
        self._agent_session_id: str | None = None

        # Voice session mode: long-lived agent process tied to the call duration.
        # True once start_voice_session() completes successfully.
        self._voice_session_started: bool = False

    @property
    def model(self) -> str:
        return "agent-service-proxy"

    @property
    def provider(self) -> str:
        return "it0-agent"

    def chat(
        self,
        *,
        chat_ctx: llm.ChatContext,
        tools: list[llm.Tool] | None = None,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
        parallel_tool_calls: NotGivenOr[bool] = NOT_GIVEN,
        tool_choice: NotGivenOr[llm.ToolChoice] = NOT_GIVEN,
        extra_kwargs: NotGivenOr[dict[str, Any]] = NOT_GIVEN,
    ) -> "AgentServiceLLMStream":
        return AgentServiceLLMStream(
            llm_instance=self,
            chat_ctx=chat_ctx,
            tools=tools or [],
            conn_options=conn_options,
        )

    # -------------------------------------------------------------------------
    # Voice session lifecycle
    # -------------------------------------------------------------------------

    async def start_voice_session(self, session_id: str | None = None) -> str:
        """
        Start a long-lived voice agent session in agent-service.

        Called ONCE when the LiveKit room is connected.  The agent-service
        spawns a background run loop that stays alive for the entire call,
        accepting injected messages via voice/inject.

        Returns the session ID to use for all subsequent inject calls.
        """
        headers: dict[str, str] = {"Content-Type": "application/json"}
        if self._auth_header:
            headers["Authorization"] = self._auth_header

        body: dict[str, Any] = {}
        if session_id:
            body["sessionId"] = session_id

        try:
            async with httpx.AsyncClient(
                timeout=httpx.Timeout(connect=10, read=15, write=10, pool=10),
            ) as client:
                resp = await client.post(
                    f"{self._agent_service_url}/api/v1/agent/sessions/voice/start",
                    json=body,
                    headers=headers,
                )

            if resp.status_code not in (200, 201):
                logger.error(
                    "start_voice_session failed: %d %s",
                    resp.status_code, resp.text[:200],
                )
                return ""

            data = resp.json()
            new_session_id: str = data.get("sessionId", "")
            self._agent_session_id = new_session_id
            self._voice_session_started = True
            logger.info("Voice session started: %s", new_session_id)
            return new_session_id

        except Exception as exc:
            logger.error("start_voice_session error: %s: %s", type(exc).__name__, exc)
            return ""

    async def inject_text_message(
        self,
        *,
        text: str = "",
        attachments: list[dict] | None = None,
    ) -> str:
        """
        Inject a text message (sent via data channel, not speech) into the agent.

        In voice session mode, uses voice/inject and collects the streamed response.
        In text mode, uses POST /tasks directly.

        Returns the complete response text for TTS playback via session.say().
        """
        if not text and not attachments:
            return ""

        try:
            if self._voice_session_started and self._agent_session_id:
                return await self._do_inject_voice(text)
            else:
                return await self._do_inject(text, attachments)
        except Exception as exc:
            logger.error("inject_text_message error: %s: %s", type(exc).__name__, exc)
            return ""

    async def _do_inject_voice(self, text: str) -> str:
        """Inject via voice/inject endpoint and collect response text from WS."""
        import time

        agent_url = self._agent_service_url
        session_id = self._agent_session_id
        headers: dict[str, str] = {"Content-Type": "application/json"}
        if self._auth_header:
            headers["Authorization"] = self._auth_header

        ws_url = agent_url.replace("http://", "ws://").replace("https://", "wss://")
        ws_url = f"{ws_url}/ws/agent"

        timeout_secs = 120
        collected_text = ""

        async with websockets.connect(
            ws_url,
            open_timeout=10,
            close_timeout=5,
            ping_interval=20,
            ping_timeout=10,
        ) as ws:
            await ws.send(json.dumps({
                "event": "subscribe_session",
                "data": {"sessionId": session_id},
            }))

            async with httpx.AsyncClient(
                timeout=httpx.Timeout(connect=10, read=15, write=10, pool=10),
            ) as client:
                resp = await client.post(
                    f"{agent_url}/api/v1/agent/sessions/{session_id}/voice/inject",
                    json={"message": text},
                    headers=headers,
                )
            if resp.status_code not in (200, 201):
                logger.error("_do_inject_voice failed: %d", resp.status_code)
                return ""

            deadline = time.time() + timeout_secs
            while time.time() < deadline:
                remaining = deadline - time.time()
                try:
                    raw = await asyncio.wait_for(ws.recv(), timeout=min(30.0, remaining))
                except asyncio.TimeoutError:
                    continue
                except websockets.exceptions.ConnectionClosed:
                    break

                try:
                    msg = json.loads(raw)
                except (json.JSONDecodeError, TypeError):
                    continue

                if msg.get("event") != "stream_event":
                    continue

                evt_data = msg.get("data", {})
                evt_type = evt_data.get("type", "")

                if evt_type == "text":
                    collected_text += evt_data.get("content", "")
                elif evt_type in ("completed", "error"):
                    break

        return collected_text

    async def _do_inject(
        self,
        text: str,
        attachments: list[dict] | None,
    ) -> str:
        """Text-mode inject: POST /tasks and collect streamed response text."""
        import time

        agent_url = self._agent_service_url
        headers: dict[str, str] = {"Content-Type": "application/json"}
        if self._auth_header:
            headers["Authorization"] = self._auth_header

        ws_url = agent_url.replace("http://", "ws://").replace("https://", "wss://")
        ws_url = f"{ws_url}/ws/agent"

        timeout_secs = 120
        engine_type = self._engine_type

        body: dict[str, Any] = {
            "prompt": text if text else "(see attachments)",
            "engineType": engine_type,
            "voiceMode": False,
        }
        if self._agent_session_id:
            body["sessionId"] = self._agent_session_id
        if attachments:
            body["attachments"] = attachments

        collected_text = ""

        async with websockets.connect(
            ws_url,
            open_timeout=10,
            close_timeout=5,
            ping_interval=20,
            ping_timeout=10,
        ) as ws:
            if self._agent_session_id:
                await ws.send(json.dumps({
                    "event": "subscribe_session",
                    "data": {"sessionId": self._agent_session_id},
                }))

            async with httpx.AsyncClient(
                timeout=httpx.Timeout(connect=10, read=30, write=10, pool=10),
            ) as client:
                resp = await client.post(
                    f"{agent_url}/api/v1/agent/tasks",
                    json=body,
                    headers=headers,
                )
            if resp.status_code not in (200, 201):
                logger.error("_do_inject task failed: %d", resp.status_code)
                return ""

            data = resp.json()
            session_id = data.get("sessionId", "")
            task_id = data.get("taskId", "")
            self._agent_session_id = session_id

            await ws.send(json.dumps({
                "event": "subscribe_session",
                "data": {"sessionId": session_id, "taskId": task_id},
            }))

            deadline = time.time() + timeout_secs
            while time.time() < deadline:
                remaining = deadline - time.time()
                try:
                    raw = await asyncio.wait_for(ws.recv(), timeout=min(30.0, remaining))
                except asyncio.TimeoutError:
                    continue
                except websockets.exceptions.ConnectionClosed:
                    break

                try:
                    msg = json.loads(raw)
                except (json.JSONDecodeError, TypeError):
                    continue

                if msg.get("event") != "stream_event":
                    continue

                evt_data = msg.get("data", {})
                evt_type = evt_data.get("type", "")

                if evt_type == "text":
                    collected_text += evt_data.get("content", "")
                elif evt_type == "completed":
                    return collected_text
                elif evt_type == "error":
                    break

        return collected_text

    async def terminate_voice_session(self) -> None:
        """
        Terminate the long-lived voice agent session.

        Called when the user hangs up (LiveKit room disconnects).
        Signals agent-service to abort any running SDK task and clean up.
        """
        if not self._voice_session_started or not self._agent_session_id:
            return

        session_id = self._agent_session_id
        headers: dict[str, str] = {"Content-Type": "application/json"}
        if self._auth_header:
            headers["Authorization"] = self._auth_header

        try:
            async with httpx.AsyncClient(
                timeout=httpx.Timeout(connect=5, read=10, write=5, pool=5),
            ) as client:
                resp = await client.delete(
                    f"{self._agent_service_url}/api/v1/agent/sessions/{session_id}/voice",
                    headers=headers,
                )
            logger.info(
                "Voice session terminated: %s (status=%d)",
                session_id, resp.status_code,
            )
        except Exception as exc:
            logger.error("terminate_voice_session error: %s: %s", type(exc).__name__, exc)
        finally:
            self._voice_session_started = False
            self._agent_session_id = None


class AgentServiceLLMStream(llm.LLMStream):
    """Streams text from agent-service via WebSocket."""

    def __init__(
        self,
        *,
        llm_instance: AgentServiceLLM,
        chat_ctx: llm.ChatContext,
        tools: list[llm.Tool],
        conn_options: APIConnectOptions,
    ):
        super().__init__(
            llm_instance,
            chat_ctx=chat_ctx,
            tools=tools,
            conn_options=conn_options,
        )
        self._llm_instance = llm_instance

    # Retry configuration (used in fallback / non-voice mode only)
    _MAX_RETRIES = 2
    _RETRY_DELAYS = [1.0, 3.0]  # seconds between retries

    async def _run(self) -> None:
        # Extract the latest user message from ChatContext
        user_text = ""
        for item in reversed(self._chat_ctx.items):
            if getattr(item, "type", None) != "message":
                continue
            if item.role == "user":
                user_text = item.text_content
                break

        if not user_text:
            # on_enter/generate_reply may call LLM without a user message;
            # look for the developer/system instruction to use as prompt
            for item in self._chat_ctx.items:
                if getattr(item, "type", None) != "message":
                    continue
                if item.role in ("developer", "system"):
                    user_text = item.text_content
                    break

        if not user_text:
            logger.warning("No user message found in chat context")
            return

        request_id = f"agent-{uuid.uuid4().hex[:12]}"

        # -----------------------------------------------------------------
        # Voice session mode: inject into the long-lived run loop
        # -----------------------------------------------------------------
        if self._llm_instance._voice_session_started:
            try:
                await self._do_stream_voice(user_text, request_id)
            except Exception as exc:
                logger.error(
                    "Voice inject stream error: %s: %s",
                    type(exc).__name__, exc, exc_info=True,
                )
                self._event_ch.send_nowait(
                    llm.ChatChunk(
                        id=request_id,
                        delta=llm.ChoiceDelta(
                            role="assistant",
                            content="抱歉，语音服务暂时不可用。",
                        ),
                    )
                )
            return

        # -----------------------------------------------------------------
        # Fallback / text mode: per-turn POST /tasks with retry
        # -----------------------------------------------------------------
        last_error: Exception | None = None

        for attempt in range(self._MAX_RETRIES + 1):
            try:
                if attempt > 0:
                    delay = self._RETRY_DELAYS[min(attempt - 1, len(self._RETRY_DELAYS) - 1)]
                    logger.info("Retry %d/%d after %.1fs", attempt, self._MAX_RETRIES, delay)
                    await asyncio.sleep(delay)

                await self._do_stream(user_text, request_id)
                return  # success

            except (httpx.ConnectError, httpx.ConnectTimeout, OSError) as exc:
                last_error = exc
                logger.warning(
                    "Agent stream attempt %d failed (network): %s: %s",
                    attempt + 1, type(exc).__name__, exc,
                )
            except websockets.exceptions.InvalidStatusCode as exc:
                last_error = exc
                logger.warning(
                    "Agent WS connect attempt %d failed: status %s",
                    attempt + 1, getattr(exc, "status_code", "?"),
                )
            except Exception as exc:
                logger.error("Agent stream error: %s: %s", type(exc).__name__, exc)
                self._event_ch.send_nowait(
                    llm.ChatChunk(
                        id=request_id,
                        delta=llm.ChoiceDelta(
                            role="assistant",
                            content="抱歉，Agent服务暂时不可用。",
                        ),
                    )
                )
                return

        logger.error(
            "Agent stream failed after %d attempts: %s",
            self._MAX_RETRIES + 1, last_error,
        )
        self._event_ch.send_nowait(
            llm.ChatChunk(
                id=request_id,
                delta=llm.ChoiceDelta(
                    role="assistant",
                    content="抱歉，Agent服务暂时不可用，请稍后再试。",
                ),
            )
        )

    # -------------------------------------------------------------------------
    # Voice session stream: inject + subscribe
    # -------------------------------------------------------------------------

    async def _do_stream_voice(self, user_text: str, request_id: str) -> None:
        """
        Voice session mode: inject the user utterance into the long-lived
        agent run loop via POST voice/inject, then subscribe to the WebSocket
        to receive the streaming response as ChatChunks.

        Unlike the per-task approach (_do_stream), no new task is created here —
        the VoiceSessionManager's run loop handles task creation internally.
        """
        import time

        agent_url = self._llm_instance._agent_service_url
        auth_header = self._llm_instance._auth_header
        session_id = self._llm_instance._agent_session_id

        if not session_id:
            logger.error("_do_stream_voice: no session_id — falling back to per-task mode")
            await self._do_stream(user_text, request_id)
            return

        headers: dict[str, str] = {"Content-Type": "application/json"}
        if auth_header:
            headers["Authorization"] = auth_header

        ws_url = agent_url.replace("http://", "ws://").replace("https://", "wss://")
        ws_url = f"{ws_url}/ws/agent"

        timeout_secs = 120

        logger.info(
            "Voice inject: session=%s text=%s",
            session_id, user_text[:80],
        )

        async with websockets.connect(
            ws_url,
            open_timeout=10,
            close_timeout=5,
            ping_interval=20,
            ping_timeout=10,
        ) as ws:
            # 1. Pre-subscribe to receive events as soon as the task starts
            await ws.send(json.dumps({
                "event": "subscribe_session",
                "data": {"sessionId": session_id},
            }))

            # 2. Inject the message (enqueues to the run loop; async return)
            async with httpx.AsyncClient(
                timeout=httpx.Timeout(connect=10, read=15, write=10, pool=10),
            ) as client:
                resp = await client.post(
                    f"{agent_url}/api/v1/agent/sessions/{session_id}/voice/inject",
                    json={"message": user_text},
                    headers=headers,
                )

            if resp.status_code not in (200, 201):
                logger.error(
                    "Voice inject failed: %d %s",
                    resp.status_code, resp.text[:200],
                )
                self._event_ch.send_nowait(
                    llm.ChatChunk(
                        id=request_id,
                        delta=llm.ChoiceDelta(
                            role="assistant",
                            content="抱歉，语音指令注入失败。",
                        ),
                    )
                )
                return

            # 3. Emit initial role delta (LiveKit convention)
            self._event_ch.send_nowait(
                llm.ChatChunk(
                    id=request_id,
                    delta=llm.ChoiceDelta(role="assistant"),
                )
            )

            # 4. Stream events → ChatChunks until completed/error
            deadline = time.time() + timeout_secs

            while time.time() < deadline:
                remaining = deadline - time.time()
                try:
                    raw = await asyncio.wait_for(
                        ws.recv(), timeout=min(30.0, remaining)
                    )
                except asyncio.TimeoutError:
                    if time.time() >= deadline:
                        logger.warning("Voice inject stream timeout after %ds", timeout_secs)
                    continue
                except websockets.exceptions.ConnectionClosed:
                    logger.warning("Voice inject WS connection closed")
                    break

                try:
                    msg = json.loads(raw)
                except (json.JSONDecodeError, TypeError):
                    continue

                if msg.get("event") != "stream_event":
                    continue

                evt_data = msg.get("data", {})
                evt_type = evt_data.get("type", "")

                if evt_type == "text":
                    content = evt_data.get("content", "")
                    if content:
                        self._event_ch.send_nowait(
                            llm.ChatChunk(
                                id=request_id,
                                delta=llm.ChoiceDelta(content=content),
                            )
                        )

                elif evt_type == "completed":
                    logger.info("Voice inject stream completed")
                    return

                elif evt_type == "error":
                    err_msg = evt_data.get("message", "Unknown error")
                    logger.error("Voice inject agent error: %s", err_msg)
                    self._event_ch.send_nowait(
                        llm.ChatChunk(
                            id=request_id,
                            delta=llm.ChoiceDelta(content=f"Agent 错误: {err_msg}"),
                        )
                    )
                    return

    # -------------------------------------------------------------------------
    # Legacy per-task stream (text chat / fallback)
    # -------------------------------------------------------------------------

    async def _do_stream(self, user_text: str, request_id: str) -> None:
        """Execute a single WS+HTTP streaming attempt (text-chat / non-voice mode)."""
        import time

        agent_url = self._llm_instance._agent_service_url
        auth_header = self._llm_instance._auth_header

        headers: dict[str, str] = {"Content-Type": "application/json"}
        if auth_header:
            headers["Authorization"] = auth_header

        ws_url = agent_url.replace("http://", "ws://").replace("https://", "wss://")
        ws_url = f"{ws_url}/ws/agent"

        timeout_secs = 120

        logger.info("Connecting to agent-service WS: %s", ws_url)
        async with websockets.connect(
            ws_url,
            open_timeout=10,
            close_timeout=5,
            ping_interval=20,
            ping_timeout=10,
        ) as ws:
            # 1. Pre-subscribe with existing session ID (for event buffering)
            if self._llm_instance._agent_session_id:
                await ws.send(json.dumps({
                    "event": "subscribe_session",
                    "data": {"sessionId": self._llm_instance._agent_session_id},
                }))

            # 2. Create agent task (with timeout)
            engine_type = self._llm_instance._engine_type
            voice_mode = engine_type == "claude_agent_sdk"

            body: dict[str, Any] = {
                "prompt": user_text,
                "engineType": engine_type,
                "voiceMode": voice_mode,
            }

            if voice_mode:
                body["systemPrompt"] = (
                    "你正在通过语音与用户实时对话。请严格遵守以下规则：\n"
                    "1. 只输出用户关注的最终答案，不要输出工具调用过程、中间步骤或技术细节\n"
                    "2. 用简洁自然的口语中文回答，像面对面对话一样\n"
                    "3. 回复要简短精炼，适合语音播报，通常1-3句话即可\n"
                    "4. 不要使用markdown格式、代码块、列表符号等文本格式"
                )

            if self._llm_instance._agent_session_id:
                body["sessionId"] = self._llm_instance._agent_session_id

            logger.info(
                "POST /tasks engine=%s voiceMode=%s user_text=%s",
                engine_type, voice_mode, user_text[:80],
            )
            async with httpx.AsyncClient(
                timeout=httpx.Timeout(connect=10, read=30, write=10, pool=10),
            ) as client:
                resp = await client.post(
                    f"{agent_url}/api/v1/agent/tasks",
                    json=body,
                    headers=headers,
                )

                if resp.status_code not in (200, 201):
                    logger.error(
                        "Task creation failed: %d %s",
                        resp.status_code, resp.text[:200],
                    )
                    self._event_ch.send_nowait(
                        llm.ChatChunk(
                            id=request_id,
                            delta=llm.ChoiceDelta(
                                role="assistant",
                                content="抱歉，Agent服务暂时不可用。",
                            ),
                        )
                    )
                    return

                data = resp.json()
                session_id = data.get("sessionId", "")
                task_id = data.get("taskId", "")
                self._llm_instance._agent_session_id = session_id
                logger.info(
                    "Task created: session=%s, task=%s", session_id, task_id
                )

            # 3. Subscribe with actual IDs
            await ws.send(json.dumps({
                "event": "subscribe_session",
                "data": {"sessionId": session_id, "taskId": task_id},
            }))

            # 4. Send initial role delta
            self._event_ch.send_nowait(
                llm.ChatChunk(
                    id=request_id,
                    delta=llm.ChoiceDelta(role="assistant"),
                )
            )

            # 5. Stream events → ChatChunk
            deadline = time.time() + timeout_secs

            while time.time() < deadline:
                remaining = deadline - time.time()
                try:
                    raw = await asyncio.wait_for(
                        ws.recv(), timeout=min(30.0, remaining)
                    )
                except asyncio.TimeoutError:
                    if time.time() >= deadline:
                        logger.warning("Agent stream timeout after %ds", timeout_secs)
                    continue
                except websockets.exceptions.ConnectionClosed:
                    logger.warning("Agent WS connection closed during streaming")
                    break

                try:
                    msg = json.loads(raw)
                except (json.JSONDecodeError, TypeError):
                    continue

                event_type = msg.get("event", "")

                if event_type == "stream_event":
                    evt_data = msg.get("data", {})
                    evt_type = evt_data.get("type", "")

                    if evt_type == "text":
                        content = evt_data.get("content", "")
                        if content:
                            self._event_ch.send_nowait(
                                llm.ChatChunk(
                                    id=request_id,
                                    delta=llm.ChoiceDelta(content=content),
                                )
                            )

                    elif evt_type == "completed":
                        logger.info("Agent stream completed")
                        return

                    elif evt_type == "error":
                        err_msg = evt_data.get("message", "Unknown error")
                        logger.error("Agent error: %s", err_msg)
                        if "aborted" in err_msg.lower() or "exited" in err_msg.lower():
                            logger.info("Clearing agent session after abort/exit")
                            self._llm_instance._agent_session_id = None
                        self._event_ch.send_nowait(
                            llm.ChatChunk(
                                id=request_id,
                                delta=llm.ChoiceDelta(
                                    content=f"Agent 错误: {err_msg}"
                                ),
                            )
                        )
                        return
