"""ModelRelay API client for LLM reasoning calls."""

import json
import os
import sys
import time
from dataclasses import dataclass, field
from typing import Any, Iterator, Literal, Optional

import httpx

MODELRELAY_BASE_URL = "https://api.modelrelay.ai/api/v1"
DEFAULT_MODEL = "gemini-3-flash-preview"
_MODEL_CACHE_TTL_SECONDS = 300
ROOT_MODEL_ENV = "RECALL_RLM_ROOT_MODEL"
SUBCALL_MODEL_ENV = "RECALL_RLM_SUBCALL_MODEL"


@dataclass(frozen=True)
class ModelResolutionPolicy:
    """Policy for resolving model names."""

    allow_env: bool = False

    @staticmethod
    def env_allowed() -> "ModelResolutionPolicy":
        return ModelResolutionPolicy(allow_env=True)


DEFAULT_MODEL_POLICY = ModelResolutionPolicy()


def _resolve_policy(policy: "ModelResolutionPolicy | None") -> ModelResolutionPolicy:
    return policy or DEFAULT_MODEL_POLICY


# MARK: - Response Dataclasses (Type Strictness)


@dataclass
class ModelInfo:
    """Typed model metadata from the /models endpoint."""

    model_id: str
    name: str | None = None
    capabilities: list[str] = field(default_factory=list)
    context_window: int | None = None
    max_output_tokens: int | None = None

    @staticmethod
    def from_dict(data: dict[str, Any]) -> "ModelInfo":
        """Parse a model entry from the API response."""
        model_id = data.get("model_id") or data.get("id") or data.get("name") or ""
        capabilities = data.get("capabilities") or data.get("capability") or data.get("features") or []
        if isinstance(capabilities, str):
            capabilities = [capabilities]
        return ModelInfo(
            model_id=str(model_id),
            name=data.get("name"),
            capabilities=[str(c) for c in capabilities],
            context_window=_safe_int(data.get("context_window")),
            max_output_tokens=_safe_int(data.get("max_output_tokens")),
        )


@dataclass
class TextContent:
    """Text content part from a response."""

    type: Literal["text"]
    text: str


@dataclass
class MessageOutput:
    """Assistant message from the response."""

    type: Literal["message"]
    role: str
    content: list[TextContent]


@dataclass
class ResponseUsage:
    """Token usage from the API response."""

    input_tokens: int = 0
    output_tokens: int = 0

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens


@dataclass
class BatchUsage:
    """Aggregated usage stats from a batch response."""

    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_requests: int = 0
    successful_requests: int = 0
    failed_requests: int = 0

    @staticmethod
    def from_dict(data: dict[str, Any]) -> "BatchUsage":
        return BatchUsage(
            total_input_tokens=_safe_int(data.get("total_input_tokens")) or 0,
            total_output_tokens=_safe_int(data.get("total_output_tokens")) or 0,
            total_requests=_safe_int(data.get("total_requests")) or 0,
            successful_requests=_safe_int(data.get("successful_requests")) or 0,
            failed_requests=_safe_int(data.get("failed_requests")) or 0,
        )


@dataclass
class BatchError:
    """Error details for a batch item."""

    status: int | None = None
    message: str | None = None
    detail: str | None = None
    code: str | None = None

    @staticmethod
    def from_dict(data: dict[str, Any]) -> "BatchError":
        return BatchError(
            status=_safe_int(data.get("status")),
            message=data.get("message"),
            detail=data.get("detail"),
            code=data.get("code"),
        )


@dataclass
class BatchResult:
    """Result for a single batch item."""

    id: str
    status: Literal["success", "error"]
    response: dict[str, Any] | None = None
    error: BatchError | None = None

    @staticmethod
    def from_dict(data: dict[str, Any]) -> "BatchResult":
        if not isinstance(data, dict):
            _invalid_batch_payload("invalid batch result payload", data)
        item_id_raw = data.get("id")
        if not item_id_raw:
            _invalid_batch_payload("missing batch item id", data)
        item_id = str(item_id_raw)
        status = data.get("status")
        if status not in ("success", "error"):
            _invalid_batch_payload(f"invalid batch status '{status}' for item {item_id}", data)
        response = data.get("response") if status == "success" else None
        if status == "success" and not isinstance(response, dict):
            _invalid_batch_payload(f"missing response for batch item {item_id}", data)
        error = data.get("error") if status == "error" else None
        if status == "error" and not isinstance(error, dict):
            _invalid_batch_payload(f"missing error for batch item {item_id}", data)
        error_obj = BatchError.from_dict(error) if status == "error" else None
        return BatchResult(
            id=item_id,
            status=status,
            response=response,
            error=error_obj,
        )


@dataclass
class BatchResponse:
    """Typed batch responses payload."""

    id: str
    results: list[BatchResult]
    usage: BatchUsage

    @staticmethod
    def from_dict(data: dict[str, Any]) -> "BatchResponse":
        if not isinstance(data, dict):
            _invalid_batch_payload("invalid batch response payload", data)
        batch_id_raw = data.get("id")
        if not batch_id_raw:
            _invalid_batch_payload("missing batch id in response", data)
        batch_id = str(batch_id_raw)
        results_raw = data.get("results")
        if not isinstance(results_raw, list):
            _invalid_batch_payload("missing results in batch response", data)
        usage_raw = data.get("usage")
        if not isinstance(usage_raw, dict):
            _invalid_batch_payload("missing usage in batch response", data)
        results = [BatchResult.from_dict(item) for item in results_raw]
        return BatchResponse(
            id=batch_id,
            results=results,
            usage=BatchUsage.from_dict(usage_raw),
        )


def _safe_int(value: Any) -> int | None:
    """Safely convert a value to int, returning None on failure."""
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def resolve_model(explicit: Optional[str], *, policy: ModelResolutionPolicy | None = None) -> str:
    """Resolve the model name, optionally applying env overrides."""
    policy = _resolve_policy(policy)
    if explicit and explicit.strip():
        return explicit.strip()
    if policy.allow_env:
        env_model = os.environ.get("RECALL_RLM_MODEL") or os.environ.get("MODELRELAY_MODEL")
        if env_model and env_model.strip():
            return env_model.strip()
    return DEFAULT_MODEL


def resolve_root_model(explicit: Optional[str], *, policy: ModelResolutionPolicy | None = None) -> str:
    """Resolve the root model name, optionally applying env overrides."""
    policy = _resolve_policy(policy)
    if explicit and explicit.strip():
        return explicit.strip()
    if policy.allow_env:
        env_model = os.environ.get(ROOT_MODEL_ENV)
        if env_model and env_model.strip():
            return env_model.strip()
    return resolve_model(None, policy=policy)


def resolve_subcall_model(
    explicit: Optional[str],
    root_fallback: Optional[str] = None,
    *,
    policy: ModelResolutionPolicy | None = None,
) -> str:
    """Resolve the subcall model name, optionally applying env overrides."""
    policy = _resolve_policy(policy)
    if explicit and explicit.strip():
        return explicit.strip()
    if policy.allow_env:
        env_model = os.environ.get(SUBCALL_MODEL_ENV)
        if env_model and env_model.strip():
            return env_model.strip()
    if root_fallback and root_fallback.strip():
        return root_fallback.strip()
    return resolve_root_model(None, policy=policy)


class ModelRelayError(Exception):
    """Raised when ModelRelay API returns an error."""

    def __init__(self, status_code: int, message: str, response_body: str = ""):
        self.status_code = status_code
        self.message = message
        self.response_body = response_body
        super().__init__(f"ModelRelay API error {status_code}: {message}")


def _truncate_payload(payload: Any, limit: int = 2000) -> str:
    """Serialize payload for error reporting, truncating to a safe size."""
    try:
        text = json.dumps(payload, ensure_ascii=True)
    except (TypeError, ValueError):
        text = str(payload)
    if len(text) > limit:
        return text[:limit] + "...(truncated)"
    return text


def _invalid_batch_payload(message: str, payload: Any) -> None:
    raise ModelRelayError(0, message, _truncate_payload(payload))


def get_api_key() -> str:
    """Get ModelRelay API key from environment."""
    key = os.environ.get("MODELRELAY_API_KEY")
    if not key:
        raise ValueError("MODELRELAY_API_KEY environment variable not set")
    return key


def _build_input(messages: list[dict]) -> list[dict]:
    """Convert OpenAI-style messages to ModelRelay input format."""
    result = []
    for msg in messages:
        role = msg["role"]
        content = msg["content"]

        # ModelRelay InputItem format:
        # {"type": "message", "role": "user", "content": [{"type": "text", "text": "..."}]}
        if isinstance(content, list):
            content_items = content
        else:
            content_items = [{"type": "text", "text": content}]

        result.append({
            "type": "message",
            "role": role,
            "content": content_items,
        })

    return result


def _extract_text_from_response(data: dict[str, Any]) -> str:
    """Extract assistant text content from a ModelRelay responses payload."""
    output = data.get("output", [])
    text_parts = []
    for item in output:
        if item.get("type") == "message" and item.get("role") == "assistant":
            for part in item.get("content", []):
                if part.get("type") == "text":
                    text_parts.append(part.get("text", ""))
    return "".join(text_parts)


def _log_event(message: str, **metadata: object) -> None:
    payload = {"message": message, **metadata}
    print(json.dumps(payload, ensure_ascii=True), file=sys.stderr, flush=True)


class TokenUsage:
    """Token usage from an API call."""
    def __init__(self, input_tokens: int = 0, output_tokens: int = 0):
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens
        self.total_tokens = input_tokens + output_tokens


# MARK: - ModelRelayClient (Injectable, Instance-based)


class ModelRelayClient:
    """Injectable ModelRelay API client with instance-based caching.

    This class encapsulates the HTTP client and model cache as instance state,
    enabling dependency injection for testing and avoiding global mutable state.

    Usage:
        # Create a client with explicit API key
        client = ModelRelayClient(api_key="your-key")
        models = client.list_models()

        # Or let it read from environment
        client = ModelRelayClient()
        result = client.responses_create([{"role": "user", "content": "Hello"}])
    """

    def __init__(
        self,
        api_key: str | None = None,
        customer_token: str | None = None,
        base_url: str = MODELRELAY_BASE_URL,
        cache_ttl_seconds: float = _MODEL_CACHE_TTL_SECONDS,
    ):
        """Initialize the client.

        Args:
            api_key: ModelRelay API key. If None, reads from MODELRELAY_API_KEY env var at init time.
            customer_token: ModelRelay customer token for PAYGO billing. Takes precedence over api_key.
                           If None, reads from MODELRELAY_CUSTOMER_TOKEN env var at init time.
            base_url: Base URL for the ModelRelay API.
            cache_ttl_seconds: TTL for model list cache.

        Environment variables are captured once at initialization, not on each API call.
        """
        # Capture env vars at init time to avoid mid-execution env reads
        self._api_key = api_key or os.environ.get("MODELRELAY_API_KEY")
        self._customer_token = customer_token or os.environ.get("MODELRELAY_CUSTOMER_TOKEN")
        self._base_url = base_url
        self._cache_ttl = cache_ttl_seconds
        self._model_cache: list[dict[str, Any]] | None = None
        self._model_cache_at: float | None = None

    def _get_customer_token(self) -> str | None:
        """Get customer token captured at init time."""
        return self._customer_token

    def _get_api_key(self) -> str:
        """Get API key captured at init time."""
        if not self._api_key:
            raise ValueError("MODELRELAY_API_KEY not provided and not found in environment at init time")
        return self._api_key

    def _get_auth_headers(self) -> dict[str, str]:
        """Get authentication headers. Customer token takes precedence over API key."""
        customer_token = self._get_customer_token()
        if customer_token:
            return {"Authorization": f"Bearer {customer_token}"}
        return {"X-ModelRelay-Api-Key": self._get_api_key()}

    def list_models(self) -> list[dict[str, Any]]:
        """Fetch available models from ModelRelay with caching.

        Returns raw dicts for backward compatibility.
        Use list_models_typed() for strongly-typed ModelInfo objects.
        """
        now = time.time()
        if self._model_cache is not None and self._model_cache_at is not None:
            if now - self._model_cache_at < self._cache_ttl:
                return self._model_cache

        headers = {
            "Accept": "application/json",
            "X-ModelRelay-Client": "rcl-rlm-python",
        }
        try:
            auth_headers = self._get_auth_headers()
            headers.update(auth_headers)
        except ValueError:
            pass  # Allow unauthenticated model list fetch

        try:
            with httpx.Client(timeout=20.0) as client:
                response = client.get(f"{self._base_url}/models", headers=headers)
                response.raise_for_status()
                payload = response.json()
        except httpx.TimeoutException as exc:
            _log_event("modelrelay_models_fetch_timeout", error=str(exc))
            return []
        except httpx.HTTPStatusError as exc:
            if exc.response.status_code in (401, 403):
                # Auth failures are configuration errors - fail fast
                _log_event("modelrelay_models_auth_failed", status_code=exc.response.status_code)
                raise ModelRelayError(
                    exc.response.status_code,
                    f"Authentication failed (status {exc.response.status_code}). Check your credentials.",
                    exc.response.text,
                )
            # Other HTTP errors (5xx, etc.) - graceful degradation
            _log_event("modelrelay_models_fetch_failed", status_code=exc.response.status_code, error=str(exc))
            return []
        except Exception as exc:  # noqa: BLE001 - network boundary
            _log_event("modelrelay_models_fetch_failed", error=str(exc))
            return []

        models = payload.get("models")
        if isinstance(models, list):
            self._model_cache = models
            self._model_cache_at = now
            return models

        _log_event("modelrelay_models_parse_failed", payload_type=str(type(payload)))
        return []

    def list_models_typed(self) -> list[ModelInfo]:
        """Fetch available models as typed ModelInfo objects."""
        return [ModelInfo.from_dict(entry) for entry in self.list_models()]

    def get_model_info(self, model_id: str) -> dict[str, Any] | None:
        """Return model metadata for a given model id (raw dict)."""
        if not model_id:
            return None
        for entry in self.list_models():
            if entry.get("model_id") == model_id:
                return entry
        return None

    def get_model_info_typed(self, model_id: str) -> ModelInfo | None:
        """Return typed model metadata for a given model id."""
        info = self.get_model_info(model_id)
        if info is None:
            return None
        return ModelInfo.from_dict(info)

    def get_model_context_window(self, model_id: str) -> int | None:
        """Return the max context window (tokens) for a model if known."""
        info = self.get_model_info_typed(model_id)
        return info.context_window if info else None

    def get_model_max_output_tokens(self, model_id: str) -> int | None:
        """Return the max output tokens for a model if known."""
        info = self.get_model_info_typed(model_id)
        return info.max_output_tokens if info else None

    def clear_cache(self) -> None:
        """Clear the model list cache."""
        self._model_cache = None
        self._model_cache_at = None

    def _make_request_with_retry(
        self,
        endpoint: str,
        request_body: dict[str, Any],
        max_retries: int = 3,
        timeout: float = 60.0,
    ) -> dict[str, Any]:
        """Make HTTP POST with exponential backoff retry for 5xx errors.

        Args:
            endpoint: API endpoint path (e.g., "/responses")
            request_body: JSON request body
            max_retries: Maximum retry attempts for 5xx errors
            timeout: Request timeout in seconds

        Returns:
            Parsed JSON response

        Raises:
            ModelRelayError: On 4xx/5xx errors after retries exhausted
        """
        auth_headers = self._get_auth_headers()
        headers = {
            **auth_headers,
            "X-ModelRelay-Client": "rcl-rlm-python",
            "Content-Type": "application/json",
            "Accept": "application/json",
        }

        last_error: ModelRelayError | None = None
        for attempt in range(max_retries):
            with httpx.Client(timeout=timeout) as client:
                response = client.post(
                    f"{self._base_url}{endpoint}",
                    headers=headers,
                    json=request_body,
                )
                if response.status_code >= 500:
                    error_text = response.text
                    last_error = ModelRelayError(response.status_code, error_text, error_text)
                    if attempt < max_retries - 1:
                        wait_time = 2**attempt
                        _log_event(
                            "modelrelay_retry",
                            status_code=response.status_code,
                            wait_seconds=wait_time,
                            attempt=attempt + 1,
                            max_retries=max_retries,
                        )
                        time.sleep(wait_time)
                        continue
                    _log_event(
                        "modelrelay_error",
                        status_code=response.status_code,
                        error_text=error_text,
                    )
                    raise last_error
                if response.status_code >= 400:
                    error_text = response.text
                    _log_event(
                        "modelrelay_error",
                        status_code=response.status_code,
                        error_text=error_text,
                    )
                    raise ModelRelayError(response.status_code, error_text, error_text)
                return response.json()

        # Should not reach here, but satisfy type checker
        if last_error:
            raise last_error
        raise ModelRelayError(0, "Request failed with no error details", "")

    def responses_create(
        self,
        messages: list[dict],
        model: Optional[str] = None,
        temperature: float = 0.0,
        max_tokens: int = 4096,
        max_retries: int = 3,
        return_usage: bool = False,
    ) -> str | tuple[str, TokenUsage]:
        """Make a responses request to ModelRelay with retry logic for transient errors."""
        model = resolve_model(model)
        request_body = {
            "model": model,
            "input": _build_input(messages),
            "max_output_tokens": max_tokens,
            "temperature": temperature,
        }

        data = self._make_request_with_retry("/responses", request_body, max_retries)
        result_text = _extract_text_from_response(data)

        if return_usage:
            usage = data.get("usage", {})
            token_usage = TokenUsage(
                input_tokens=usage.get("input_tokens", 0),
                output_tokens=usage.get("output_tokens", 0),
            )
            return result_text, token_usage

        return result_text

    def batch_responses(
        self,
        requests: list[dict[str, Any]],
        options: dict[str, Any] | None = None,
        max_retries: int = 3,
    ) -> dict[str, Any]:
        """Make a batch responses request to ModelRelay with retry logic for transient errors."""
        if not requests:
            raise ValueError("batch responses requests must not be empty")

        request_body: dict[str, Any] = {"requests": requests}
        if options:
            request_body["options"] = options

        return self._make_request_with_retry("/responses/batch", request_body, max_retries)

    def batch_responses_typed(
        self,
        requests: list[dict[str, Any]],
        options: dict[str, Any] | None = None,
        max_retries: int = 3,
    ) -> BatchResponse:
        """Make a batch responses request to ModelRelay and parse typed output."""
        payload = self.batch_responses(
            requests=requests,
            options=options,
            max_retries=max_retries,
        )
        return BatchResponse.from_dict(payload)

    def responses_stream(
        self,
        messages: list[dict],
        model: Optional[str] = None,
        temperature: float = 0.0,
        max_tokens: int = 4096,
    ) -> Iterator[str]:
        """Stream a responses request to ModelRelay."""
        model = resolve_model(model)
        auth_headers = self._get_auth_headers()

        request_body = {
            "model": model,
            "input": _build_input(messages),
            "max_output_tokens": max_tokens,
            "temperature": temperature,
        }

        with httpx.Client(timeout=120.0) as client:
            with client.stream(
                "POST",
                f"{self._base_url}/responses",
                headers={
                    **auth_headers,
                    "X-ModelRelay-Client": "rcl-rlm-python",
                    "Content-Type": "application/json",
                    "Accept": "application/x-ndjson; profile=\"responses-stream/v2\"",
                },
                json=request_body,
            ) as response:
                response.raise_for_status()
                for line in response.iter_lines():
                    if not line:
                        continue
                    try:
                        event = json.loads(line)
                        if event.get("type") == "content_delta":
                            delta = event.get("delta", {})
                            if delta.get("type") == "text":
                                yield delta.get("content", "")
                    except json.JSONDecodeError as exc:
                        _log_event("stream_json_decode_error", line=line[:100] if line else "", error=str(exc))
                        continue


# MARK: - Default Client Singleton (Backward Compatibility)


_default_client: ModelRelayClient | None = None


def _get_default_client() -> ModelRelayClient:
    """Get or create the default singleton client for backward compatibility."""
    global _default_client
    if _default_client is None:
        _default_client = ModelRelayClient()
    return _default_client


# MARK: - Module-Level Functions (Convenience Wrappers)
#
# These use a module-level singleton client for convenience.
# For better testability, create a ModelRelayClient instance directly.


def list_models() -> list[dict[str, Any]]:
    """Fetch available models from ModelRelay.

    Returns raw dicts for backward compatibility.
    Use ModelRelayClient().list_models_typed() for strongly-typed objects.
    """
    return _get_default_client().list_models()


def list_models_typed() -> list[ModelInfo]:
    """Fetch available models as typed ModelInfo objects."""
    return _get_default_client().list_models_typed()


def get_model_info(model_id: str) -> dict[str, Any] | None:
    """Return model metadata for a given model id (raw dict)."""
    return _get_default_client().get_model_info(model_id)


def get_model_info_typed(model_id: str) -> ModelInfo | None:
    """Return typed model metadata for a given model id."""
    return _get_default_client().get_model_info_typed(model_id)


def get_model_context_window(model_id: str) -> int | None:
    """Return the max context window (tokens) for a model if known."""
    return _get_default_client().get_model_context_window(model_id)


def get_model_max_output_tokens(model_id: str) -> int | None:
    """Return the max output tokens for a model if known."""
    return _get_default_client().get_model_max_output_tokens(model_id)


def resolve_vision_model(explicit: Optional[str], *, policy: ModelResolutionPolicy | None = None) -> str:
    """Resolve a vision-capable model name, optionally applying env overrides."""
    policy = _resolve_policy(policy)
    if explicit and explicit.strip():
        return explicit.strip()
    if policy.allow_env:
        env_model = os.environ.get("RECALL_RLM_VISION_MODEL") or os.environ.get("MODELRELAY_VISION_MODEL")
        if env_model and env_model.strip():
            return env_model.strip()

    for entry in list_models():
        capabilities = entry.get("capabilities") or entry.get("capability") or entry.get("features") or []
        if isinstance(capabilities, str):
            capabilities = [capabilities]
        if any(str(value).lower() == "vision" for value in capabilities):
            model_id = entry.get("model_id") or entry.get("id") or entry.get("name")
            if model_id:
                return str(model_id)

    return resolve_model(None, policy=policy)


def responses_create(
    messages: list[dict],
    model: Optional[str] = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
    max_retries: int = 3,
    return_usage: bool = False,
) -> str | tuple[str, TokenUsage]:
    """Make a responses request to ModelRelay with retry logic for transient errors."""
    return _get_default_client().responses_create(
        messages=messages,
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
        max_retries=max_retries,
        return_usage=return_usage,
    )


def responses_stream(
    messages: list[dict],
    model: Optional[str] = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
) -> Iterator[str]:
    """Stream a responses request to ModelRelay."""
    return _get_default_client().responses_stream(
        messages=messages,
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
    )


def batch_responses(
    requests: list[dict[str, Any]],
    options: dict[str, Any] | None = None,
    max_retries: int = 3,
) -> dict[str, Any]:
    """Make a batch responses request to ModelRelay."""
    return _get_default_client().batch_responses(
        requests=requests,
        options=options,
        max_retries=max_retries,
    )


def batch_responses_typed(
    requests: list[dict[str, Any]],
    options: dict[str, Any] | None = None,
    max_retries: int = 3,
) -> BatchResponse:
    """Make a batch responses request to ModelRelay and parse typed output."""
    return _get_default_client().batch_responses_typed(
        requests=requests,
        options=options,
        max_retries=max_retries,
    )
