"""SageLLM Generator - 统一 LLM 推理算子。

通过 ``sage.libs.llm`` 协议层（L3 接口）调用 LLM 推理服务，不直接依赖
sagellm 包。底层使用内联的 OpenAI-compatible HTTP 客户端，
兼容 SageLLM、vLLM、Ollama 等后端。

See: sage.libs.llm.protocol.LLMServiceProtocol
"""

from __future__ import annotations

import logging
import os
from collections.abc import AsyncGenerator, Sequence
from dataclasses import dataclass, field
from typing import Any

from sage.common.config.ports import SagePorts
from sage.common.core.functions import MapFunction as MapOperator
from sage.libs.llm.types import LLMRequest, LLMResponse, LLMStreamChunk


class _OpenAICompatHTTPClient:
    """Minimal inline OpenAI-compatible HTTP client for sage-middleware.

    Implements the same generate()/stream() surface as LLMServiceProtocol
    without depending on the deleted sage.libs.llm.adapters package.
    Requires the 'openai' package (``pip install openai``).
    """

    def __init__(
        self,
        base_url: str,
        api_key: str = "EMPTY",
        model: str = "",
        timeout: float = 120.0,
    ) -> None:
        try:
            import openai  # type: ignore[import]
        except ImportError as exc:
            raise ImportError(
                "'openai' package is required for SageLLMGenerator: pip install openai"
            ) from exc
        self._client = openai.OpenAI(
            base_url=base_url,
            api_key=api_key,
            timeout=timeout,
        )
        self._model = model

    def generate(self, request: LLMRequest) -> LLMResponse:
        import openai  # type: ignore[import]

        model = request.model or self._model
        try:
            completion = self._client.completions.create(
                model=model,
                prompt=request.prompt,
                max_tokens=request.max_tokens,
                temperature=request.temperature if request.temperature is not None else 1.0,
                top_p=request.top_p if request.top_p is not None else 1.0,
            )
            choice = completion.choices[0]
            usage = completion.usage
            return LLMResponse(
                text=choice.text,
                prompt_tokens=usage.prompt_tokens if usage else 0,
                completion_tokens=usage.completion_tokens if usage else 0,
                model=model,
            )
        except openai.OpenAIError as exc:
            raise RuntimeError(f"OpenAI-compat request failed: {exc}") from exc

    def stream(self, request: LLMRequest):
        import openai  # type: ignore[import]

        model = request.model or self._model
        try:
            stream = self._client.completions.create(
                model=model,
                prompt=request.prompt,
                max_tokens=request.max_tokens,
                temperature=request.temperature if request.temperature is not None else 1.0,
                top_p=request.top_p if request.top_p is not None else 1.0,
                stream=True,
            )
            for chunk in stream:
                delta = chunk.choices[0].text if chunk.choices else ""
                yield LLMStreamChunk(delta=delta)
        except openai.OpenAIError as exc:
            raise RuntimeError(f"OpenAI-compat stream failed: {exc}") from exc

    def is_available(self) -> bool:
        try:
            self._client.models.list()
            return True
        except Exception:
            return False


logger = logging.getLogger(__name__)


def _normalize_input(data: Any) -> tuple[dict[str, Any], str, dict[str, Any]]:
    """Normalize input into ``(context, prompt, options)``."""
    context: dict[str, Any] = {}
    prompt: str = ""
    options: dict[str, Any] = {}

    if isinstance(data, str):
        prompt = data
    elif isinstance(data, dict):
        prompt = data.get("prompt", "")
        options = dict(data.get("options", {}))
        context = {key: value for key, value in data.items() if key not in ("prompt", "options")}
    elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
        if len(data) >= 1:
            first = data[0]
            if isinstance(first, str):
                prompt = first
            elif isinstance(first, dict):
                context = dict(first)
        if len(data) >= 2:
            second = data[1]
            if isinstance(second, str):
                prompt = second
            elif isinstance(second, dict) and "prompt" in second:
                prompt = second["prompt"]
                options.update(second.get("options", {}))
        if len(data) >= 3 and isinstance(data[2], dict):
            options.update(data[2])
    else:
        prompt = str(data)

    return context, prompt, options


@dataclass
class SageLLMGenerator(MapOperator):
    """SageLLM unified generator in service-port mode.

    All inference requests are sent to a single OpenAI-compatible SageLLM
    endpoint. Middleware no longer creates local backends (mock/cuda/ascend).
    """

    model_path: str = ""
    base_url: str = ""
    api_key: str = "EMPTY"

    max_tokens: int = 2048
    max_new_tokens: int = 128
    temperature: float = 0.7
    top_p: float = 0.95
    top_k: int = 50

    engine_id: str = ""
    timeout: float = 120.0
    default_options: dict[str, Any] = field(default_factory=dict)

    # Historical fields kept in config surface; no local backend behavior.
    backend_type: str = "service"
    device_map: str = ""
    dtype: str = ""
    device: str = ""
    load_in_8bit: bool = False
    load_in_4bit: bool = False
    trust_remote_code: bool = False

    _backend: _OpenAICompatHTTPClient | None = field(default=None, init=False, repr=False)
    _initialized: bool = field(default=False, init=False, repr=False)

    def __post_init__(self) -> None:
        super().__init__()
        if not self.engine_id:
            self.engine_id = f"sage-llm-{id(self)}"

    def _resolve_base_url(self) -> str:
        """Resolve OpenAI-compatible inference base URL."""
        if self.base_url:
            return self.base_url.rstrip("/")

        env_base_url = os.getenv("SAGE_MIDDLEWARE_INFERENCE_BASE_URL", "").strip()
        if env_base_url:
            return env_base_url.rstrip("/")

        return f"http://127.0.0.1:{SagePorts.SAGELLM_SERVE_PORT}/v1"

    def _ensure_client(self) -> None:
        """Ensure inference backend is initialized."""
        if self._initialized and self._backend is not None:
            return

        resolved_base_url = self._resolve_base_url()
        self._backend = _OpenAICompatHTTPClient(
            base_url=resolved_base_url,
            api_key=self.api_key or "EMPTY",
            model=self.model_path or "",
            timeout=self.timeout,
        )
        self._initialized = True
        logger.info(
            "SageLLMGenerator initialized via sage.libs.llm protocol on base_url=%s engine_id=%s",
            resolved_base_url,
            self.engine_id,
        )

    def _build_generation_params(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
        """Build generation params from defaults and user options."""
        params = {
            "prompt": prompt,
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "top_k": self.top_k,
        }
        params.update(self.default_options)
        params.update({key: value for key, value in options.items() if value is not None})
        return params

    def execute(self, data: Any) -> dict[str, Any]:
        """Execute a single completion request."""
        self._ensure_client()

        context, prompt, options = _normalize_input(data)
        params = self._build_generation_params(prompt, options)

        if not prompt:
            logger.warning("Empty prompt received, returning empty result")
            return {"text": "", "usage": {}, "context": context}

        try:
            request = LLMRequest(
                prompt=prompt,
                model=self.model_path or "",
                max_tokens=int(params.get("max_tokens", self.max_tokens)),
                temperature=params.get("temperature"),
                top_p=params.get("top_p"),
                extra={
                    k: v
                    for k, v in params.items()
                    if k not in ("prompt", "max_tokens", "temperature", "top_p")
                },
            )
            response = self._backend.generate(request)  # type: ignore[union-attr]
            usage = {
                "prompt_tokens": response.prompt_tokens,
                "completion_tokens": response.completion_tokens,
                "total_tokens": response.prompt_tokens + response.completion_tokens,
            }
            output = {"text": response.text, "usage": usage}
            if context:
                output["context"] = context
            return output
        except Exception as error:
            logger.error("Generation failed: %s", error)
            raise RuntimeError(f"SageLLM generation failed: {error}") from error

    async def stream_async(self, data: Any) -> AsyncGenerator[dict[str, Any], None]:
        """Stream generation via OpenAI-compatible API."""
        self._ensure_client()

        _context, prompt, options = _normalize_input(data)
        params = self._build_generation_params(prompt, options)

        if not prompt:
            logger.warning("Empty prompt received for streaming")
            yield {"text": "", "done": True, "usage": {}}
            return

        try:
            request = LLMRequest(
                prompt=prompt,
                model=self.model_path or "",
                max_tokens=int(params.get("max_tokens", self.max_tokens)),
                temperature=params.get("temperature"),
                top_p=params.get("top_p"),
                stream=True,
            )
            for chunk in self._backend.stream(request):  # type: ignore[union-attr]
                if chunk.delta:
                    yield {"text": chunk.delta, "done": False}

            yield {"text": "", "done": True, "usage": {}}

        except Exception as error:
            logger.error("Streaming generation failed: %s", error)
            yield {"text": "", "done": True, "error": str(error)}

    def shutdown(self) -> None:
        """Release backend resources."""
        self._backend = None
        self._initialized = False

    def __del__(self) -> None:
        try:
            self.shutdown()
        except Exception:
            pass


__all__ = ["SageLLMGenerator"]
