# SPDX-License-Identifier: Apache-2.0
"""Client wrapper for coordinator-based pipelines."""

from __future__ import annotations

import asyncio
import uuid
from dataclasses import replace
from typing import Any, AsyncIterator, Callable

import numpy as np

from sglang_omni.client.audio import (
    DEFAULT_SAMPLE_RATE,
    FORMAT_MIME_TYPES,
    audio_to_base64,
    encode_audio,
    to_numpy,
)
from sglang_omni.client.types import (
    AbortLevel,
    AbortResult,
    ClientError,
    CompletionAudio,
    CompletionResult,
    CompletionStreamChunk,
    GenerateChunk,
    GenerateRequest,
    SpeechResult,
    UsageInfo,
)
from sglang_omni.pipeline.coordinator import Coordinator
from sglang_omni.proto import OmniRequest, RequestState, StreamMessage


class Client:
    """Internal client used by API adapters."""

    def __init__(
        self,
        coordinator: Coordinator,
        result_builder: Callable[[str, Any], GenerateChunk] | None = None,
        stream_builder: Callable[[str, StreamMessage], GenerateChunk] | None = None,
    ) -> None:
        self._coordinator = coordinator
        self._result_builder = result_builder or self._default_result_builder
        self._stream_builder = stream_builder or self._default_stream_builder

    # ------------------------------------------------------------------
    # Low-level generate (backward compatible)
    # ------------------------------------------------------------------

    async def generate(
        self,
        request: GenerateRequest,
        request_id: str | None = None,
    ) -> AsyncIterator[GenerateChunk]:
        req_id = request_id or str(uuid.uuid4())
        omni_request = self._build_omni_request(request)
        if request.stream:
            async for msg in self._coordinator.stream(req_id, omni_request):
                if isinstance(msg, StreamMessage):
                    yield self._stream_builder(req_id, msg)
                else:
                    yield self._result_builder(req_id, msg.result)
            return

        result = await self._coordinator.submit(req_id, omni_request)
        yield self._result_builder(req_id, result)

    # ------------------------------------------------------------------
    # High-level: non-streaming completion
    # ------------------------------------------------------------------

    async def completion(
        self,
        request: GenerateRequest,
        *,
        request_id: str,
        audio_format: str = "wav",
    ) -> CompletionResult:
        """Run a non-streaming completion and return an aggregated result.

        Iterates ``generate()``, accumulates text, concatenates audio chunks,
        and encodes audio to base64.

        Raises:
            ClientError: If the pipeline produces no response at all.
        """
        text_parts: list[str] = []
        audio_chunks: list[Any] = []
        sample_rate: int | None = None
        last_chunk: GenerateChunk | None = None
        finish_reason: str | None = None
        logprobs_parts: list[Any] = []
        saw_output_token_logprobs = False
        omni_rollout: dict[str, Any] | None = None
        weight_version: str | None = None

        async for chunk in self.generate(request, request_id=request_id):
            last_chunk = chunk
            if chunk.text:
                text_parts.append(chunk.text)
            if chunk.audio_data is not None:
                audio_chunks.append(chunk.audio_data)
            if chunk.sample_rate is not None:
                sample_rate = chunk.sample_rate
            if chunk.finish_reason is not None:
                finish_reason = chunk.finish_reason
            if chunk.output_token_logprobs is not None:
                saw_output_token_logprobs = True
                logprobs_parts.extend(chunk.output_token_logprobs)
            if chunk.omni_rollout is not None:
                omni_rollout = chunk.omni_rollout
            if chunk.weight_version is not None:
                weight_version = chunk.weight_version

        if last_chunk is None:
            raise ClientError("No response from pipeline")

        full_text = "".join(text_parts)

        audio: CompletionAudio | None = None
        if audio_chunks:
            if len(audio_chunks) == 1:
                combined = audio_chunks[0]
            else:
                arrays = [to_numpy(c) for c in audio_chunks]
                axis = -1 if arrays[0].ndim > 1 else 0
                combined = np.concatenate(arrays, axis=axis)
            audio_b64 = audio_to_base64(
                combined,
                sample_rate=sample_rate or DEFAULT_SAMPLE_RATE,
                output_format=audio_format,
            )
            audio = CompletionAudio(
                id=f"audio-{request_id}",
                data=audio_b64,
                transcript=full_text if full_text else None,
            )

        return CompletionResult(
            request_id=request_id,
            text=full_text,
            audio=audio,
            finish_reason=finish_reason or "stop",
            usage=last_chunk.usage,
            output_token_logprobs=(
                logprobs_parts if saw_output_token_logprobs else None
            ),
            omni_rollout=omni_rollout,
            weight_version=weight_version,
        )

    # ------------------------------------------------------------------
    # High-level: streaming completion
    # ------------------------------------------------------------------

    async def completion_stream(
        self,
        request: GenerateRequest,
        *,
        request_id: str,
        audio_format: str = "wav",
    ) -> AsyncIterator[CompletionStreamChunk]:
        """Iterate ``generate()`` and yield high-level stream chunks.

        Audio data is base64-encoded before yielding so that callers never
        need to touch numpy / raw bytes.
        """
        async for chunk in self.generate(request, request_id=request_id):
            audio_b64: str | None = None
            if chunk.modality == "audio" and chunk.audio_data is not None:
                audio_b64 = audio_to_base64(
                    chunk.audio_data,
                    sample_rate=chunk.sample_rate or DEFAULT_SAMPLE_RATE,
                    output_format=audio_format,
                )

            yield CompletionStreamChunk(
                request_id=request_id,
                text=chunk.text,
                modality=chunk.modality,
                audio_b64=audio_b64,
                finish_reason=chunk.finish_reason,
                usage=chunk.usage,
                stage_name=chunk.stage_name,
            )

    # ------------------------------------------------------------------
    # High-level: text-to-speech
    # ------------------------------------------------------------------

    async def speech(
        self,
        request: GenerateRequest,
        *,
        request_id: str,
        response_format: str = "wav",
        speed: float = 1.0,
        allow_format_fallback: bool = True,
    ) -> SpeechResult:
        """Run a TTS request and return encoded audio bytes.

        Raises:
            ClientError: If the pipeline produces no audio output.
        """
        audio_chunks: list[Any] = []
        sample_rate: int | None = None
        last_chunk: GenerateChunk | None = None
        extra_params = dict(request.extra_params)
        extra_params.pop("stream", None)
        request = replace(request, stream=False, extra_params=extra_params)

        async for chunk in self.generate(request, request_id=request_id):
            if chunk.audio_data is not None:
                audio_chunks.append(chunk.audio_data)
            if chunk.sample_rate is not None:
                sample_rate = chunk.sample_rate
            last_chunk = chunk

        if not audio_chunks:
            raise ClientError("No audio output generated from the pipeline.")

        if len(audio_chunks) == 1:
            audio_data = audio_chunks[0]
        else:
            arrays = [to_numpy(c) for c in audio_chunks]
            axis = -1 if arrays[0].ndim > 1 else 0
            audio_data = np.concatenate(arrays, axis=axis)

        encode_kwargs: dict[str, Any] = {
            "response_format": response_format,
            "speed": speed,
            "allow_format_fallback": allow_format_fallback,
        }
        if sample_rate is not None:
            encode_kwargs["sample_rate"] = sample_rate

        audio_bytes, mime_type = await asyncio.to_thread(
            encode_audio, audio_data, **encode_kwargs
        )

        # Derive actual format from MIME type (encode_audio may fall back
        # to WAV if the requested codec is unavailable).
        actual_format = response_format
        for ext, mt in FORMAT_MIME_TYPES.items():
            if mt == mime_type:
                actual_format = ext
                break

        return SpeechResult(
            audio_bytes=audio_bytes,
            mime_type=mime_type,
            format=actual_format,
            sample_rate=sample_rate,
            usage=last_chunk.usage if last_chunk else None,
        )

    # ------------------------------------------------------------------
    # Other operations
    # ------------------------------------------------------------------

    async def abort(
        self,
        request_id: str,
        level: AbortLevel = AbortLevel.SOFT,
    ) -> AbortResult:
        success = await self._coordinator.abort(request_id)
        return AbortResult(success=success, level_applied=level)

    async def get_status(self, request_id: str) -> RequestState | None:
        info = self._coordinator.get_request_info(request_id)
        if info is None:
            return None
        return info.state

    def health(self) -> dict[str, Any]:
        return self._coordinator.health()

    async def admin(
        self,
        action: str,
        payload: dict[str, Any] | None = None,
        *,
        stages: list[str] | None = None,
        timeout_s: float = 60.0,
    ) -> dict[str, Any]:
        return await self._coordinator.admin(
            action,
            payload,
            stages=stages,
            timeout_s=timeout_s,
        )

    async def model_info(
        self,
        *,
        stages: list[str] | None = None,
        timeout_s: float = 30.0,
    ) -> dict[str, Any]:
        return await self._coordinator.model_info(
            stages=stages,
            timeout_s=timeout_s,
        )

    async def pause_generation(
        self,
        payload: dict[str, Any] | None = None,
        *,
        stages: list[str] | None = None,
        timeout_s: float = 60.0,
    ) -> dict[str, Any]:
        return await self._coordinator.pause_generation(
            payload,
            stages=stages,
            timeout_s=timeout_s,
        )

    async def continue_generation(
        self,
        payload: dict[str, Any] | None = None,
        *,
        stages: list[str] | None = None,
        timeout_s: float = 60.0,
    ) -> dict[str, Any]:
        return await self._coordinator.continue_generation(
            payload,
            stages=stages,
            timeout_s=timeout_s,
        )

    async def update_weights_from_disk(
        self,
        payload: dict[str, Any],
        *,
        stages: list[str] | None = None,
        timeout_s: float = 120.0,
    ) -> dict[str, Any]:
        return await self._coordinator.update_weights_from_disk(
            payload,
            stages=stages,
            timeout_s=timeout_s,
        )

    async def init_weights_update_group(
        self,
        payload: dict[str, Any],
        *,
        stages: list[str] | None = None,
        timeout_s: float = 300.0,
    ) -> dict[str, Any]:
        return await self._coordinator.init_weights_update_group(
            payload,
            stages=stages,
            timeout_s=timeout_s,
        )

    async def destroy_weights_update_group(
        self,
        payload: dict[str, Any],
        *,
        stages: list[str] | None = None,
        timeout_s: float = 300.0,
    ) -> dict[str, Any]:
        return await self._coordinator.destroy_weights_update_group(
            payload,
            stages=stages,
            timeout_s=timeout_s,
        )

    async def update_weights_from_distributed(
        self,
        payload: dict[str, Any],
        *,
        stages: list[str] | None = None,
        timeout_s: float = 300.0,
    ) -> dict[str, Any]:
        return await self._coordinator.update_weights_from_distributed(
            payload,
            stages=stages,
            timeout_s=timeout_s,
        )

    async def weights_checker(
        self,
        payload: dict[str, Any] | None = None,
        *,
        stages: list[str] | None = None,
        timeout_s: float = 120.0,
    ) -> dict[str, Any]:
        return await self._coordinator.weights_checker(
            payload,
            stages=stages,
            timeout_s=timeout_s,
        )

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    @staticmethod
    def _set_audio_data(chunk: GenerateChunk, data: dict[str, Any]) -> None:
        audio_data = data.get("audio_data") or data.get("audio")
        if audio_data is None and data.get("audio_waveform") is not None:
            raw = data.get("audio_waveform")
            if isinstance(raw, memoryview):
                raw = raw.tobytes()
            dtype = np.dtype(data.get("audio_waveform_dtype", "float32"))
            arr = np.frombuffer(raw, dtype=dtype)
            shape = data.get("audio_waveform_shape")
            if shape:
                arr = arr.reshape(shape)
            audio_data = arr.copy()
        if audio_data is not None:
            chunk.audio_data = audio_data
            chunk.modality = "audio"
        sample_rate = data.get("sample_rate")
        if sample_rate is not None:
            chunk.sample_rate = sample_rate

    @staticmethod
    def _build_usage_info(data: dict[str, Any]) -> UsageInfo | None:
        usage = dict(data.get("usage") or {})
        if "prompt_tokens" not in usage and data.get("prompt_tokens") is not None:
            usage["prompt_tokens"] = data.get("prompt_tokens")
        if (
            "completion_tokens" not in usage
            and data.get("completion_tokens") is not None
        ):
            usage["completion_tokens"] = data.get("completion_tokens")
        if "total_tokens" not in usage:
            prompt_tokens = usage.get("prompt_tokens")
            completion_tokens = usage.get("completion_tokens")
            if prompt_tokens is not None or completion_tokens is not None:
                usage["total_tokens"] = (prompt_tokens or 0) + (completion_tokens or 0)
        if "engine_time_s" not in usage and data.get("engine_time_s") is not None:
            usage["engine_time_s"] = data.get("engine_time_s")
        return UsageInfo.from_dict(usage)

    @staticmethod
    def _build_omni_request(request: GenerateRequest) -> OmniRequest:
        inputs = _extract_inputs(request)
        params = _build_params(request)
        metadata = dict(request.metadata)
        if request.model:
            metadata.setdefault("model", request.model)
        if request.output_modalities:
            metadata["output_modalities"] = request.output_modalities
        return OmniRequest(inputs=inputs, params=params, metadata=metadata)

    @staticmethod
    def _default_result_builder(request_id: str, result: Any) -> GenerateChunk:
        chunk = GenerateChunk(request_id=request_id, finish_reason="stop")
        if isinstance(result, GenerateChunk):
            result.request_id = request_id
            return result
        if isinstance(result, dict):
            # Multi-terminal merged result, e.g. decode + code2wav/talker/
            # talker_stream.
            audio_result = None
            if "decode" in result:
                for audio_stage in ("code2wav", "talker", "talker_stream"):
                    if audio_stage in result:
                        audio_result = result[audio_stage] or {}
                        break
            if audio_result is not None:
                decode_result = result["decode"] or {}
                text = decode_result.get("text")
                if isinstance(text, str):
                    chunk.text = text
                finish_reason = decode_result.get("finish_reason")
                if finish_reason is not None:
                    chunk.finish_reason = finish_reason
                output_token_logprobs = decode_result.get("output_token_logprobs")
                if output_token_logprobs is not None:
                    chunk.output_token_logprobs = output_token_logprobs
                omni_rollout = decode_result.get("omni_rollout")
                if omni_rollout is not None:
                    chunk.omni_rollout = omni_rollout
                weight_version = decode_result.get("weight_version")
                if weight_version is not None:
                    chunk.weight_version = weight_version
                Client._set_audio_data(chunk, audio_result)
                chunk.usage = Client._build_usage_info(
                    decode_result
                ) or Client._build_usage_info(audio_result)
                return chunk
            text = result.get("text")
            if isinstance(text, str):
                chunk.text = text
            token_ids = result.get("token_ids")
            if token_ids is not None:
                if not isinstance(token_ids, (list, tuple)):
                    token_ids = token_ids.tolist()
                chunk.token_ids = list(token_ids)
            logprobs = result.get("logprobs")
            if logprobs is not None:
                chunk.logprobs = logprobs
            output_token_logprobs = result.get("output_token_logprobs")
            if output_token_logprobs is not None:
                chunk.output_token_logprobs = output_token_logprobs
            omni_rollout = result.get("omni_rollout")
            if omni_rollout is not None:
                chunk.omni_rollout = omni_rollout
            weight_version = result.get("weight_version")
            if weight_version is not None:
                chunk.weight_version = weight_version
            finish_reason = result.get("finish_reason")
            if finish_reason is not None:
                chunk.finish_reason = finish_reason
            chunk.stage_id = result.get("stage_id")
            chunk.stage_name = result.get("stage_name")
            modality = result.get("modality")
            if modality is not None:
                chunk.modality = modality
            Client._set_audio_data(chunk, result)
            chunk.usage = Client._build_usage_info(result)
            return chunk
        if isinstance(result, str):
            chunk.text = result
            return chunk
        chunk.text = str(result)
        return chunk

    @staticmethod
    def _default_stream_builder(request_id: str, msg: StreamMessage) -> GenerateChunk:
        chunk = GenerateChunk(request_id=request_id)
        chunk.stage_name = msg.stage_name or msg.from_stage
        chunk.stage_id = msg.stage_id
        if msg.modality:
            chunk.modality = msg.modality

        data = msg.chunk
        if isinstance(data, GenerateChunk):
            data.request_id = request_id
            if data.stage_name is None:
                data.stage_name = chunk.stage_name
            if data.stage_id is None:
                data.stage_id = chunk.stage_id
            if not data.modality and chunk.modality:
                data.modality = chunk.modality
            return data
        if isinstance(data, dict):
            text = data.get("text")
            if isinstance(text, str):
                chunk.text = text
            token_ids = data.get("token_ids")
            if token_ids is not None:
                if not isinstance(token_ids, (list, tuple)):
                    token_ids = token_ids.tolist()
                chunk.token_ids = list(token_ids)
            logprobs = data.get("logprobs")
            if logprobs is not None:
                chunk.logprobs = logprobs
            output_token_logprobs = data.get("output_token_logprobs")
            if output_token_logprobs is not None:
                chunk.output_token_logprobs = output_token_logprobs
            omni_rollout = data.get("omni_rollout")
            if omni_rollout is not None:
                chunk.omni_rollout = omni_rollout
            weight_version = data.get("weight_version")
            if weight_version is not None:
                chunk.weight_version = weight_version
            finish_reason = data.get("finish_reason")
            if finish_reason is not None:
                chunk.finish_reason = finish_reason
            chunk.usage = Client._build_usage_info(data)
            stage_name = data.get("stage_name")
            if stage_name is not None:
                chunk.stage_name = stage_name
            stage_id = data.get("stage_id")
            if stage_id is not None:
                chunk.stage_id = stage_id
            modality = data.get("modality")
            if modality is not None:
                chunk.modality = modality
            Client._set_audio_data(chunk, data)
            return chunk
        if isinstance(data, str):
            chunk.text = data
            return chunk
        if isinstance(data, int):
            chunk.token_ids = [data]
            return chunk
        chunk.text = str(data)
        return chunk


def _extract_inputs(request: GenerateRequest) -> Any:
    choices = [
        request.prompt is not None,
        request.prompt_token_ids is not None,
        request.messages is not None,
    ]
    if sum(choices) != 1:
        raise ValueError(
            "GenerateRequest requires exactly one input: "
            "prompt, prompt_token_ids, or messages."
        )
    if request.prompt is not None:
        return request.prompt
    if request.prompt_token_ids is not None:
        return list(request.prompt_token_ids)

    # Build messages list
    messages = [msg.to_dict() for msg in request.messages or []]

    # Check if we have audios, images, or videos in metadata
    audios = request.metadata.get("audios")
    images = request.metadata.get("images")
    videos = request.metadata.get("videos")

    # If we have any media, return a dict with messages and media
    # Otherwise, return just the messages list (for backward compatibility)
    if audios or images or videos:
        result = {"messages": messages}
        if images:
            result["images"] = images
        if audios:
            result["audios"] = audios
        if videos:
            result["videos"] = videos
        for key in (
            "video_fps",
            "video_max_frames",
            "video_min_pixels",
            "video_max_pixels",
            "video_total_pixels",
        ):
            value = request.metadata.get(key)
            if value is not None:
                result[key] = value
        return result
    return messages


def _build_params(request: GenerateRequest) -> dict[str, Any]:
    params = request.sampling.to_dict()
    max_new_tokens = request.sampling.max_new_tokens
    if request.max_tokens is not None:
        max_new_tokens = request.max_tokens
    if max_new_tokens is None:
        params.pop("max_new_tokens", None)
    else:
        params["max_new_tokens"] = max_new_tokens
    params["stream"] = request.stream
    if request.stage_sampling:
        params["stage_sampling"] = {
            key: value.to_dict() for key, value in request.stage_sampling.items()
        }
    if request.stage_params:
        params["stage_params"] = request.stage_params
    if request.extra_params:
        params.update(request.extra_params)
    return params