"""RND1 diffusion language model provider for local completions.

RND1 is a 30B sparse MoE diffusion LM (3B active parameters) that generates
text bidirectionally. Unlike autoregressive models, diffusion LMs iteratively
refine the entire output simultaneously, which can improve coherence for
certain tasks.

This provider connects to a locally-served RND1 instance via HTTP, typically
using vLLM or SGLang as the inference backend with an OpenAI-compatible API.

IMPORTANT LIMITATIONS:
    - RND1-Base is NOT instruction-tuned. It cannot reliably:
        * Follow tool-calling conventions (function calls)
        * Generate structured output (JSON, XML)
        * Adhere to specific formatting instructions
    - For tasks requiring structured output, use the two-stage pattern:
        RND1 (free-form analysis) -> Structured Model (Claude/GPT) for actions

Best suited for:
    - Free-form text generation and completion
    - Holistic query analysis (complexity estimation, approach suggestion)
    - Tasks where bidirectional context improves coherence
"""

from __future__ import annotations

import json
import logging
from collections.abc import Generator
from typing import Any

import httpx

logger = logging.getLogger(__name__)

from ember._internal.exceptions import ProviderAPIError
from ember.models.providers.base import BaseProvider
from ember.models.schemas import ChatResponse, UsageStats


class RND1Provider(BaseProvider):
    """Keyless provider for local RND1 diffusion model inference.

    RND1 is deployed as a local server (vLLM/SGLang) exposing an
    OpenAI-compatible API. This provider handles the HTTP communication
    and diffusion-specific parameters.

    Configuration (via ~/.ember/config.yaml):
        providers:
          rnd1:
            base_url: "http://localhost:8080"
            timeout_ms: 60000
            num_diffusion_steps: 256
            temperature: 0.01
            default_model: "rnd1-base-0910"

    Diffusion-specific parameters:
        - num_diffusion_steps: Number of denoising steps (default 256).
          Higher values improve quality but increase latency.
        - temperature: Controls randomness in diffusion (default 0.01).
          Lower values produce more deterministic outputs.

    Capability Limitations:
        RND1-Base is NOT instruction-tuned. The following capabilities
        are NOT reliably supported:
        - supports_tool_use: False
        - supports_structured_output: False
        - supports_system_prompt: False (weak adherence)
    """

    requires_api_key: bool = False

    # Capability flags - RND1-Base is NOT instruction-tuned
    supports_tool_use: bool = False
    supports_structured_output: bool = False
    supports_function_calling: bool = False

    def _get_api_key_from_env(self) -> str:
        return ""

    @staticmethod
    def _base_url(**kwargs: Any) -> str:
        """Resolve the RND1 server base URL from kwargs or config."""
        if "base_url" in kwargs:
            candidate = kwargs["base_url"]
            if not isinstance(candidate, str):
                raise ProviderAPIError(
                    "base_url must be a string",
                    context={"provider": "rnd1", "value_type": type(candidate).__name__},
                )
            base_url = candidate.strip()
            return base_url or "http://localhost:8080"

        candidate = _config_value("base_url")
        if candidate is None:
            return "http://localhost:8080"
        if not isinstance(candidate, str):
            type_name = type(candidate).__name__
            raise TypeError(f"providers.rnd1.base_url must be a string, got {type_name}")
        base_url = candidate.strip()
        return base_url or "http://localhost:8080"

    @staticmethod
    def _timeout(**kwargs: Any) -> float:
        """Resolve timeout in seconds from kwargs or config.

        Default is 60 seconds since diffusion models are slower than AR models.
        """
        if "timeout" in kwargs:
            try:
                return float(kwargs["timeout"])
            except (TypeError, ValueError) as exc:
                raise ProviderAPIError(
                    "timeout must be numeric (seconds)",
                    context={"provider": "rnd1", "value": kwargs["timeout"]},
                ) from exc

        config_ms = _config_value("timeout_ms")
        if config_ms is None:
            return 60.0  # Default 60s for diffusion (slower than AR models)
        if isinstance(config_ms, bool):
            raise TypeError("providers.rnd1.timeout_ms must be numeric")
        if isinstance(config_ms, (int, float)):
            return float(config_ms) / 1000.0
        if isinstance(config_ms, str):
            try:
                return float(config_ms) / 1000.0
            except ValueError as exc:
                raise TypeError("providers.rnd1.timeout_ms must be numeric") from exc
        type_name = type(config_ms).__name__
        raise TypeError(f"providers.rnd1.timeout_ms must be numeric, got {type_name}")

    @staticmethod
    def _num_diffusion_steps(**kwargs: Any) -> int:
        """Resolve number of diffusion steps from kwargs or config.

        Higher step counts improve output quality but increase latency.
        Default is 256 steps which provides good quality/speed balance.
        """
        if "num_diffusion_steps" in kwargs:
            try:
                return int(kwargs["num_diffusion_steps"])
            except (TypeError, ValueError) as exc:
                raise ProviderAPIError(
                    "num_diffusion_steps must be an integer",
                    context={"provider": "rnd1", "value": kwargs["num_diffusion_steps"]},
                ) from exc

        config_val = _config_value("num_diffusion_steps")
        if config_val is None:
            return 256  # Default for good quality/speed balance
        if isinstance(config_val, bool):
            raise TypeError("providers.rnd1.num_diffusion_steps must be an integer")
        if isinstance(config_val, int):
            return config_val
        if isinstance(config_val, str):
            try:
                return int(config_val)
            except ValueError as exc:
                raise TypeError("providers.rnd1.num_diffusion_steps must be an integer") from exc
        type_name = type(config_val).__name__
        raise TypeError(f"providers.rnd1.num_diffusion_steps must be an integer, got {type_name}")

    @staticmethod
    def _diffusion_temperature(**kwargs: Any) -> float:
        """Resolve diffusion temperature from kwargs or config.

        Lower temperatures produce more deterministic outputs.
        Default is 0.01 for near-deterministic generation.
        """
        if "diffusion_temperature" in kwargs:
            try:
                return float(kwargs["diffusion_temperature"])
            except (TypeError, ValueError) as exc:
                raise ProviderAPIError(
                    "diffusion_temperature must be numeric",
                    context={"provider": "rnd1", "value": kwargs["diffusion_temperature"]},
                ) from exc

        config_val = _config_value("temperature")
        if config_val is None:
            return 0.01  # Default for near-deterministic diffusion
        if isinstance(config_val, bool):
            raise TypeError("providers.rnd1.temperature must be numeric")
        if isinstance(config_val, (int, float)):
            return float(config_val)
        if isinstance(config_val, str):
            try:
                return float(config_val)
            except ValueError as exc:
                raise TypeError("providers.rnd1.temperature must be numeric") from exc
        type_name = type(config_val).__name__
        raise TypeError(f"providers.rnd1.temperature must be numeric, got {type_name}")

    @staticmethod
    def _estimate_tokens(text: str) -> int:
        """Estimate token count for usage statistics."""
        return estimate_tokens(text, _config_value("token_char_ratio"), "rnd1")

    def _choose_default_model(self) -> str:
        """Return the configured default model or fallback."""
        default_model = _config_value("default_model")
        if default_model is None:
            return "rnd1-base-0910"
        if not isinstance(default_model, str):
            type_name = type(default_model).__name__
            raise TypeError(f"providers.rnd1.default_model must be a string, got {type_name}")
        candidate = default_model.strip()
        return candidate or "rnd1-base-0910"

    def complete(self, prompt: str, model: str, **kwargs: Any) -> ChatResponse:
        """Execute a completion request against the local RND1 server.

        The server is expected to expose an OpenAI-compatible chat completions
        endpoint at /v1/chat/completions. Diffusion-specific parameters are
        passed via extra_body when supported.

        Args:
            prompt: The user prompt text.
            model: Model identifier (e.g., "rnd1-base-0910").
            **kwargs: Additional options including:
                - temperature: Sampling temperature for the response.
                - max_tokens: Maximum tokens to generate.
                - num_diffusion_steps: Diffusion denoising steps.
                - diffusion_temperature: Temperature for diffusion process.
                - system: System prompt content.
                - context: Alias for system prompt.

        Returns:
            ChatResponse with generated text and usage statistics.

        Raises:
            ProviderAPIError: On connection or HTTP errors.
        """
        base_url = self._base_url(**kwargs)
        timeout = self._timeout(**kwargs)

        requested_model = model
        if model in {"auto", "default"}:
            requested_model = self._choose_default_model()

        # Build messages for chat completion format
        messages: list[dict[str, str]] = []
        system = kwargs.pop("context", None) or kwargs.pop("system", None)
        if system:
            messages.append({"role": "system", "content": str(system)})
        messages.append({"role": "user", "content": str(prompt)})

        # Standard OpenAI-compatible parameters
        body: dict[str, Any] = {
            "model": requested_model,
            "messages": messages,
        }

        if "temperature" in kwargs:
            body["temperature"] = kwargs.pop("temperature")
        if "max_tokens" in kwargs:
            body["max_tokens"] = kwargs.pop("max_tokens")
        if "stop" in kwargs:
            body["stop"] = kwargs.pop("stop")

        # Diffusion-specific parameters via extra_body (vLLM/SGLang convention)
        extra_body: dict[str, Any] = {}
        num_steps = self._num_diffusion_steps(**kwargs)
        diff_temp = self._diffusion_temperature(**kwargs)
        extra_body["num_diffusion_steps"] = num_steps
        extra_body["diffusion_temperature"] = diff_temp

        if extra_body:
            body["extra_body"] = extra_body

        stream = bool(kwargs.pop("stream", False))

        url = f"{base_url}/v1/chat/completions"

        try:
            if stream:
                gen = self.stream_complete(
                    prompt,
                    requested_model,
                    base_url=base_url,
                    timeout=timeout,
                    system=system,
                    extra_body=extra_body,
                    **{k: v for k, v in kwargs.items() if k not in ("base_url", "timeout")},
                )
                chunks: list[str] = []
                try:
                    while True:
                        chunks.append(next(gen))
                except StopIteration as exc:
                    final = exc.value
                    if not isinstance(final, ChatResponse):
                        raise ProviderAPIError(
                            "Streaming response was not finalized",
                            context={"provider": "rnd1", "model": requested_model},
                        ) from exc
                    if not final.data:
                        final.data = "".join(chunks)
                    return final
            else:
                with httpx.Client(timeout=timeout) as client:
                    resp = client.post(url, json=body)
                    resp.raise_for_status()
                    data = resp.json()

                    # Parse OpenAI-compatible response
                    choices = data.get("choices", [])
                    if not choices:
                        raise ProviderAPIError(
                            "No choices in response",
                            context={"provider": "rnd1", "model": requested_model},
                        )
                    message = choices[0].get("message", {})
                    text = message.get("content", "")

                    # Extract usage if provided, otherwise estimate
                    usage_data = data.get("usage", {})
                    if usage_data:
                        pt = usage_data.get("prompt_tokens", self._estimate_tokens(str(prompt)))
                        ct = usage_data.get("completion_tokens", self._estimate_tokens(text))
                    else:
                        pt = self._estimate_tokens(str(prompt))
                        ct = self._estimate_tokens(text)

                    usage = UsageStats(prompt_tokens=pt, completion_tokens=ct, total_tokens=pt + ct)

                    return ChatResponse(
                        data=text or "",
                        usage=usage,
                        model_id=requested_model,
                        raw_output=data,
                    )

        except httpx.ConnectError as e:
            raise ProviderAPIError(
                "Could not reach RND1 server",
                context={
                    "model": model,
                    "error_type": "connection_refused",
                    "hint": "Start the RND1 server with: python scripts/serve_rnd1.py",
                    "base_url": base_url,
                },
            ) from e
        except httpx.HTTPStatusError as e:
            detail = e.response.text if e.response is not None else str(e)
            raise ProviderAPIError(
                f"RND1 HTTP error: {e}",
                context={"model": model, "detail": detail, "base_url": base_url},
            ) from e
        except httpx.RequestError as e:
            raise ProviderAPIError(
                f"RND1 request failed: {e}",
                context={"model": model, "base_url": base_url},
            ) from e
        except Exception as e:
            raise ProviderAPIError(
                f"Unexpected RND1 error: {e}",
                context={"model": model, "base_url": base_url},
            ) from e

    def stream_complete(
        self, prompt: str, model: str, **kwargs: Any
    ) -> Generator[str, None, ChatResponse]:
        """Stream completion chunks from the RND1 server.

        Args:
            prompt: The user prompt text.
            model: Model identifier.
            **kwargs: Additional options.

        Yields:
            Text chunks as they arrive.

        Returns:
            Final ChatResponse via generator return value.
        """
        base_url = self._base_url(**kwargs)
        timeout = self._timeout(**kwargs)
        system = kwargs.get("system")
        extra_body = kwargs.get("extra_body", {})

        if model in {"auto", "default"}:
            model = self._choose_default_model()

        messages: list[dict[str, str]] = []
        if system:
            messages.append({"role": "system", "content": str(system)})
        messages.append({"role": "user", "content": str(prompt)})

        body: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": True,
        }

        if "temperature" in kwargs:
            body["temperature"] = kwargs["temperature"]
        if "max_tokens" in kwargs:
            body["max_tokens"] = kwargs["max_tokens"]
        if extra_body:
            body["extra_body"] = extra_body

        url = f"{base_url}/v1/chat/completions"
        text_accum: list[str] = []

        try:
            with httpx.Client(timeout=timeout) as client:
                with client.stream("POST", url, json=body) as resp:
                    resp.raise_for_status()
                    for line in resp.iter_lines():
                        if not line:
                            continue
                        # SSE format: "data: {...}"
                        if line.startswith("data: "):
                            line = line[6:]
                        if line == "[DONE]":
                            break
                        try:
                            part = json.loads(line)
                        except Exception:
                            logger.debug("RND1 stream: skipping unparseable line: %s", line[:200])
                            continue

                        choices = part.get("choices", [])
                        if choices:
                            delta = choices[0].get("delta", {})
                            content = delta.get("content", "")
                            if content:
                                text_accum.append(content)
                                yield content

            # Finalize
            text = "".join(text_accum)
            pt = self._estimate_tokens(str(prompt))
            ct = self._estimate_tokens(text)
            usage = UsageStats(prompt_tokens=pt, completion_tokens=ct, total_tokens=pt + ct)
            final = ChatResponse(
                data=text,
                usage=usage,
                model_id=model,
                raw_output={"streamed": True, "endpoint": url},
            )
            return final
        except httpx.ConnectError as e:
            raise ProviderAPIError(
                "Could not reach RND1 server",
                context={
                    "model": model,
                    "error_type": "connection_refused",
                    "hint": "Start the RND1 server with: python scripts/serve_rnd1.py",
                    "base_url": base_url,
                },
            ) from e
        except httpx.RequestError as e:
            raise ProviderAPIError(
                f"RND1 request failed: {e}",
                context={"model": model, "base_url": base_url},
            ) from e


from ember.models.providers._config import estimate_tokens, make_config_helpers

_config_value, _config_str = make_config_helpers("rnd1")
