from __future__ import annotations

from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, ConfigDict, Field

from ...adapters.llm_adapter 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)

    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}
</{step_str}>"""
        elif self.system_message:
            return ""  # empty string
        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}
</{step_str}>"""


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 = ""
    read_state_images: list[dict[str, Any]] = Field(default_factory=list)
    # Compaction state
    compacted_memory: str = ""
    compaction_count: int = 0
    last_compaction_step: int | None = None

    model_config = ConfigDict(arbitrary_types_allowed=True)
