import base64
import json
import re
from typing import Any, overload

from browser_use.llm.messages import (
	AssistantMessage,
	BaseMessage,
	ContentPartImageParam,
	ContentPartRefusalParam,
	ContentPartTextParam,
	SystemMessage,
	ToolCall,
	UserMessage,
)


class AWSBedrockMessageSerializer:
	"""Serializer for converting between custom message types and AWS Bedrock message format."""

	@staticmethod
	def _is_base64_image(url: str) -> bool:
		"""Check if the URL is a base64 encoded image."""
		return url.startswith('data:image/')

	@staticmethod
	def _is_url_image(url: str) -> bool:
		"""Check if the URL is a regular HTTP/HTTPS image URL."""
		return url.startswith(('http://', 'https://')) and any(
			url.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp']
		)

	@staticmethod
	def _parse_base64_url(url: str) -> tuple[str, bytes]:
		"""Parse a base64 data URL to extract format and raw bytes."""
		# Format: data:image/jpeg;base64,<data>
		if not url.startswith('data:'):
			raise ValueError(f'Invalid base64 URL: {url}')

		header, data = url.split(',', 1)

		# Extract format from mime type
		mime_match = re.search(r'image/(\w+)', header)
		if mime_match:
			format_name = mime_match.group(1).lower()
			# Map common formats
			format_mapping = {'jpg': 'jpeg', 'jpeg': 'jpeg', 'png': 'png', 'gif': 'gif', 'webp': 'webp'}
			image_format = format_mapping.get(format_name, 'jpeg')
		else:
			image_format = 'jpeg'  # Default format

		# Decode base64 data
		try:
			image_bytes = base64.b64decode(data)
		except Exception as e:
			raise ValueError(f'Failed to decode base64 image data: {e}')

		return image_format, image_bytes

	@staticmethod
	def _download_and_convert_image(url: str) -> tuple[str, bytes]:
		"""Download an image from URL and convert to base64 bytes."""
		try:
			import httpx
		except ImportError:
			raise ImportError('httpx not available. Please install it to use URL images with AWS Bedrock.')

		try:
			response = httpx.get(url, timeout=30)
			response.raise_for_status()

			# Detect format from content type or URL
			content_type = response.headers.get('content-type', '').lower()
			if 'jpeg' in content_type or url.lower().endswith(('.jpg', '.jpeg')):
				image_format = 'jpeg'
			elif 'png' in content_type or url.lower().endswith('.png'):
				image_format = 'png'
			elif 'gif' in content_type or url.lower().endswith('.gif'):
				image_format = 'gif'
			elif 'webp' in content_type or url.lower().endswith('.webp'):
				image_format = 'webp'
			else:
				image_format = 'jpeg'  # Default format

			return image_format, response.content

		except Exception as e:
			raise ValueError(f'Failed to download image from {url}: {e}')

	@staticmethod
	def _serialize_content_part_text(part: ContentPartTextParam) -> dict[str, Any]:
		"""Convert a text content part to AWS Bedrock format."""
		return {'text': part.text}

	@staticmethod
	def _serialize_content_part_image(part: ContentPartImageParam) -> dict[str, Any]:
		"""Convert an image content part to AWS Bedrock format."""
		url = part.image_url.url

		if AWSBedrockMessageSerializer._is_base64_image(url):
			# Handle base64 encoded images
			image_format, image_bytes = AWSBedrockMessageSerializer._parse_base64_url(url)
		elif AWSBedrockMessageSerializer._is_url_image(url):
			# Download and convert URL images
			image_format, image_bytes = AWSBedrockMessageSerializer._download_and_convert_image(url)
		else:
			raise ValueError(f'Unsupported image URL format: {url}')

		return {
			'image': {
				'format': image_format,
				'source': {
					'bytes': image_bytes,
				},
			}
		}

	@staticmethod
	def _serialize_user_content(
		content: str | list[ContentPartTextParam | ContentPartImageParam],
	) -> list[dict[str, Any]]:
		"""Serialize content for user messages."""
		if isinstance(content, str):
			return [{'text': content}]

		content_blocks: list[dict[str, Any]] = []
		for part in content:
			if part.type == 'text':
				content_blocks.append(AWSBedrockMessageSerializer._serialize_content_part_text(part))
			elif part.type == 'image_url':
				content_blocks.append(AWSBedrockMessageSerializer._serialize_content_part_image(part))

		return content_blocks

	@staticmethod
	def _serialize_system_content(
		content: str | list[ContentPartTextParam],
	) -> list[dict[str, Any]]:
		"""Serialize content for system messages."""
		if isinstance(content, str):
			return [{'text': content}]

		content_blocks: list[dict[str, Any]] = []
		for part in content:
			if part.type == 'text':
				content_blocks.append(AWSBedrockMessageSerializer._serialize_content_part_text(part))

		return content_blocks

	@staticmethod
	def _serialize_assistant_content(
		content: str | list[ContentPartTextParam | ContentPartRefusalParam] | None,
	) -> list[dict[str, Any]]:
		"""Serialize content for assistant messages."""
		if content is None:
			return []
		if isinstance(content, str):
			return [{'text': content}]

		content_blocks: list[dict[str, Any]] = []
		for part in content:
			if part.type == 'text':
				content_blocks.append(AWSBedrockMessageSerializer._serialize_content_part_text(part))
			# Skip refusal content parts - AWS Bedrock doesn't need them

		return content_blocks

	@staticmethod
	def _serialize_tool_call(tool_call: ToolCall) -> dict[str, Any]:
		"""Convert a tool call to AWS Bedrock format."""
		try:
			arguments = json.loads(tool_call.function.arguments)
		except json.JSONDecodeError:
			# If arguments aren't valid JSON, wrap them
			arguments = {'arguments': tool_call.function.arguments}

		return {
			'toolUse': {
				'toolUseId': tool_call.id,
				'name': tool_call.function.name,
				'input': arguments,
			}
		}

	# region - Serialize overloads
	@overload
	@staticmethod
	def serialize(message: UserMessage) -> dict[str, Any]: ...

	@overload
	@staticmethod
	def serialize(message: SystemMessage) -> SystemMessage: ...

	@overload
	@staticmethod
	def serialize(message: AssistantMessage) -> dict[str, Any]: ...

	@staticmethod
	def serialize(message: BaseMessage) -> dict[str, Any] | SystemMessage:
		"""Serialize a custom message to AWS Bedrock format."""

		if isinstance(message, UserMessage):
			return {
				'role': 'user',
				'content': AWSBedrockMessageSerializer._serialize_user_content(message.content),
			}

		elif isinstance(message, SystemMessage):
			# System messages are handled separately in AWS Bedrock
			return message

		elif isinstance(message, AssistantMessage):
			content_blocks: list[dict[str, Any]] = []

			# Add content blocks if present
			if message.content is not None:
				content_blocks.extend(AWSBedrockMessageSerializer._serialize_assistant_content(message.content))

			# Add tool use blocks if present
			if message.tool_calls:
				for tool_call in message.tool_calls:
					content_blocks.append(AWSBedrockMessageSerializer._serialize_tool_call(tool_call))

			# AWS Bedrock requires at least one content block
			if not content_blocks:
				content_blocks = [{'text': ''}]

			return {
				'role': 'assistant',
				'content': content_blocks,
			}

		else:
			raise ValueError(f'Unknown message type: {type(message)}')

	@staticmethod
	def serialize_messages(messages: list[BaseMessage]) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None]:
		"""
		Serialize a list of messages, extracting any system message.

		Returns:
			Tuple of (bedrock_messages, system_message) where system_message is extracted
			from any SystemMessage in the list.
		"""
		bedrock_messages: list[dict[str, Any]] = []
		system_message: list[dict[str, Any]] | None = None

		for message in messages:
			if isinstance(message, SystemMessage):
				# Extract system message content
				system_message = AWSBedrockMessageSerializer._serialize_system_content(message.content)
			else:
				# Serialize and add to regular messages
				serialized = AWSBedrockMessageSerializer.serialize(message)
				bedrock_messages.append(serialized)

		return bedrock_messages, system_message