from __future__ import annotations

"""Azure OpenAI provider implementation."""

import json
import os
from typing import Any, AsyncIterator

import httpx

from agentrspy.core import (
    CompletionRequest,
    CompletionResponse,
    LlmProvider,
    Message,
    ProviderError,
    Role,
    StopReason,
    StreamChunk,
    ToolCall,
    ToolDefinition,
    Usage,
    map_stop_reason,
    parse_sse_chunk,
)


class AzureOpenAiProvider(LlmProvider):
    """Azure OpenAI provider."""

    def __init__(
        self,
        *,
        client: httpx.AsyncClient,
        api_key: str,
        base_url: str,
        model: str,
    ) -> None:
        self.client = client
        self.api_key = api_key
        self.base_url = base_url
        self.model = model

    @classmethod
    def from_env(cls) -> "AzureOpenAiBuilder":
        return AzureOpenAiBuilder(
            api_key=os.getenv("AZURE_OPENAI_KEY"),
            base_url=os.getenv("AZURE_OPENAI_ENDPOINT"),
            model=os.getenv("AZURE_OPENAI_MODEL"),
        )

    @classmethod
    def builder(cls) -> "AzureOpenAiBuilder":
        return AzureOpenAiBuilder()

    def _request_body(self, request: CompletionRequest, stream: bool) -> dict[str, Any]:
        body: dict[str, Any] = {
            "model": request.model or self.model,
            "messages": [_map_message(message) for message in request.messages],
            "temperature": 1.0,
            "max_completion_tokens": request.max_tokens,
            "stream": stream,
        }
        if request.tools is not None:
            body["tools"] = [_map_tool(tool) for tool in request.tools]
        if request.system is not None:
            body["messages"].insert(0, {"role": "system", "content": request.system})
        body.update(request.extra)
        return body

    async def _send(self, request: CompletionRequest, stream: bool) -> httpx.Response:
        try:
            response = await self.client.post(
                f"{self.base_url.rstrip('/')}/chat/completions",
                headers={"Authorization": f"Bearer {self.api_key}"},
                json=self._request_body(request, stream),
                timeout=60.0,
            )
        except httpx.HTTPError as error:
            raise ProviderError.http(str(error)) from error
        if response.status_code >= 400:
            raise ProviderError.api(response.status_code, response.text)
        return response

    async def complete(self, request: CompletionRequest) -> CompletionResponse:
        response = await self._send(request, stream=False)
        try:
            payload = response.json()
        except ValueError as error:
            raise ProviderError.parse(str(error)) from error
        return _map_response(payload)

    async def stream(self, request: CompletionRequest) -> AsyncIterator[StreamChunk]:
        response = await self._send(request, stream=True)

        async def iterator() -> AsyncIterator[StreamChunk]:
            async for line in response.aiter_lines():
                if not line:
                    continue
                parsed = parse_sse_chunk(line.encode("utf-8"))
                if parsed is not None:
                    yield parsed

        return iterator()

    def name(self) -> str:
        return "azureopenai"


class AzureOpenAiBuilder:
    """Builder for `AzureOpenAiProvider`."""

    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        model: str | None = None,
        client: httpx.AsyncClient | None = None,
    ) -> None:
        self._api_key = api_key
        self._base_url = base_url
        self._model = model
        self._client = client or httpx.AsyncClient()

    def api_key(self, api_key: str) -> "AzureOpenAiBuilder":
        self._api_key = api_key
        return self

    def base_url(self, base_url: str) -> "AzureOpenAiBuilder":
        self._base_url = base_url
        return self

    def model(self, model: str) -> "AzureOpenAiBuilder":
        self._model = model
        return self

    def client(self, client: httpx.AsyncClient) -> "AzureOpenAiBuilder":
        self._client = client
        return self

    def build(self) -> AzureOpenAiProvider:
        if not self._api_key:
            raise ProviderError.missing_api_key("AZURE_OPENAI_KEY")
        return AzureOpenAiProvider(
            client=self._client,
            api_key=self._api_key,
            base_url=self._base_url or "https://api.azureopenai.com/v1",
            model=self._model or "gpt-4o-mini",
        )


def _map_message(message: Message) -> dict[str, Any]:
    role = {
        Role.SYSTEM: "system",
        Role.USER: "user",
        Role.ASSISTANT: "assistant",
        Role.TOOL: "tool",
    }[message.role]
    value: dict[str, Any] = {"role": role, "content": message.text_content()}
    if message.tool_calls:
        value["tool_calls"] = [
            _map_tool_call(tool_call) for tool_call in message.tool_calls
        ]
    if message.tool_call_id:
        value["tool_call_id"] = message.tool_call_id
    return value


def _map_tool(tool: ToolDefinition) -> dict[str, Any]:
    return {
        "type": "function",
        "function": {
            "name": tool.name,
            "description": tool.description,
            "parameters": tool.schema,
        },
    }


def _map_tool_call(tool_call: ToolCall) -> dict[str, Any]:
    return {
        "id": tool_call.id,
        "type": "function",
        "function": {
            "name": tool_call.name,
            "arguments": json.dumps(tool_call.arguments),
        },
    }


def _map_response(payload: dict[str, Any]) -> CompletionResponse:
    choices = payload.get("choices")
    if not isinstance(choices, list) or not choices:
        raise ProviderError.invalid_response("missing choices")
    choice = choices[0]
    message_value = choice.get("message")
    if not isinstance(message_value, dict):
        raise ProviderError.invalid_response("missing message")
    stop_reason = (
        map_stop_reason(choice.get("finish_reason", "stop"))
        if choice.get("finish_reason")
        else StopReason.STOP
    )
    tool_calls = None
    raw_tool_calls = message_value.get("tool_calls")
    if isinstance(raw_tool_calls, list):
        tool_calls = []
        for call in raw_tool_calls:
            function = call.get("function") or {}
            arguments = function.get("arguments")
            try:
                parsed_arguments = (
                    json.loads(arguments) if isinstance(arguments, str) else {}
                )
            except json.JSONDecodeError:
                parsed_arguments = {}
            tool_calls.append(
                ToolCall(
                    id=str(call.get("id") or ""),
                    name=str(function.get("name") or ""),
                    arguments=parsed_arguments,
                )
            )
        if not tool_calls:
            tool_calls = None
    message = Message.assistant(str(message_value.get("content") or ""))
    message.tool_calls = tool_calls
    usage_value = payload.get("usage") or {}
    usage = Usage(
        input_tokens=int(usage_value.get("prompt_tokens") or 0),
        output_tokens=int(usage_value.get("completion_tokens") or 0),
        total_tokens=int(usage_value.get("total_tokens") or 0),
    )
    return CompletionResponse(
        message=message,
        stop_reason=stop_reason,
        usage=usage,
        model=str(payload.get("model") or ""),
        raw=payload,
    )
