import base64
import dataclasses
from dataclasses import asdict, dataclass
from enum import Enum
from io import BytesIO
from typing import (
    Any,
    Dict,
    List,
    Mapping,
    Optional,
    Sequence,
    Union,
    Literal,
    AsyncGenerator,
)

from pydantic import BaseModel
from aleph_alpha_client.structured_output import ResponseFormat
from PIL.Image import Image


class Role(str, Enum):
    """A role used for a message in a chat."""

    User = "user"
    Assistant = "assistant"
    System = "system"
    Tool = "tool"


@dataclass(frozen=True)
class FunctionCall:
    name: str
    arguments: str


@dataclass(frozen=True)
class ToolCall:
    id: str
    type: str
    function: FunctionCall

    @staticmethod
    def from_json(json: Dict[str, Any]) -> "ToolCall":
        function = json["function"]
        return ToolCall(
            id=json["id"],
            type=json["type"],
            function=FunctionCall(
                name=function["name"], arguments=function["arguments"]
            ),
        )

    def to_json(self) -> Mapping[str, Any]:
        return {
            "id": self.id,
            "type": self.type,
            "function": {
                "name": self.function.name,
                "arguments": self.function.arguments,
            },
        }


@dataclass(frozen=True)
class Message:
    """
    Describes a message in a chat.

    Parameters:
        role (Role, required):
            The role of the message.

        content (str | List[Union[str | Image]], required):
            The content of the message.
    """

    role: Role
    content: Union[str, List[Union[str, Image]]]
    tool_call_id: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None

    def to_json(self) -> Mapping[str, Any]:
        result = {
            "role": self.role.value,
            "content": _message_content_to_json(self.content),
        }
        if self.tool_calls is not None:
            result["tool_calls"] = [t.to_json() for t in self.tool_calls]
        if self.tool_call_id is not None:
            result["tool_call_id"] = self.tool_call_id
        return result


# We introduce a more specific message type because chat responses can only
# contain text at the moment. This enables static type checking to proof that
# `content` is always a string.
@dataclass(frozen=True)
class TextMessage:
    """
    Describes a text message in a chat.

    Parameters:
        role (Role, required):
            The role of the message.

        content (str, required):
            The content of the message.
    """

    role: Role
    content: str
    tool_calls: Optional[List[ToolCall]] = None

    @staticmethod
    def from_json(json: Dict[str, Any]) -> "TextMessage":
        tool_calls = json.get("tool_calls")
        return TextMessage(
            role=Role(json["role"]),
            content=json["content"],
            tool_calls=None
            if tool_calls is None
            else [ToolCall.from_json(tool_call) for tool_call in tool_calls],
        )

    # In multi-turn conversations the returned TextMessage is part of the chat
    # history and converted to the prompt. As such, it requires conversion to
    # json again. Here, the message content is a string, but can reuse the
    # _message_content_to_json function nonetheless.
    def to_json(self) -> Mapping[str, Any]:
        result = {
            "role": self.role.value,
            "content": _message_content_to_json(self.content),
        }
        if self.tool_calls is not None:
            result["tool_calls"] = [t.to_json() for t in self.tool_calls]
        return result


def _message_content_to_json(
    content: Union[str, List[Union[str, Image]]],
) -> Union[str, List[Mapping[str, Any]]]:
    if isinstance(content, str):
        return content
    else:
        result: List[Mapping[str, Any]] = []
        for chunk in content:
            if isinstance(chunk, str):
                result.append({"type": "text", "text": chunk})
            elif isinstance(chunk, Image):
                result.append(
                    {
                        "type": "image_url",
                        "image_url": {"url": _image_to_data_uri(chunk)},
                    }
                )
            else:
                raise ValueError(
                    "The item in the prompt is not valid. Try either a string or an Image."
                )
        return result


def _image_to_data_uri(image: Image) -> str:
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    base_64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{base_64}"


@dataclass(frozen=True)
class StreamOptions:
    """
    Additional options to affect the streaming behavior.
    """

    # If set, an additional chunk will be streamed before the data: [DONE] message.
    # The usage field on this chunk shows the token usage statistics for the entire
    # request, and the choices field will always be an empty array.
    include_usage: bool


@dataclass(frozen=True)
class ToolFunction:
    type: Literal["function"]
    function: Any


@dataclass(frozen=True)
class ChatRequest:
    """
    Describes a chat request.

    Only supports a subset of the parameters of `CompletionRequest` for simplicity.
    See `CompletionRequest` for documentation on the parameters.
    """

    model: str
    messages: Sequence[Union[Message, TextMessage]]
    maximum_tokens: Optional[int] = None
    temperature: Optional[float] = None
    top_k: Optional[int] = None
    top_p: Optional[float] = None
    stream_options: Optional[StreamOptions] = None
    steering_concepts: Optional[List[str]] = None
    response_format: Optional[ResponseFormat] = None

    tools: Optional[List[Any]] = None
    tool_choice: Optional[Union[Literal["auto", "required", "none"], ToolFunction]] = (
        None
    )
    parallel_tool_calls: Optional[bool] = None

    def to_json(self) -> Mapping[str, Any]:
        payload = {k: v for k, v in asdict(self).items() if v is not None}
        payload["messages"] = [message.to_json() for message in self.messages]
        if self.response_format:
            # Handle Pydantic models by converting them to JSONSchema first
            if isinstance(self.response_format, type) and issubclass(
                self.response_format, BaseModel
            ):
                # This is a Pydantic model, convert it to JSONSchema
                from aleph_alpha_client.structured_output import JSONSchema

                json_schema = JSONSchema.from_pydantic(self.response_format)
                payload["response_format"] = json_schema.to_json()
            else:
                # This is already a JSONSchema or compatible object
                payload["response_format"] = self.response_format.to_json()
        return payload


class FinishReason(str, Enum):
    """
    The reason the model stopped generating tokens.

    This will be `stop` if the model hit a natural stop point or a provided stop
    sequence or length if the maximum number of tokens specified in the request
    was reached. If the API is unable to understand the stop reason emitted by
    one of the workers, content_filter is returned.
    """

    Stop = "stop"
    Length = "length"
    ContentFilter = "content_filter"
    ToolCalls = "tool_calls"


@dataclass(frozen=True)
class ChatResponse:
    """
    A simplified version of the chat response.

    As the `ChatRequest` does not support the `n` parameter (allowing for multiple return values),
    the `ChatResponse` assumes there to be only one choice.
    """

    finish_reason: FinishReason
    message: TextMessage

    @staticmethod
    def from_json(json: Dict[str, Any]) -> "ChatResponse":
        first_choice = json["choices"][0]
        return ChatResponse(
            finish_reason=FinishReason(first_choice["finish_reason"]),
            message=TextMessage.from_json(first_choice["message"]),
        )


@dataclass(frozen=True)
class Usage:
    """
    Usage statistics for the completion request.

    When streaming is enabled, this field will be null by default.
    To include an additional usage-only message in the response stream, set stream_options.include_usage to true.
    """

    # Number of tokens in the generated completion.
    completion_tokens: int

    # Number of tokens in the prompt.
    prompt_tokens: int

    # Total number of tokens used in the request (prompt + completion).
    total_tokens: int

    @staticmethod
    def from_json(json: Dict[str, Any]) -> "Usage":
        return Usage(
            completion_tokens=json["completion_tokens"],
            prompt_tokens=json["prompt_tokens"],
            total_tokens=json["total_tokens"],
        )


@dataclass(frozen=True)
class ChatStreamChunk:
    """
    A streamed chat completion chunk.

    Parameters:
        content (str, required):
            The content of the current chat completion. Will be empty for the first chunk of every completion stream and non-empty for the remaining chunks.

        role (Role, optional):
            The role of the current chat completion. Will be assistant for the first chunk of every completion stream and missing for the remaining chunks.
    """

    content: str
    role: Optional[Role]
    tool_calls: Optional[List[ToolCall]]

    @staticmethod
    def from_json(choice: Dict[str, Any]) -> "ChatStreamChunk":
        delta = choice["delta"]
        tool_calls = delta.get("tool_calls")
        return ChatStreamChunk(
            content=delta["content"],
            role=Role(delta.get("role")) if delta.get("role") else None,
            tool_calls=None
            if tool_calls is None
            else [ToolCall.from_json(tool_call) for tool_call in tool_calls],
        )


def stream_chat_item_from_json(
    json: Dict[str, Any],
) -> Union[ChatStreamChunk, FinishReason, Usage]:
    """Parse a chat event into one of the three possible types.

    This function takes two assumptions:

    1. If neither a finish reason nor a usage is present, the chunk contains a message.
    2. The chunk only carries information relevant to one of the three possible types, so no information is lost by choosing any one of them.
    """
    if (usage := json.get("usage")) is not None:
        return Usage.from_json(usage)

    first_choice = json["choices"][0]
    if (finish_reason := first_choice.get("finish_reason")) is not None:
        return FinishReason(finish_reason)

    return ChatStreamChunk.from_json(first_choice)


@dataclasses.dataclass
class ToolCallState:
    id: str
    arguments: str
    type: str
    function_name: str


async def process_chat_stream(
    s: AsyncGenerator[Dict[str, Any], None],
) -> AsyncGenerator[Union[ChatStreamChunk, Usage, ToolCall, FinishReason], None]:
    tool_calls: dict[str, ToolCallState] = {}
    # This loop collects all tool calls and emits them at the end of the stream *before* the finish reason
    # We have to do this because tool calls are streamed and distributed over different chunks
    async for json in s:
        if (usage := json.get("usage")) is not None:
            yield Usage.from_json(usage)
            continue

        first_choice = json["choices"][0]
        if (finish_reason := first_choice.get("finish_reason")) is not None:
            # Stream done, we have to yield all tool calls before the finish reason
            for tool in tool_calls.values():
                yield ToolCall(
                    id=tool.id,
                    type=tool.type,
                    function=FunctionCall(
                        name=tool.function_name,
                        arguments=tool.arguments,
                    ),
                )
            yield FinishReason(finish_reason)
            continue

        delta = first_choice.get("delta")
        if (tool_calls_json := delta.pop("tool_calls", None)) is not None:
            for tool_call in tool_calls_json:
                # We received a tool call delta, append all deltas to our state
                # or create the state if it is the first message for this tool call
                index = tool_call["index"]
                if index in tool_calls:
                    tool_calls[index].arguments += tool_call["function"]["arguments"]
                else:
                    function = tool_call["function"]
                    tool_calls[index] = ToolCallState(
                        id=tool_call["id"],
                        type=tool_call["type"],
                        function_name=function["name"],
                        arguments=function.get("arguments", ""),
                    )

            content = delta.get("content")
            if content is None or content == "":
                # We skip empty content here because our current api always returns content,
                # even for tool_call messages where this is not needed
                continue

        yield stream_chat_item_from_json(json)
