from typing import List, Optional

from litellm import Usage as LitellmUsage
from litellm.types.utils import Message as LitellmMessage
from litellm.utils import Choices as LitellmChoice
from litellm.utils import ModelResponse as LitellmModelResponse
from rustic_ai.core.guild.agent_ext.depends.llm.models import (
    AssistantMessage,
    ChatCompletionMessageToolCall,
    ChatCompletionResponse,
    Choice,
    CompletionUsage,
    FinishReason,
    FunctionCall,
    ToolType,
)


class MessageUtils:
    @staticmethod
    def from_litellm(litellm_message: LitellmMessage) -> AssistantMessage:
        tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
        if hasattr(litellm_message, "tool_calls") and litellm_message.tool_calls:
            for tool_call in litellm_message.tool_calls:
                tool_calls = []
                if isinstance(tool_call, ChatCompletionMessageToolCall):
                    tool_calls.append(
                        ChatCompletionMessageToolCall(
                            id=tool_call.id,
                            type=ToolType.function,
                            function=FunctionCall(name=tool_call.function.name, arguments=tool_call.function.arguments),
                        )
                    )
        return AssistantMessage(
            content=litellm_message.content,
            tool_calls=tool_calls,
        )


class ChoiceUtils:
    @staticmethod
    def from_litellm(litellm_choice: LitellmChoice) -> Choice:
        return Choice(
            finish_reason=FinishReason(litellm_choice.finish_reason),
            index=litellm_choice.index,
            message=MessageUtils.from_litellm(litellm_choice.message),
        )


class ResponseUtils:
    @staticmethod
    def from_litellm(litellm_response: LitellmModelResponse) -> ChatCompletionResponse:
        cusage = None
        usage: LitellmUsage = litellm_response.get("usage")
        if usage:
            cusage = CompletionUsage(
                completion_tokens=usage.completion_tokens,
                prompt_tokens=usage.prompt_tokens,
                total_tokens=usage.total_tokens,
            )

        return ChatCompletionResponse(
            id=litellm_response.id,
            choices=[
                ChoiceUtils.from_litellm(choice)
                for choice in litellm_response.choices
                if isinstance(choice, LitellmChoice)
            ],
            created=litellm_response.created,
            model=litellm_response.model,
            system_fingerprint=litellm_response.system_fingerprint,
            usage=cusage,
        )
