"""LiteLLM wrapper for unified LLM access."""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator

from ._constants import ENV_VAR_MAP
from ._errors import TraciaError, TraciaErrorCode, sanitize_error_message
from ._types import (
    ContentPart,
    FinishReason,
    LLMProvider,
    LocalPromptMessage,
    TextPart,
    ToolCall,
    ToolCallPart,
    ToolChoice,
    ToolDefinition,
)

if TYPE_CHECKING:
    from litellm import ModelResponse


@dataclass
class CompletionResult:
    """Result from an LLM completion."""

    text: str
    input_tokens: int
    output_tokens: int
    total_tokens: int
    tool_calls: list[ToolCall]
    finish_reason: FinishReason
    provider: LLMProvider


# Model to provider mapping for common models
_MODEL_PROVIDER_MAP: dict[str, LLMProvider] = {
    # OpenAI
    "gpt-3.5-turbo": LLMProvider.OPENAI,
    "gpt-4": LLMProvider.OPENAI,
    "gpt-4-turbo": LLMProvider.OPENAI,
    "gpt-4o": LLMProvider.OPENAI,
    "gpt-4o-mini": LLMProvider.OPENAI,
    "gpt-4.1": LLMProvider.OPENAI,
    "gpt-4.1-mini": LLMProvider.OPENAI,
    "gpt-4.1-nano": LLMProvider.OPENAI,
    "gpt-4.5-preview": LLMProvider.OPENAI,
    "gpt-5": LLMProvider.OPENAI,
    "o1": LLMProvider.OPENAI,
    "o1-mini": LLMProvider.OPENAI,
    "o1-preview": LLMProvider.OPENAI,
    "o3": LLMProvider.OPENAI,
    "o3-mini": LLMProvider.OPENAI,
    "o4-mini": LLMProvider.OPENAI,
    # Anthropic
    "claude-3-haiku-20240307": LLMProvider.ANTHROPIC,
    "claude-3-sonnet-20240229": LLMProvider.ANTHROPIC,
    "claude-3-opus-20240229": LLMProvider.ANTHROPIC,
    "claude-3-5-haiku-20241022": LLMProvider.ANTHROPIC,
    "claude-3-5-sonnet-20241022": LLMProvider.ANTHROPIC,
    "claude-sonnet-4-20250514": LLMProvider.ANTHROPIC,
    "claude-opus-4-20250514": LLMProvider.ANTHROPIC,
    # Google
    "gemini-2.0-flash": LLMProvider.GOOGLE,
    "gemini-2.0-flash-lite": LLMProvider.GOOGLE,
    "gemini-2.5-pro": LLMProvider.GOOGLE,
    "gemini-2.5-flash": LLMProvider.GOOGLE,
}


def resolve_provider(model: str, explicit_provider: LLMProvider | None) -> LLMProvider:
    """Resolve the provider for a model.

    Args:
        model: The model name.
        explicit_provider: Explicitly specified provider.

    Returns:
        The resolved provider.

    Raises:
        TraciaError: If the provider cannot be determined.
    """
    if explicit_provider is not None:
        return explicit_provider

    # Check the model registry
    if model in _MODEL_PROVIDER_MAP:
        return _MODEL_PROVIDER_MAP[model]

    # Try prefix-based detection
    if model.startswith("gpt-") or model.startswith("o1") or model.startswith("o3") or model.startswith("o4"):
        return LLMProvider.OPENAI
    if model.startswith("claude-"):
        return LLMProvider.ANTHROPIC
    if model.startswith("gemini-"):
        return LLMProvider.GOOGLE

    raise TraciaError(
        code=TraciaErrorCode.UNSUPPORTED_MODEL,
        message=f"Cannot determine provider for model '{model}'. Please specify the provider explicitly.",
    )


def get_litellm_model(model: str, provider: LLMProvider) -> str:
    """Get the litellm-compatible model name.

    LiteLLM requires a ``gemini/`` prefix to route Google AI Studio models
    correctly. Without it, litellm defaults to the Vertex AI path which
    requires Application Default Credentials instead of an API key.

    Args:
        model: The user-facing model name (e.g. ``gemini-2.0-flash``).
        provider: The resolved provider.

    Returns:
        The model string suitable for ``litellm.completion()``.
    """
    if provider == LLMProvider.GOOGLE and not model.startswith("gemini/"):
        return f"gemini/{model}"
    return model


def get_provider_api_key(
    provider: LLMProvider,
    provider_api_key: str | None = None,
) -> str:
    """Get the API key for a provider.

    Args:
        provider: The LLM provider.
        provider_api_key: Explicitly provided API key.

    Returns:
        The API key.

    Raises:
        TraciaError: If no API key is found.
    """
    if provider_api_key:
        return provider_api_key

    env_var = ENV_VAR_MAP.get(provider.value)
    if env_var:
        key = os.environ.get(env_var)
        if key:
            return key

    raise TraciaError(
        code=TraciaErrorCode.MISSING_PROVIDER_API_KEY,
        message=f"No API key found for provider '{provider.value}'. "
        f"Set the {ENV_VAR_MAP.get(provider.value, 'PROVIDER_API_KEY')} environment variable "
        "or pass provider_api_key parameter.",
    )


def convert_messages(
    messages: list[LocalPromptMessage],
) -> list[dict[str, Any]]:
    """Convert Tracia messages to LiteLLM/OpenAI format.

    Args:
        messages: The Tracia messages.

    Returns:
        Messages in LiteLLM format.
    """
    result: list[dict[str, Any]] = []

    for msg in messages:
        # Handle tool role
        if msg.role == "tool":
            content = msg.content if isinstance(msg.content, str) else str(msg.content)
            result.append({
                "role": "tool",
                "tool_call_id": msg.tool_call_id,
                "content": content,
            })
            continue

        # Handle developer role (map to system)
        role = "system" if msg.role == "developer" else msg.role

        # Handle string content
        if isinstance(msg.content, str):
            result.append({"role": role, "content": msg.content})
            continue

        # Handle list content (text parts and tool calls)
        content_parts: list[dict[str, Any]] = []
        tool_calls: list[dict[str, Any]] = []

        for part in msg.content:
            if isinstance(part, TextPart) or (isinstance(part, dict) and part.get("type") == "text"):
                text = part.text if isinstance(part, TextPart) else part.get("text", "")
                content_parts.append({"type": "text", "text": text})
            elif isinstance(part, ToolCallPart) or (isinstance(part, dict) and part.get("type") == "tool_call"):
                if isinstance(part, ToolCallPart):
                    tc_id = part.id
                    tc_name = part.name
                    tc_args = part.arguments
                else:
                    tc_id = part.get("id", "")
                    tc_name = part.get("name", "")
                    tc_args = part.get("arguments", {})

                tool_calls.append({
                    "id": tc_id,
                    "type": "function",
                    "function": {
                        "name": tc_name,
                        "arguments": json.dumps(tc_args) if isinstance(tc_args, dict) else tc_args,
                    },
                })

        # Build the message
        msg_dict: dict[str, Any] = {"role": role}

        if content_parts:
            # If only text parts, we can simplify
            if len(content_parts) == 1 and not tool_calls:
                msg_dict["content"] = content_parts[0]["text"]
            else:
                msg_dict["content"] = content_parts
        elif not tool_calls:
            msg_dict["content"] = ""

        if tool_calls:
            msg_dict["tool_calls"] = tool_calls

        result.append(msg_dict)

    return result


def convert_tools(tools: list[ToolDefinition] | None) -> list[dict[str, Any]] | None:
    """Convert tool definitions to LiteLLM format.

    Args:
        tools: The tool definitions.

    Returns:
        Tools in LiteLLM format.
    """
    if not tools:
        return None

    result = []
    for tool in tools:
        result.append({
            "type": "function",
            "function": {
                "name": tool.name,
                "description": tool.description,
                "parameters": tool.parameters.model_dump(exclude_none=True),
            },
        })
    return result


def convert_tool_choice(tool_choice: ToolChoice | None) -> str | dict[str, Any] | None:
    """Convert tool choice to LiteLLM format.

    Args:
        tool_choice: The tool choice.

    Returns:
        Tool choice in LiteLLM format.
    """
    if tool_choice is None:
        return None

    if isinstance(tool_choice, str):
        return tool_choice

    if isinstance(tool_choice, dict) and "tool" in tool_choice:
        return {"type": "function", "function": {"name": tool_choice["tool"]}}

    return None


def parse_finish_reason(reason: str | None) -> FinishReason:
    """Parse the finish reason from LiteLLM response.

    Args:
        reason: The raw finish reason.

    Returns:
        The normalized finish reason.
    """
    if reason == "tool_calls":
        return "tool_calls"
    if reason == "length":
        return "max_tokens"
    return "stop"


def extract_tool_calls(response: "ModelResponse") -> list[ToolCall]:
    """Extract tool calls from a LiteLLM response.

    Args:
        response: The LiteLLM response.

    Returns:
        The extracted tool calls.
    """
    tool_calls: list[ToolCall] = []

    choices = getattr(response, "choices", [])
    if not choices:
        return tool_calls

    message = getattr(choices[0], "message", None)
    if not message:
        return tool_calls

    raw_tool_calls = getattr(message, "tool_calls", None)
    if not raw_tool_calls:
        return tool_calls

    for tc in raw_tool_calls:
        func = getattr(tc, "function", None)
        if func:
            try:
                args = json.loads(func.arguments) if isinstance(func.arguments, str) else func.arguments
            except json.JSONDecodeError:
                args = {}

            tool_calls.append(ToolCall(
                id=tc.id,
                name=func.name,
                arguments=args,
            ))

    return tool_calls


class LLMClient:
    """Client for making LLM calls via LiteLLM."""

    def complete(
        self,
        model: str,
        messages: list[LocalPromptMessage],
        *,
        provider: LLMProvider | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | None = None,
        stop: list[str] | None = None,
        tools: list[ToolDefinition] | None = None,
        tool_choice: ToolChoice | None = None,
        api_key: str | None = None,
        timeout: float | None = None,
    ) -> CompletionResult:
        """Make a synchronous completion request.

        Args:
            model: The model name.
            messages: The messages to send.
            provider: The LLM provider.
            temperature: Sampling temperature.
            max_tokens: Maximum output tokens.
            top_p: Top-p sampling.
            stop: Stop sequences.
            tools: Tool definitions.
            tool_choice: Tool choice setting.
            api_key: Provider API key.
            timeout: Request timeout in seconds.

        Returns:
            The completion result.

        Raises:
            TraciaError: If the request fails.
        """
        try:
            import litellm
        except ImportError as e:
            raise TraciaError(
                code=TraciaErrorCode.MISSING_PROVIDER_SDK,
                message="litellm is not installed. Install it with: pip install litellm",
            ) from e

        resolved_provider = resolve_provider(model, provider)
        resolved_api_key = get_provider_api_key(resolved_provider, api_key)

        # Build the request
        litellm_messages = convert_messages(messages)
        litellm_tools = convert_tools(tools)
        litellm_tool_choice = convert_tool_choice(tool_choice)

        request_kwargs: dict[str, Any] = {
            "model": get_litellm_model(model, resolved_provider),
            "messages": litellm_messages,
            "api_key": resolved_api_key,
        }

        if temperature is not None:
            request_kwargs["temperature"] = temperature
        if max_tokens is not None:
            request_kwargs["max_tokens"] = max_tokens
        if top_p is not None:
            request_kwargs["top_p"] = top_p
        if stop is not None:
            request_kwargs["stop"] = stop
        if litellm_tools is not None:
            request_kwargs["tools"] = litellm_tools
        if litellm_tool_choice is not None:
            request_kwargs["tool_choice"] = litellm_tool_choice
        if timeout is not None:
            request_kwargs["timeout"] = timeout

        try:
            response = litellm.completion(**request_kwargs)
        except Exception as e:
            error_msg = sanitize_error_message(str(e))
            raise TraciaError(
                code=TraciaErrorCode.PROVIDER_ERROR,
                message=f"LLM provider error: {error_msg}",
            ) from e

        # Extract result
        usage = getattr(response, "usage", None)
        choices = getattr(response, "choices", [])
        message = choices[0].message if choices else None
        content = getattr(message, "content", "") or ""
        finish_reason = choices[0].finish_reason if choices else "stop"

        return CompletionResult(
            text=content,
            input_tokens=getattr(usage, "prompt_tokens", 0) if usage else 0,
            output_tokens=getattr(usage, "completion_tokens", 0) if usage else 0,
            total_tokens=getattr(usage, "total_tokens", 0) if usage else 0,
            tool_calls=extract_tool_calls(response),
            finish_reason=parse_finish_reason(finish_reason),
            provider=resolved_provider,
        )

    async def acomplete(
        self,
        model: str,
        messages: list[LocalPromptMessage],
        *,
        provider: LLMProvider | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | None = None,
        stop: list[str] | None = None,
        tools: list[ToolDefinition] | None = None,
        tool_choice: ToolChoice | None = None,
        api_key: str | None = None,
        timeout: float | None = None,
    ) -> CompletionResult:
        """Make an asynchronous completion request.

        Args:
            model: The model name.
            messages: The messages to send.
            provider: The LLM provider.
            temperature: Sampling temperature.
            max_tokens: Maximum output tokens.
            top_p: Top-p sampling.
            stop: Stop sequences.
            tools: Tool definitions.
            tool_choice: Tool choice setting.
            api_key: Provider API key.
            timeout: Request timeout in seconds.

        Returns:
            The completion result.

        Raises:
            TraciaError: If the request fails.
        """
        try:
            import litellm
        except ImportError as e:
            raise TraciaError(
                code=TraciaErrorCode.MISSING_PROVIDER_SDK,
                message="litellm is not installed. Install it with: pip install litellm",
            ) from e

        resolved_provider = resolve_provider(model, provider)
        resolved_api_key = get_provider_api_key(resolved_provider, api_key)

        # Build the request
        litellm_messages = convert_messages(messages)
        litellm_tools = convert_tools(tools)
        litellm_tool_choice = convert_tool_choice(tool_choice)

        request_kwargs: dict[str, Any] = {
            "model": get_litellm_model(model, resolved_provider),
            "messages": litellm_messages,
            "api_key": resolved_api_key,
        }

        if temperature is not None:
            request_kwargs["temperature"] = temperature
        if max_tokens is not None:
            request_kwargs["max_tokens"] = max_tokens
        if top_p is not None:
            request_kwargs["top_p"] = top_p
        if stop is not None:
            request_kwargs["stop"] = stop
        if litellm_tools is not None:
            request_kwargs["tools"] = litellm_tools
        if litellm_tool_choice is not None:
            request_kwargs["tool_choice"] = litellm_tool_choice
        if timeout is not None:
            request_kwargs["timeout"] = timeout

        try:
            response = await litellm.acompletion(**request_kwargs)
        except Exception as e:
            error_msg = sanitize_error_message(str(e))
            raise TraciaError(
                code=TraciaErrorCode.PROVIDER_ERROR,
                message=f"LLM provider error: {error_msg}",
            ) from e

        # Extract result
        usage = getattr(response, "usage", None)
        choices = getattr(response, "choices", [])
        message = choices[0].message if choices else None
        content = getattr(message, "content", "") or ""
        finish_reason = choices[0].finish_reason if choices else "stop"

        return CompletionResult(
            text=content,
            input_tokens=getattr(usage, "prompt_tokens", 0) if usage else 0,
            output_tokens=getattr(usage, "completion_tokens", 0) if usage else 0,
            total_tokens=getattr(usage, "total_tokens", 0) if usage else 0,
            tool_calls=extract_tool_calls(response),
            finish_reason=parse_finish_reason(finish_reason),
            provider=resolved_provider,
        )

    def stream(
        self,
        model: str,
        messages: list[LocalPromptMessage],
        *,
        provider: LLMProvider | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | None = None,
        stop: list[str] | None = None,
        tools: list[ToolDefinition] | None = None,
        tool_choice: ToolChoice | None = None,
        api_key: str | None = None,
        timeout: float | None = None,
    ) -> tuple[Iterator[str], list[CompletionResult], LLMProvider]:
        """Make a streaming completion request.

        Args:
            model: The model name.
            messages: The messages to send.
            provider: The LLM provider.
            temperature: Sampling temperature.
            max_tokens: Maximum output tokens.
            top_p: Top-p sampling.
            stop: Stop sequences.
            tools: Tool definitions.
            tool_choice: Tool choice setting.
            api_key: Provider API key.
            timeout: Request timeout in seconds.

        Returns:
            A tuple of (chunk iterator, result holder list, provider).

        Raises:
            TraciaError: If the request fails.
        """
        try:
            import litellm
        except ImportError as e:
            raise TraciaError(
                code=TraciaErrorCode.MISSING_PROVIDER_SDK,
                message="litellm is not installed. Install it with: pip install litellm",
            ) from e

        resolved_provider = resolve_provider(model, provider)
        resolved_api_key = get_provider_api_key(resolved_provider, api_key)

        # Build the request
        litellm_messages = convert_messages(messages)
        litellm_tools = convert_tools(tools)
        litellm_tool_choice = convert_tool_choice(tool_choice)

        request_kwargs: dict[str, Any] = {
            "model": get_litellm_model(model, resolved_provider),
            "messages": litellm_messages,
            "api_key": resolved_api_key,
            "stream": True,
        }

        if temperature is not None:
            request_kwargs["temperature"] = temperature
        if max_tokens is not None:
            request_kwargs["max_tokens"] = max_tokens
        if top_p is not None:
            request_kwargs["top_p"] = top_p
        if stop is not None:
            request_kwargs["stop"] = stop
        if litellm_tools is not None:
            request_kwargs["tools"] = litellm_tools
        if litellm_tool_choice is not None:
            request_kwargs["tool_choice"] = litellm_tool_choice
        if timeout is not None:
            request_kwargs["timeout"] = timeout

        result_holder: list[CompletionResult] = []

        def generate_chunks() -> Iterator[str]:
            full_text = ""
            input_tokens = 0
            output_tokens = 0
            total_tokens = 0
            tool_calls: list[ToolCall] = []
            finish_reason: FinishReason = "stop"
            tool_call_chunks: dict[int, dict[str, Any]] = {}

            try:
                response = litellm.completion(**request_kwargs)

                for chunk in response:
                    choices = getattr(chunk, "choices", [])
                    if not choices:
                        continue

                    delta = getattr(choices[0], "delta", None)
                    if delta:
                        content = getattr(delta, "content", None)
                        if content:
                            full_text += content
                            yield content

                        # Handle streaming tool calls
                        delta_tool_calls = getattr(delta, "tool_calls", None)
                        if delta_tool_calls:
                            for tc in delta_tool_calls:
                                idx = tc.index
                                if idx not in tool_call_chunks:
                                    tool_call_chunks[idx] = {
                                        "id": "",
                                        "name": "",
                                        "arguments": "",
                                    }
                                if tc.id:
                                    tool_call_chunks[idx]["id"] = tc.id
                                if tc.function:
                                    if tc.function.name:
                                        tool_call_chunks[idx]["name"] = tc.function.name
                                    if tc.function.arguments:
                                        tool_call_chunks[idx]["arguments"] += tc.function.arguments

                    chunk_finish = getattr(choices[0], "finish_reason", None)
                    if chunk_finish:
                        finish_reason = parse_finish_reason(chunk_finish)

                    # Extract usage from final chunk
                    usage = getattr(chunk, "usage", None)
                    if usage:
                        input_tokens = getattr(usage, "prompt_tokens", 0)
                        output_tokens = getattr(usage, "completion_tokens", 0)
                        total_tokens = getattr(usage, "total_tokens", 0)

                # Convert accumulated tool calls
                for idx in sorted(tool_call_chunks.keys()):
                    tc_data = tool_call_chunks[idx]
                    try:
                        args = json.loads(tc_data["arguments"]) if tc_data["arguments"] else {}
                    except json.JSONDecodeError:
                        args = {}
                    tool_calls.append(ToolCall(
                        id=tc_data["id"],
                        name=tc_data["name"],
                        arguments=args,
                    ))

                result_holder.append(CompletionResult(
                    text=full_text,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    total_tokens=total_tokens,
                    tool_calls=tool_calls,
                    finish_reason=finish_reason,
                    provider=resolved_provider,
                ))

            except Exception as e:
                error_msg = sanitize_error_message(str(e))
                raise TraciaError(
                    code=TraciaErrorCode.PROVIDER_ERROR,
                    message=f"LLM provider error: {error_msg}",
                ) from e

        return generate_chunks(), result_holder, resolved_provider

    async def astream(
        self,
        model: str,
        messages: list[LocalPromptMessage],
        *,
        provider: LLMProvider | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | None = None,
        stop: list[str] | None = None,
        tools: list[ToolDefinition] | None = None,
        tool_choice: ToolChoice | None = None,
        api_key: str | None = None,
        timeout: float | None = None,
    ) -> tuple[AsyncIterator[str], list[CompletionResult], LLMProvider]:
        """Make an async streaming completion request.

        Args:
            model: The model name.
            messages: The messages to send.
            provider: The LLM provider.
            temperature: Sampling temperature.
            max_tokens: Maximum output tokens.
            top_p: Top-p sampling.
            stop: Stop sequences.
            tools: Tool definitions.
            tool_choice: Tool choice setting.
            api_key: Provider API key.
            timeout: Request timeout in seconds.

        Returns:
            A tuple of (async chunk iterator, result holder list, provider).

        Raises:
            TraciaError: If the request fails.
        """
        try:
            import litellm
        except ImportError as e:
            raise TraciaError(
                code=TraciaErrorCode.MISSING_PROVIDER_SDK,
                message="litellm is not installed. Install it with: pip install litellm",
            ) from e

        resolved_provider = resolve_provider(model, provider)
        resolved_api_key = get_provider_api_key(resolved_provider, api_key)

        # Build the request
        litellm_messages = convert_messages(messages)
        litellm_tools = convert_tools(tools)
        litellm_tool_choice = convert_tool_choice(tool_choice)

        request_kwargs: dict[str, Any] = {
            "model": get_litellm_model(model, resolved_provider),
            "messages": litellm_messages,
            "api_key": resolved_api_key,
            "stream": True,
        }

        if temperature is not None:
            request_kwargs["temperature"] = temperature
        if max_tokens is not None:
            request_kwargs["max_tokens"] = max_tokens
        if top_p is not None:
            request_kwargs["top_p"] = top_p
        if stop is not None:
            request_kwargs["stop"] = stop
        if litellm_tools is not None:
            request_kwargs["tools"] = litellm_tools
        if litellm_tool_choice is not None:
            request_kwargs["tool_choice"] = litellm_tool_choice
        if timeout is not None:
            request_kwargs["timeout"] = timeout

        result_holder: list[CompletionResult] = []

        async def generate_chunks() -> AsyncIterator[str]:
            full_text = ""
            input_tokens = 0
            output_tokens = 0
            total_tokens = 0
            tool_calls: list[ToolCall] = []
            finish_reason: FinishReason = "stop"
            tool_call_chunks: dict[int, dict[str, Any]] = {}

            try:
                response = await litellm.acompletion(**request_kwargs)

                async for chunk in response:
                    choices = getattr(chunk, "choices", [])
                    if not choices:
                        continue

                    delta = getattr(choices[0], "delta", None)
                    if delta:
                        content = getattr(delta, "content", None)
                        if content:
                            full_text += content
                            yield content

                        # Handle streaming tool calls
                        delta_tool_calls = getattr(delta, "tool_calls", None)
                        if delta_tool_calls:
                            for tc in delta_tool_calls:
                                idx = tc.index
                                if idx not in tool_call_chunks:
                                    tool_call_chunks[idx] = {
                                        "id": "",
                                        "name": "",
                                        "arguments": "",
                                    }
                                if tc.id:
                                    tool_call_chunks[idx]["id"] = tc.id
                                if tc.function:
                                    if tc.function.name:
                                        tool_call_chunks[idx]["name"] = tc.function.name
                                    if tc.function.arguments:
                                        tool_call_chunks[idx]["arguments"] += tc.function.arguments

                    chunk_finish = getattr(choices[0], "finish_reason", None)
                    if chunk_finish:
                        finish_reason = parse_finish_reason(chunk_finish)

                    # Extract usage from final chunk
                    usage = getattr(chunk, "usage", None)
                    if usage:
                        input_tokens = getattr(usage, "prompt_tokens", 0)
                        output_tokens = getattr(usage, "completion_tokens", 0)
                        total_tokens = getattr(usage, "total_tokens", 0)

                # Convert accumulated tool calls
                for idx in sorted(tool_call_chunks.keys()):
                    tc_data = tool_call_chunks[idx]
                    try:
                        args = json.loads(tc_data["arguments"]) if tc_data["arguments"] else {}
                    except json.JSONDecodeError:
                        args = {}
                    tool_calls.append(ToolCall(
                        id=tc_data["id"],
                        name=tc_data["name"],
                        arguments=args,
                    ))

                result_holder.append(CompletionResult(
                    text=full_text,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    total_tokens=total_tokens,
                    tool_calls=tool_calls,
                    finish_reason=finish_reason,
                    provider=resolved_provider,
                ))

            except Exception as e:
                error_msg = sanitize_error_message(str(e))
                raise TraciaError(
                    code=TraciaErrorCode.PROVIDER_ERROR,
                    message=f"LLM provider error: {error_msg}",
                ) from e

        return generate_chunks(), result_holder, resolved_provider


def build_assistant_message(
    text: str,
    tool_calls: list[ToolCall],
) -> LocalPromptMessage:
    """Build an assistant message from completion result.

    Args:
        text: The text content.
        tool_calls: Any tool calls made.

    Returns:
        The assistant message.
    """
    if not tool_calls:
        return LocalPromptMessage(role="assistant", content=text)

    content: list[ContentPart] = []

    if text:
        content.append(TextPart(type="text", text=text))

    for tc in tool_calls:
        content.append(ToolCallPart(
            type="tool_call",
            id=tc.id,
            name=tc.name,
            arguments=tc.arguments,
        ))

    return LocalPromptMessage(role="assistant", content=content)
