"""Alibaba Cloud Model Studio OpenAI-compatible chat adapter."""

from __future__ import annotations

import json
from collections.abc import AsyncIterator, Mapping, Sequence
from time import perf_counter
from typing import Any

import httpx

from app.adapters.bailian._base import (
    BailianHttpAdapter,
    extract_request_id,
    invalid_request,
    invalid_response,
    parse_usage,
    response_model,
    sanitized_error,
)
from app.ports.model_providers import (
    ChatCompletionResult,
    ChatMessage,
    ChatStreamEvent,
    ModelProviderError,
    ProviderErrorKind,
    ProviderUsage,
)

_ALLOWED_ROLES = frozenset({"system", "user", "assistant"})


class BailianChatAdapter(BailianHttpAdapter):
    """Call chat completions with thinking and web search forcibly disabled."""

    def __init__(
        self,
        *,
        api_key: str,
        base_url: str,
        model: str = "deepseek-v4-flash",
        http_client: httpx.AsyncClient | None = None,
        timeout_seconds: float = 60.0,
        max_retries: int = 0,
        retry_base_seconds: float = 0.5,
    ) -> None:
        if not model or model != model.strip():
            raise invalid_request("chat.configuration", "invalid_model")
        super().__init__(
            api_key=api_key,
            base_url=base_url,
            expected_path="/compatible-mode/v1",
            http_client=http_client,
            timeout_seconds=timeout_seconds,
            max_retries=max_retries,
            retry_base_seconds=retry_base_seconds,
        )
        self._model = model

    async def complete(
        self,
        messages: Sequence[ChatMessage],
        *,
        max_tokens: int,
    ) -> ChatCompletionResult:
        operation = "chat.complete"
        validated_messages = self._validate_messages(
            messages,
            max_tokens=max_tokens,
            operation=operation,
        )
        payload = self._payload(validated_messages, max_tokens=max_tokens, stream=False)
        sensitive_values = tuple(message.content for message in validated_messages)
        response = await self._post_json(
            operation=operation,
            path="chat/completions",
            payload=payload,
            sensitive_values=sensitive_values,
        )
        content, finish_reason = self._parse_completion(
            response.body,
            operation=operation,
        )
        return ChatCompletionResult(
            content=content,
            finish_reason=finish_reason,
            model=response_model(
                response.body,
                self._model,
                sensitive_values=(*sensitive_values, self._api_key),
            ),
            request_id=extract_request_id(
                response.body,
                sensitive_values=(*sensitive_values, self._api_key),
            ),
            usage=parse_usage(response.body.get("usage")),
            elapsed_ms=response.elapsed_ms,
        )

    async def stream(
        self,
        messages: Sequence[ChatMessage],
        *,
        max_tokens: int,
    ) -> AsyncIterator[ChatStreamEvent]:
        operation = "chat.stream"
        validated_messages = self._validate_messages(
            messages,
            max_tokens=max_tokens,
            operation=operation,
        )
        payload = self._payload(validated_messages, max_tokens=max_tokens, stream=True)
        sensitive_values = tuple(message.content for message in validated_messages)
        started = perf_counter()
        for attempt in range(self._max_retries + 1):
            emitted = False
            try:
                async with self._client.stream(
                    "POST",
                    self._url("chat/completions"),
                    headers=self._headers(),
                    json=payload,
                ) as response:
                    if response.status_code >= 400:
                        await response.aread()
                        try:
                            self._raise_http_error(
                                operation=operation,
                                response=response,
                                sensitive_values=(*sensitive_values, self._api_key),
                            )
                        except ModelProviderError as error:
                            if await self._maybe_retry(
                                error,
                                attempt=attempt,
                                response=response,
                            ):
                                continue
                            raise

                    async for line in response.aiter_lines():
                        if not line or line.startswith(":"):
                            continue
                        if not line.startswith("data:"):
                            raise invalid_response(operation, "invalid_sse_event")
                        raw_data = line[5:].strip()
                        if raw_data == "[DONE]":
                            return
                        event_body = self._decode_stream_event(raw_data, operation=operation)
                        event = self._parse_stream_event(
                            event_body,
                            operation=operation,
                            sensitive_values=sensitive_values,
                            elapsed_ms=(perf_counter() - started) * 1000,
                        )
                        emitted = True
                        yield event
                    return
            except ModelProviderError as error:
                if not emitted and await self._maybe_retry(
                    error,
                    attempt=attempt,
                    response=None,
                ):
                    continue
                raise
            except httpx.TimeoutException:
                timeout_error = sanitized_error(
                    operation=operation,
                    kind=ProviderErrorKind.TIMEOUT,
                    provider_code="request_timeout",
                    retryable=True,
                )
                if not emitted and await self._maybe_retry(
                    timeout_error,
                    attempt=attempt,
                    response=None,
                ):
                    continue
                raise timeout_error from None
            except httpx.HTTPError:
                transport_error = sanitized_error(
                    operation=operation,
                    kind=ProviderErrorKind.TRANSPORT,
                    provider_code="transport_error",
                    retryable=True,
                )
                if not emitted and await self._maybe_retry(
                    transport_error,
                    attempt=attempt,
                    response=None,
                ):
                    continue
                raise transport_error from None

        raise AssertionError("bounded chat retry loop exhausted unexpectedly")

    def _validate_messages(
        self,
        messages: object,
        *,
        max_tokens: int,
        operation: str,
    ) -> tuple[ChatMessage, ...]:
        if isinstance(messages, (str, bytes)) or not isinstance(messages, Sequence):
            raise invalid_request(operation, "invalid_message_collection")
        validated = tuple(messages)
        if not validated:
            raise invalid_request(operation, "empty_messages")
        if isinstance(max_tokens, bool) or not isinstance(max_tokens, int) or max_tokens <= 0:
            raise invalid_request(operation, "invalid_max_tokens")
        for message in validated:
            if (
                not isinstance(message, ChatMessage)
                or message.role not in _ALLOWED_ROLES
                or not isinstance(message.content, str)
                or not message.content
            ):
                raise invalid_request(operation, "invalid_message")
        return validated

    def _payload(
        self,
        messages: tuple[ChatMessage, ...],
        *,
        max_tokens: int,
        stream: bool,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "model": self._model,
            "messages": [
                {"role": message.role, "content": message.content} for message in messages
            ],
            "max_tokens": max_tokens,
            "stream": stream,
            "enable_thinking": False,
            "enable_search": False,
        }
        if stream:
            payload["stream_options"] = {"include_usage": True}
        return payload

    def _parse_completion(
        self,
        body: Mapping[str, Any],
        *,
        operation: str,
    ) -> tuple[str, str | None]:
        choices = body.get("choices")
        if not isinstance(choices, list) or len(choices) != 1:
            raise invalid_response(operation, "invalid_choices")
        choice = choices[0]
        if not isinstance(choice, Mapping):
            raise invalid_response(operation, "invalid_choice")
        message = choice.get("message")
        if not isinstance(message, Mapping):
            raise invalid_response(operation, "invalid_message")
        content = message.get("content")
        if not isinstance(content, str):
            raise invalid_response(operation, "invalid_content")
        finish_reason = choice.get("finish_reason")
        if finish_reason is not None and not isinstance(finish_reason, str):
            raise invalid_response(operation, "invalid_finish_reason")
        return content, finish_reason

    def _decode_stream_event(
        self,
        raw_data: str,
        *,
        operation: str,
    ) -> Mapping[str, Any]:
        try:
            decoded = json.loads(raw_data)
        except (TypeError, ValueError):
            raise invalid_response(operation, "invalid_stream_json") from None
        if not isinstance(decoded, Mapping):
            raise invalid_response(operation, "invalid_stream_object")
        return decoded

    def _parse_stream_event(
        self,
        body: Mapping[str, Any],
        *,
        operation: str,
        sensitive_values: tuple[str, ...],
        elapsed_ms: float,
    ) -> ChatStreamEvent:
        choices = body.get("choices")
        delta = ""
        finish_reason: str | None = None
        if choices not in (None, []):
            if not isinstance(choices, list) or len(choices) != 1:
                raise invalid_response(operation, "invalid_stream_choices")
            choice = choices[0]
            if not isinstance(choice, Mapping):
                raise invalid_response(operation, "invalid_stream_choice")
            raw_delta = choice.get("delta")
            if not isinstance(raw_delta, Mapping):
                raise invalid_response(operation, "invalid_stream_delta")
            content = raw_delta.get("content", "")
            if not isinstance(content, str):
                raise invalid_response(operation, "invalid_stream_content")
            delta = content
            raw_finish_reason = choice.get("finish_reason")
            if raw_finish_reason is not None and not isinstance(raw_finish_reason, str):
                raise invalid_response(operation, "invalid_stream_finish_reason")
            finish_reason = raw_finish_reason

        usage = parse_usage(body.get("usage"))
        if choices in (None, []) and usage == ProviderUsage():
            raise invalid_response(operation, "empty_stream_event")
        return ChatStreamEvent(
            delta=delta,
            finish_reason=finish_reason,
            model=response_model(
                body,
                self._model,
                sensitive_values=(*sensitive_values, self._api_key),
            ),
            request_id=extract_request_id(
                body,
                sensitive_values=(*sensitive_values, self._api_key),
            ),
            usage=usage,
            elapsed_ms=elapsed_ms,
        )