"""RND1 conductor adapter for local diffusion model inference.

This adapter connects to a locally-served RND1 instance. Partners must
start the server using: python scripts/serve_rnd1.py

RND1 is a diffusion language model that generates text bidirectionally
through iterative denoising. The server exposes an OpenAI-compatible API.

Example usage:

    # Start RND1 server first:
    # python scripts/serve_rnd1.py --port 8080

    from partner_inference_harness.conductor import RND1Conductor, ConductorRequest

    conductor = RND1Conductor(base_url="http://localhost:8080")
    response = conductor.complete(ConductorRequest(
        query="What is machine learning?",
        max_tokens=256,
    ))
    print(response.text)
"""

from __future__ import annotations

import logging
import time
from typing import Any

import httpx

from .interface import (
    ConductorError,
    ConductorInterface,
    ConductorRequest,
    ConductorResponse,
)

logger = logging.getLogger(__name__)


class RND1Conductor(ConductorInterface):
    """Conductor backed by local RND1 diffusion model.

    Connects to an RND1 server running locally via vLLM or SGLang.
    Uses the OpenAI-compatible /v1/chat/completions endpoint.

    Attributes:
        base_url: Base URL of the RND1 server.
        timeout: Request timeout in seconds.
        num_diffusion_steps: Number of diffusion denoising steps.
    """

    DEFAULT_MODEL = "rnd1-base-0910"
    DEFAULT_BASE_URL = "http://localhost:8080"
    DEFAULT_TIMEOUT_SECONDS = 60.0
    DEFAULT_DIFFUSION_STEPS = 256

    def __init__(
        self,
        base_url: str = DEFAULT_BASE_URL,
        timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
        num_diffusion_steps: int = DEFAULT_DIFFUSION_STEPS,
        model: str = DEFAULT_MODEL,
    ) -> None:
        """Initialize RND1 conductor.

        Args:
            base_url: Base URL of the RND1 server.
            timeout_seconds: Request timeout in seconds.
            num_diffusion_steps: Number of diffusion denoising steps.
            model: Model identifier for the served model.
        """
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout_seconds
        self.num_diffusion_steps = num_diffusion_steps
        self.model = model
        self._client = httpx.Client(timeout=timeout_seconds)

    @property
    def backend_type(self) -> str:
        """Return 'local' indicating local model backend."""
        return "local"

    def complete(self, request: ConductorRequest) -> ConductorResponse:
        """Execute a completion request through RND1.

        Uses the OpenAI-compatible /v1/chat/completions endpoint.

        Args:
            request: The completion request containing query and parameters.

        Returns:
            ConductorResponse with generated text and metadata.

        Raises:
            ConductorError: If the completion request fails.
        """
        start_time = time.perf_counter()

        messages = self._build_messages(request)
        payload = self._build_payload(request, messages)

        try:
            response = self._client.post(
                f"{self.base_url}/v1/chat/completions",
                json=payload,
            )
            response.raise_for_status()
            data = response.json()
        except httpx.ConnectError as e:
            raise ConductorError(
                message=f"Failed to connect to RND1 server at {self.base_url}: {e}",
                backend="rnd1",
                retriable=True,
            ) from e
        except httpx.TimeoutException as e:
            raise ConductorError(
                message=f"RND1 request timed out after {self.timeout}s: {e}",
                backend="rnd1",
                retriable=True,
            ) from e
        except httpx.HTTPStatusError as e:
            raise ConductorError(
                message=f"RND1 server error: {e.response.status_code} - {e.response.text}",
                backend="rnd1",
                retriable=e.response.status_code >= 500,
            ) from e

        elapsed_ms = (time.perf_counter() - start_time) * 1000

        return self._parse_response(data, elapsed_ms)

    def health_check(self) -> bool:
        """Check if the RND1 server is healthy.

        Returns:
            True if the server responds to health check.
        """
        try:
            response = self._client.get(
                f"{self.base_url}/health",
                timeout=5.0,
            )
            return response.status_code == 200
        except Exception:
            return False

    def _build_messages(self, request: ConductorRequest) -> list[dict[str, str]]:
        """Build messages array for the API request.

        Args:
            request: The conductor request.

        Returns:
            List of message dictionaries.
        """
        messages: list[dict[str, str]] = []

        if request.context:
            messages.append({
                "role": "system",
                "content": request.context,
            })

        messages.append({
            "role": "user",
            "content": request.query,
        })

        return messages

    def _build_payload(
        self,
        request: ConductorRequest,
        messages: list[dict[str, str]],
    ) -> dict[str, Any]:
        """Build the API request payload.

        Args:
            request: The conductor request.
            messages: The messages array.

        Returns:
            Complete request payload dictionary.
        """
        payload: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "max_tokens": request.max_tokens,
            "temperature": request.temperature,
        }

        # RND1-specific parameters (passed via extra_body for some servers)
        if self.num_diffusion_steps != self.DEFAULT_DIFFUSION_STEPS:
            payload["extra_body"] = {
                "num_diffusion_steps": self.num_diffusion_steps,
            }

        return payload

    def _parse_response(
        self,
        data: dict[str, Any],
        elapsed_ms: float,
    ) -> ConductorResponse:
        """Parse the API response into a ConductorResponse.

        Args:
            data: Raw API response data.
            elapsed_ms: Request latency in milliseconds.

        Returns:
            Parsed ConductorResponse.
        """
        choices = data.get("choices", [])
        if not choices:
            raise ConductorError(
                message="RND1 returned empty choices",
                backend="rnd1",
                retriable=False,
            )

        message = choices[0].get("message", {})
        text = message.get("content", "")

        usage = data.get("usage", {})
        metadata = {
            "finish_reason": choices[0].get("finish_reason"),
            "prompt_tokens": usage.get("prompt_tokens"),
            "completion_tokens": usage.get("completion_tokens"),
            "total_tokens": usage.get("total_tokens"),
            "model": data.get("model", self.model),
        }

        return ConductorResponse(
            text=text,
            model_used=f"rnd1:{self.model}",
            latency_ms=elapsed_ms,
            metadata=metadata,
        )

    def close(self) -> None:
        """Close the HTTP client."""
        self._client.close()

    def __enter__(self) -> "RND1Conductor":
        """Context manager entry."""
        return self

    def __exit__(self, *args: Any) -> None:
        """Context manager exit."""
        self.close()
