from __future__ import annotations

from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, ConfigDict, Field

from browser_use.llm.messages import (
	BaseMessage,
)

if TYPE_CHECKING:
	pass


class HistoryItem(BaseModel):
	"""Represents a single agent history item with its data and string representation"""

	step_number: int | None = None
	evaluation_previous_goal: str | None = None
	memory: str | None = None
	next_goal: str | None = None
	action_results: str | None = None
	error: str | None = None
	system_message: str | None = None

	model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

	def model_post_init(self, __context) -> None:
		"""Validate that error and system_message are not both provided"""
		if self.error is not None and self.system_message is not None:
			raise ValueError('Cannot have both error and system_message at the same time')

	def to_string(self) -> str:
		"""Get string representation of the history item"""
		step_str = 'step' if self.step_number is not None else 'step_unknown'

		if self.error:
			return f"""<{step_str}>
{self.error}"""
		elif self.system_message:
			return self.system_message
		else:
			content_parts = []

			# Only include evaluation_previous_goal if it's not None/empty
			if self.evaluation_previous_goal:
				content_parts.append(f'{self.evaluation_previous_goal}')

			# Always include memory
			if self.memory:
				content_parts.append(f'{self.memory}')

			# Only include next_goal if it's not None/empty
			if self.next_goal:
				content_parts.append(f'{self.next_goal}')

			if self.action_results:
				content_parts.append(self.action_results)

			content = '\n'.join(content_parts)

			return f"""<{step_str}>
{content}"""


class MessageHistory(BaseModel):
	"""History of messages"""

	system_message: BaseMessage | None = None
	state_message: BaseMessage | None = None
	context_messages: list[BaseMessage] = Field(default_factory=list)
	model_config = ConfigDict(arbitrary_types_allowed=True)

	def get_messages(self) -> list[BaseMessage]:
		"""Get all messages in the correct order: system -> state -> contextual"""
		messages = []
		if self.system_message:
			messages.append(self.system_message)
		if self.state_message:
			messages.append(self.state_message)
		messages.extend(self.context_messages)

		return messages


class MessageManagerState(BaseModel):
	"""Holds the state for MessageManager"""

	history: MessageHistory = Field(default_factory=MessageHistory)
	tool_id: int = 1
	agent_history_items: list[HistoryItem] = Field(
		default_factory=lambda: [HistoryItem(step_number=0, system_message='Agent initialized')]
	)
	read_state_description: str = ''
	# Images to include in the next state message (cleared after each step)
	read_state_images: list[dict[str, Any]] = Field(default_factory=list)
	compacted_memory: str | None = None
	compaction_count: int = 0
	last_compaction_step: int | None = None

	model_config = ConfigDict(arbitrary_types_allowed=True)