"""Hybrid / mixed inference client: combines LLM and Embedding calls (#15).

``MixedInferenceClient`` dispatches each request to either the LLM endpoint
or the Embedding endpoint based on the detected or declared request type.

Architecture
------------
Previously two separate clients existed:
- ``IntelligentEmbeddingClient`` (embedding server)
- ``IntelligentLLMClient``       (sageLLM backend)

This module unifies them:

    client = MixedInferenceClient(llm_url="http://localhost:8000",
                                   embedding_url="http://localhost:8001")
    # LLM completion
    text = client.complete("What is 2+2?")["text"]
    # Embedding
    vec  = client.embed("Hello world")

Both endpoints are routed through a single object with a simple strategy:
- ``text``        → ``complete()`` calls LLM
- ``embedding``   → ``embed()`` calls the embedding server

The class also exposes ``dispatch()`` for mixed workloads: pass a list of
:class:`MixedRequest` objects and it will route each automatically.

Issue
-----
#15 混合推理策略优化 — merge IntelligentEmbeddingClient + IntelligentLLMClient
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from sagellm.client import UnifiedInferenceClient

logger = logging.getLogger(__name__)


class RequestKind(str, Enum):
    """Type of inference request.

    Attributes:
        LLM: Standard text-generation / chat-completion request.
        EMBEDDING: Embedding vector request.
    """

    LLM = "llm"
    EMBEDDING = "embedding"


@dataclass
class MixedRequest:
    """A single request for mixed-inference dispatch.

    Attributes:
        kind: :class:`RequestKind` that determines routing.
        content: Prompt string (for ``LLM``) or text to embed (for
            ``EMBEDDING``).
        model: Optional model override.
        max_tokens: Max tokens for LLM requests.
        temperature: Sampling temperature for LLM requests.
        metadata: Arbitrary caller-supplied metadata preserved in the result.
    """

    kind: RequestKind
    content: str
    model: str | None = None
    max_tokens: int = 128
    temperature: float = 0.0
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class MixedResult:
    """Result from a :class:`MixedInferenceClient` dispatch.

    Attributes:
        kind: Original :class:`RequestKind`.
        text: Generated text (for LLM), empty string for embedding-only.
        embedding: Embedding vector (for EMBEDDING), empty list for LLM-only.
        usage: Token usage dict (mirrors OpenAI usage schema).
        metadata: Original request metadata, unchanged.
        raw: Full raw response from the backend.
    """

    kind: RequestKind
    text: str = ""
    embedding: list[float] = field(default_factory=list)
    usage: dict[str, int] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)
    raw: dict[str, Any] = field(default_factory=dict)


class MixedInferenceClient:
    """Unified LLM + Embedding inference client.

    Merges ``IntelligentLLMClient`` and ``IntelligentEmbeddingClient``:
    all inference, whether text-generation or embedding, goes through this
    single client.

    Args:
        llm_url: Base URL of the sageLLM server
            (default ``http://localhost:8000``).
        embedding_url: Base URL of the embedding server
            (default ``http://localhost:8001``).
        llm_model: Default LLM model name.
        embedding_model: Default embedding model name.
        api_key: Shared API key (used for both backends).
        timeout_s: Request timeout in seconds.
        offline_mode: When ``True``, return canned responses without
            making real HTTP requests.  Ideal for unit tests.

    Examples::

        # Normal usage
        client = MixedInferenceClient(
            llm_url="http://localhost:8000",
            embedding_url="http://localhost:8001",
        )
        text_result = client.complete("Tell me a joke")
        vec = client.embed("Hello world")

        # Batch mixed dispatch
        requests = [
            MixedRequest(kind=RequestKind.LLM, content="What is 2+2?"),
            MixedRequest(kind=RequestKind.EMBEDDING, content="hello world"),
        ]
        results = client.dispatch(requests)

        # Offline / test mode
        client = MixedInferenceClient(offline_mode=True)
        assert client.health()["llm"] is True
    """

    def __init__(
        self,
        *,
        llm_url: str = "http://localhost:8000",
        embedding_url: str = "http://localhost:8001",
        llm_model: str | None = None,
        embedding_model: str | None = None,
        api_key: str | None = None,
        timeout_s: float = 60.0,
        offline_mode: bool = False,
    ) -> None:
        self._llm_client = UnifiedInferenceClient(
            base_url=llm_url,
            model=llm_model,
            api_key=api_key,
            timeout_s=timeout_s,
            offline_mode=offline_mode,
        )
        self._embedding_url = embedding_url.rstrip("/")
        self._embedding_model = embedding_model or "default"
        self._api_key = api_key
        self._timeout_s = timeout_s
        self._offline_mode = offline_mode

    # ------------------------------------------------------------------
    # LLM interface
    # ------------------------------------------------------------------

    def complete(
        self,
        prompt: str,
        *,
        model: str | None = None,
        max_tokens: int = 128,
        temperature: float = 0.0,
    ) -> dict[str, Any]:
        """Generate a text completion (LLM path).

        Equivalent to ``UnifiedInferenceClient.complete()``.

        Args:
            prompt: Input text.
            model: Optional model override.
            max_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.

        Returns:
            Dict with ``"text"``, ``"usage"``, and ``"raw"`` keys.
        """
        return self._llm_client.complete(
            prompt, model=model, max_tokens=max_tokens, temperature=temperature
        )

    def chat(
        self,
        messages: list[dict[str, str]],
        *,
        model: str | None = None,
        max_tokens: int = 128,
        temperature: float = 0.0,
    ) -> dict[str, Any]:
        """Multi-turn chat (LLM path).

        Args:
            messages: Conversation history as ``[{"role": ..., "content": ...}]``.
            model: Optional model override.
            max_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.

        Returns:
            Dict with ``"text"``, ``"usage"``, and ``"raw"`` keys.
        """
        return self._llm_client.chat(
            messages, model=model, max_tokens=max_tokens, temperature=temperature
        )

    # ------------------------------------------------------------------
    # Embedding interface
    # ------------------------------------------------------------------

    def embed(
        self,
        text: str | list[str],
        *,
        model: str | None = None,
    ) -> list[list[float]]:
        """Obtain embedding vectors for the given text(s).

        Args:
            text: A single string or a list of strings to embed.
            model: Optional model override.

        Returns:
            A list of embedding vectors (one per input string).
            In ``offline_mode`` returns a list of zero-vectors of dim 4.
        """
        if isinstance(text, str):
            text = [text]

        if self._offline_mode:
            return [[0.0, 0.0, 0.0, 0.0] for _ in text]

        # HTTP request to embedding endpoint  (OpenAI-compatible)
        import json
        from urllib import error as urllib_error
        from urllib import request as urllib_request

        resolved_model = model or self._embedding_model
        payload = {"model": resolved_model, "input": text}
        url = f"{self._embedding_url}/v1/embeddings"

        headers: dict[str, str] = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"

        req = urllib_request.Request(
            url=url,
            data=json.dumps(payload).encode("utf-8"),
            headers=headers,
            method="POST",
        )
        try:
            with urllib_request.urlopen(req, timeout=self._timeout_s) as resp:  # noqa: S310
                data = json.loads(resp.read().decode("utf-8"))
        except urllib_error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"Embedding server returned HTTP {exc.code}: {body}") from exc
        except urllib_error.URLError as exc:
            raise ConnectionError(
                f"Cannot connect to embedding server at {self._embedding_url}: {exc.reason}"
            ) from exc

        return [item["embedding"] for item in data.get("data", [])]

    # ------------------------------------------------------------------
    # Mixed dispatch
    # ------------------------------------------------------------------

    def dispatch(self, requests: list[MixedRequest]) -> list[MixedResult]:
        """Route a batch of mixed requests to the correct backend.

        Each :class:`MixedRequest` with ``kind=LLM`` is sent to the LLM
        endpoint; those with ``kind=EMBEDDING`` to the embedding endpoint.

        Args:
            requests: List of :class:`MixedRequest` objects.

        Returns:
            List of :class:`MixedResult` objects in the same order as input.
        """
        results: list[MixedResult] = []
        for req in requests:
            if req.kind == RequestKind.LLM:
                raw = self._llm_client.complete(
                    req.content,
                    model=req.model,
                    max_tokens=req.max_tokens,
                    temperature=req.temperature,
                )
                results.append(
                    MixedResult(
                        kind=RequestKind.LLM,
                        text=raw.get("text", ""),
                        usage=raw.get("usage", {}),
                        metadata=req.metadata,
                        raw=raw,
                    )
                )
            elif req.kind == RequestKind.EMBEDDING:
                vecs = self.embed(req.content, model=req.model)
                results.append(
                    MixedResult(
                        kind=RequestKind.EMBEDDING,
                        embedding=vecs[0] if vecs else [],
                        metadata=req.metadata,
                        raw={"embeddings": vecs},
                    )
                )
            else:
                raise ValueError(f"Unknown request kind: {req.kind!r}")

        return results

    # ------------------------------------------------------------------
    # Health check
    # ------------------------------------------------------------------

    def health(self) -> dict[str, bool]:
        """Check the health of both backends.

        Returns:
            A dict with ``"llm"`` and ``"embedding"`` boolean flags.
        """
        if self._offline_mode:
            return {"llm": True, "embedding": True}

        from urllib import error as urllib_error
        from urllib import request as urllib_request

        def _check(url: str) -> bool:
            try:
                req = urllib_request.Request(url=f"{url}/health", method="GET")
                with urllib_request.urlopen(req, timeout=5.0) as r:  # noqa: S310
                    return r.status < 400
            except (urllib_error.URLError, urllib_error.HTTPError, Exception):
                return False

        return {
            "llm": self._llm_client.health(),
            "embedding": _check(self._embedding_url),
        }

    def __repr__(self) -> str:
        mode = " offline=True" if self._offline_mode else ""
        return (
            f"MixedInferenceClient("
            f"llm={self._llm_client.base_url!r}, "
            f"embedding={self._embedding_url!r}"
            f"{mode})"
        )
