"""
This implementation is based on the OpenAI types, while removing all the parts that are not needed for Browser Use.
"""

# region - Content parts
from typing import Literal, Union

from pydantic import BaseModel


def _truncate(text: str, max_length: int = 50) -> str:
	"""Truncate text to max_length characters, adding ellipsis if truncated."""
	if len(text) <= max_length:
		return text
	return text[: max_length - 3] + '...'


def _format_image_url(url: str, max_length: int = 50) -> str:
	"""Format image URL for display, truncating if necessary."""
	if url.startswith('data:'):
		# Base64 image
		media_type = url.split(';')[0].split(':')[1] if ';' in url else 'image'
		return f'<base64 {media_type}>'
	else:
		# Regular URL
		return _truncate(url, max_length)


class ContentPartTextParam(BaseModel):
	text: str
	type: Literal['text'] = 'text'

	def __str__(self) -> str:
		return f'Text: {_truncate(self.text)}'

	def __repr__(self) -> str:
		return f'ContentPartTextParam(text={_truncate(self.text)})'


class ContentPartRefusalParam(BaseModel):
	refusal: str
	type: Literal['refusal'] = 'refusal'

	def __str__(self) -> str:
		return f'Refusal: {_truncate(self.refusal)}'

	def __repr__(self) -> str:
		return f'ContentPartRefusalParam(refusal={_truncate(repr(self.refusal), 50)})'


SupportedImageMediaType = Literal['image/jpeg', 'image/png', 'image/gif', 'image/webp']


class ImageURL(BaseModel):
	url: str
	"""Either a URL of the image or the base64 encoded image data."""
	detail: Literal['auto', 'low', 'high'] = 'auto'
	"""Specifies the detail level of the image.

    Learn more in the
    [Vision guide](https://platform.openai.com/docs/guides/vision#low-or-high-fidelity-image-understanding).
    """
	# needed for Anthropic
	media_type: SupportedImageMediaType = 'image/png'

	def __str__(self) -> str:
		url_display = _format_image_url(self.url)
		return f'🖼️  Image[{self.media_type}, detail={self.detail}]: {url_display}'

	def __repr__(self) -> str:
		url_repr = _format_image_url(self.url, 30)
		return f'ImageURL(url={repr(url_repr)}, detail={repr(self.detail)}, media_type={repr(self.media_type)})'


class ContentPartImageParam(BaseModel):
	image_url: ImageURL
	type: Literal['image_url'] = 'image_url'

	def __str__(self) -> str:
		return str(self.image_url)

	def __repr__(self) -> str:
		return f'ContentPartImageParam(image_url={repr(self.image_url)})'


class Function(BaseModel):
	arguments: str
	"""
    The arguments to call the function with, as generated by the model in JSON
    format. Note that the model does not always generate valid JSON, and may
    hallucinate parameters not defined by your function schema. Validate the
    arguments in your code before calling your function.
    """
	name: str
	"""The name of the function to call."""

	def __str__(self) -> str:
		args_preview = _truncate(self.arguments, 80)
		return f'{self.name}({args_preview})'

	def __repr__(self) -> str:
		args_repr = _truncate(repr(self.arguments), 50)
		return f'Function(name={repr(self.name)}, arguments={args_repr})'


class ToolCall(BaseModel):
	id: str
	"""The ID of the tool call."""
	function: Function
	"""The function that the model called."""
	type: Literal['function'] = 'function'
	"""The type of the tool. Currently, only `function` is supported."""

	def __str__(self) -> str:
		return f'ToolCall[{self.id}]: {self.function}'

	def __repr__(self) -> str:
		return f'ToolCall(id={repr(self.id)}, function={repr(self.function)})'


# endregion


# region - Message types
class _MessageBase(BaseModel):
	"""Base class for all message types"""

	role: Literal['user', 'system', 'assistant']

	cache: bool = False
	"""Whether to cache this message. This is only applicable when using Anthropic models.
	"""


class UserMessage(_MessageBase):
	role: Literal['user'] = 'user'
	"""The role of the messages author, in this case `user`."""

	content: str | list[ContentPartTextParam | ContentPartImageParam]
	"""The contents of the user message."""

	name: str | None = None
	"""An optional name for the participant.

    Provides the model information to differentiate between participants of the same
    role.
    """

	@property
	def text(self) -> str:
		"""
		Automatically parse the text inside content, whether it's a string or a list of content parts.
		"""
		if isinstance(self.content, str):
			return self.content
		elif isinstance(self.content, list):
			return '\n'.join([part.text for part in self.content if part.type == 'text'])
		else:
			return ''

	def __str__(self) -> str:
		return f'UserMessage(content={self.text})'

	def __repr__(self) -> str:
		return f'UserMessage(content={repr(self.text)})'


class SystemMessage(_MessageBase):
	role: Literal['system'] = 'system'
	"""The role of the messages author, in this case `system`."""

	content: str | list[ContentPartTextParam]
	"""The contents of the system message."""

	name: str | None = None

	@property
	def text(self) -> str:
		"""
		Automatically parse the text inside content, whether it's a string or a list of content parts.
		"""
		if isinstance(self.content, str):
			return self.content
		elif isinstance(self.content, list):
			return '\n'.join([part.text for part in self.content if part.type == 'text'])
		else:
			return ''

	def __str__(self) -> str:
		return f'SystemMessage(content={self.text})'

	def __repr__(self) -> str:
		return f'SystemMessage(content={repr(self.text)})'


class AssistantMessage(_MessageBase):
	role: Literal['assistant'] = 'assistant'
	"""The role of the messages author, in this case `assistant`."""

	content: str | list[ContentPartTextParam | ContentPartRefusalParam] | None
	"""The contents of the assistant message."""

	name: str | None = None

	refusal: str | None = None
	"""The refusal message by the assistant."""

	tool_calls: list[ToolCall] = []
	"""The tool calls generated by the model, such as function calls."""

	@property
	def text(self) -> str:
		"""
		Automatically parse the text inside content, whether it's a string or a list of content parts.
		"""
		if isinstance(self.content, str):
			return self.content
		elif isinstance(self.content, list):
			text = ''
			for part in self.content:
				if part.type == 'text':
					text += part.text
				elif part.type == 'refusal':
					text += f'[Refusal] {part.refusal}'
			return text
		else:
			return ''

	def __str__(self) -> str:
		return f'AssistantMessage(content={self.text})'

	def __repr__(self) -> str:
		return f'AssistantMessage(content={repr(self.text)})'


BaseMessage = Union[UserMessage, SystemMessage, AssistantMessage]

# endregion