"""OpenAI-compatible inference client.

Uses httpx.AsyncClient directly (NOT the openai SDK). This keeps wire payloads
explicit, permits provider-specific extra_body, and simplifies redacted request
auditing.

Requirements: 2.1, 2.2, 2.3, 2.4, 2.5, 2.7, 2.9
"""
from __future__ import annotations

import json
import logging
import os
import time
from typing import Any

import httpx

from services.shared.inference.models import (
    ErrorCategory,
    InferenceResult,
    InferenceTarget,
    StructuredGenerationRequest,
    TokenUsage,
)

logger = logging.getLogger(__name__)

# Status codes that trigger retry
_RETRYABLE_STATUS_CODES = {429, 500, 502, 503, 504}

# Sensitive keys that must never appear in logs
_SENSITIVE_KEYS = frozenset({"authorization", "x-api-key", "api-key"})


def _resolve_auth_secret(secret_ref: str | None) -> str | None:
    """Resolve an authentication secret from environment variables.

    The secret_ref is treated as an environment variable name.
    Returns None if the ref is None or the env var is not set.
    """
    if not secret_ref:
        return None
    return os.environ.get(secret_ref)


def _redact_headers(headers: dict[str, str]) -> dict[str, str]:
    """Return a copy of headers with sensitive values redacted."""
    redacted = {}
    for key, value in headers.items():
        if key.lower() in _SENSITIVE_KEYS:
            redacted[key] = "***REDACTED***"
        else:
            redacted[key] = value
    return redacted


class OpenAICompatibleClient:
    """Client for OpenAI-compatible /v1/chat/completions endpoints.

    Supports:
    - Bearer and configurable authentication via runtime secret resolution
    - Standard response_format.json_schema payloads
    - Configurable vLLM structured_outputs extra-body payloads
    - Explicit JSON-object and prompt-only fallback policies
    - Retry on transient errors (5xx, 429, timeout)
    - Local JSON schema revalidation
    - Metadata capture (request_id, usage, finish_reason, retries)
    """

    def __init__(
        self,
        target: InferenceTarget,
        *,
        http_client: httpx.AsyncClient | None = None,
    ) -> None:
        self._target = target
        self._owns_client = http_client is None
        self._http = http_client or httpx.AsyncClient(timeout=target.timeout_seconds)

    @property
    def target(self) -> InferenceTarget:
        """The resolved inference target."""
        return self._target

    async def generate(
        self,
        request: StructuredGenerationRequest,
    ) -> InferenceResult:
        """Send a structured generation request and return the result.

        Chooses the strongest structured-output mode the target supports:
        1. json_schema — sends the actual schema with strict mode
        2. json_object — only if target explicitly declares json_object capability
        3. prompt_only — fallback when no structured output is available
        """
        structured_mode = self._choose_structured_mode(request)
        headers = self._build_headers()
        body = self._build_body(request, structured_mode)
        url = f"{self._target.base_url.rstrip('/')}/v1/chat/completions"

        retries = 0
        max_retries = self._target.max_retries
        start_time = time.monotonic()
        last_error: str | None = None
        last_error_category: str | None = None

        while True:
            try:
                response = await self._http.post(
                    url,
                    json=body,
                    headers=headers,
                    timeout=request.timeout_seconds,
                )
            except httpx.TimeoutException:
                last_error = "Request timed out"
                last_error_category = ErrorCategory.TIMEOUT
                if retries < max_retries:
                    retries += 1
                    logger.warning(
                        "Timeout on attempt %d/%d to %s",
                        retries,
                        max_retries,
                        url,
                    )
                    continue
                return self._error_result(
                    last_error,
                    last_error_category,
                    retries,
                    start_time,
                    structured_mode,
                )
            except httpx.ConnectError as exc:
                last_error = f"Connection error: {exc}"
                last_error_category = ErrorCategory.CONNECTION_ERROR
                if retries < max_retries:
                    retries += 1
                    logger.warning(
                        "Connection error on attempt %d/%d to %s: %s",
                        retries,
                        max_retries,
                        url,
                        exc,
                    )
                    continue
                return self._error_result(
                    last_error,
                    last_error_category,
                    retries,
                    start_time,
                    structured_mode,
                )

            # Handle retryable HTTP status codes
            if response.status_code in _RETRYABLE_STATUS_CODES:
                if response.status_code == 429:
                    last_error_category = ErrorCategory.RATE_LIMIT
                    last_error = "Rate limited (429)"
                else:
                    last_error_category = ErrorCategory.SERVER_ERROR
                    last_error = f"Server error ({response.status_code})"

                if retries < max_retries:
                    retries += 1
                    logger.warning(
                        "HTTP %d on attempt %d/%d to %s",
                        response.status_code,
                        retries,
                        max_retries,
                        url,
                    )
                    continue
                return self._error_result(
                    last_error,
                    last_error_category,
                    retries,
                    start_time,
                    structured_mode,
                )

            # Handle non-retryable errors
            if response.status_code == 401 or response.status_code == 403:
                return self._error_result(
                    f"Authentication failed ({response.status_code})",
                    ErrorCategory.AUTHENTICATION,
                    retries,
                    start_time,
                    structured_mode,
                )

            if response.status_code >= 400:
                return self._error_result(
                    f"Client error ({response.status_code})",
                    ErrorCategory.INVALID_RESPONSE,
                    retries,
                    start_time,
                    structured_mode,
                )

            # Success path
            break

        latency_ms = int((time.monotonic() - start_time) * 1000)
        return self._parse_response(
            response,
            request,
            structured_mode,
            retries,
            latency_ms,
        )

    def _choose_structured_mode(
        self, request: StructuredGenerationRequest
    ) -> str:
        """Choose the strongest structured-output mode available."""
        caps = self._target.capabilities

        if request.json_schema and caps.json_schema:
            return "json_schema"
        if request.json_schema and caps.json_object:
            return "json_object"
        if request.json_schema:
            # Schema requested but endpoint supports neither — prompt-only fallback
            return "prompt_only"
        return "none"

    def _build_headers(self) -> dict[str, str]:
        """Build request headers including authentication."""
        headers: dict[str, str] = {
            "Content-Type": "application/json",
        }

        # Add extra headers from target config
        headers.update(self._target.extra_headers)

        # Resolve auth secret
        secret = _resolve_auth_secret(self._target.auth_secret_ref)
        if secret:
            scheme = self._target.auth_scheme.lower()
            if scheme == "bearer":
                headers["Authorization"] = f"Bearer {secret}"
            else:
                # Custom auth scheme (e.g., "X-API-Key: <value>")
                headers[self._target.auth_scheme] = secret

        return headers

    def _build_body(
        self,
        request: StructuredGenerationRequest,
        structured_mode: str,
    ) -> dict[str, Any]:
        """Build the request body for /v1/chat/completions."""
        messages = [
            {"role": msg.role, "content": msg.content}
            for msg in request.messages
        ]

        body: dict[str, Any] = {
            "model": self._target.model,
            "messages": messages,
            "temperature": request.temperature,
            "max_tokens": request.max_output_tokens,
        }

        # Add seed if the target supports it
        if self._target.capabilities.seed and request.seed is not None:
            body["seed"] = request.seed

        # Add response_format based on chosen mode
        if structured_mode == "json_schema" and request.json_schema:
            body["response_format"] = {
                "type": "json_schema",
                "json_schema": {
                    "name": request.json_schema.get("title", "response"),
                    "strict": True,
                    "schema": request.json_schema,
                },
            }
        elif structured_mode == "json_object":
            body["response_format"] = {"type": "json_object"}

        # Add extra_body from target config (vLLM structured_outputs, etc.)
        if self._target.extra_body:
            body.update(self._target.extra_body)

        return body

    def _parse_response(
        self,
        response: httpx.Response,
        request: StructuredGenerationRequest,
        structured_mode: str,
        retries: int,
        latency_ms: int,
    ) -> InferenceResult:
        """Parse a successful response into InferenceResult."""
        # Extract request ID from response headers
        request_id = response.headers.get("x-request-id")

        try:
            data = response.json()
        except (json.JSONDecodeError, ValueError):
            return InferenceResult(
                content="",
                endpoint_id=self._target.endpoint_id,
                deployment_id=self._target.deployment_id,
                model=self._target.model,
                protocol=self._target.protocol,
                structured_mode=structured_mode,  # type: ignore[arg-type]
                latency_ms=latency_ms,
                retries=retries,
                request_id=request_id,
                error="Invalid JSON in response body",
                error_category=ErrorCategory.INVALID_RESPONSE,
            )

        # Extract content from choices
        choices = data.get("choices", [])
        if not choices:
            return InferenceResult(
                content="",
                endpoint_id=self._target.endpoint_id,
                deployment_id=self._target.deployment_id,
                model=self._target.model,
                protocol=self._target.protocol,
                structured_mode=structured_mode,  # type: ignore[arg-type]
                latency_ms=latency_ms,
                retries=retries,
                request_id=request_id,
                error="Empty choices in response",
                error_category=ErrorCategory.INVALID_RESPONSE,
            )

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

        # Extract usage metadata
        usage_data = data.get("usage", {})
        input_tokens = usage_data.get("prompt_tokens")
        output_tokens = usage_data.get("completion_tokens")
        total_tokens = usage_data.get("total_tokens")

        # Parse JSON content if structured output was requested
        parsed: dict[str, Any] | None = None
        schema_valid: bool | None = None

        if structured_mode in ("json_schema", "json_object") and content:
            try:
                parsed = json.loads(content)
            except (json.JSONDecodeError, ValueError):
                # Content is not valid JSON
                schema_valid = False

        # Validate parsed JSON against supplied schema
        if parsed is not None and request.json_schema:
            schema_valid = self._validate_schema(parsed, request.json_schema)

        return InferenceResult(
            content=content,
            parsed=parsed,
            endpoint_id=self._target.endpoint_id,
            deployment_id=self._target.deployment_id,
            model=self._target.model,
            protocol=self._target.protocol,
            structured_mode=structured_mode,  # type: ignore[arg-type]
            latency_ms=latency_ms,
            usage=TokenUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
            ),
            request_id=request_id,
            finish_reason=finish_reason,
            retries=retries,
            schema_valid=schema_valid,
        )

    def _validate_schema(
        self, data: dict[str, Any], schema: dict[str, Any]
    ) -> bool:
        """Revalidate parsed JSON locally against the supplied schema.

        Returns True if valid, False otherwise.
        """
        try:
            import jsonschema

            jsonschema.validate(instance=data, schema=schema)
            return True
        except Exception:
            logger.debug("Schema validation failed for response data")
            return False

    def _error_result(
        self,
        error: str,
        error_category: str,
        retries: int,
        start_time: float,
        structured_mode: str,
    ) -> InferenceResult:
        """Build an error InferenceResult."""
        latency_ms = int((time.monotonic() - start_time) * 1000)
        return InferenceResult(
            content="",
            endpoint_id=self._target.endpoint_id,
            deployment_id=self._target.deployment_id,
            model=self._target.model,
            protocol=self._target.protocol,
            structured_mode=structured_mode,  # type: ignore[arg-type]
            latency_ms=latency_ms,
            retries=retries,
            error=error,
            error_category=error_category,
        )

    async def close(self) -> None:
        """Close the underlying HTTP client if we own it."""
        if self._owns_client:
            await self._http.aclose()

    def __repr__(self) -> str:
        return (
            f"OpenAICompatibleClient("
            f"model={self._target.model!r}, "
            f"base_url={self._target.base_url!r})"
        )