from dataclasses import dataclass
from typing import Any, Optional

from pydantic import BaseModel

from agent_amory_core.dom.history_tree_processor.service import DOMHistoryElement
from agent_amory_core.dom.views import DOMState


# Pydantic
class TabInfo(BaseModel):
    """Represents information about a browser tab"""

    page_id: int
    url: str
    title: str


@dataclass
class BrowserState(DOMState):
    url: str
    title: str
    tabs: list[TabInfo]
    screenshot: Optional[str] = None


@dataclass
class BrowserStateHistory:
    url: str
    title: str
    tabs: list[TabInfo]
    interacted_element: list[DOMHistoryElement | None] | list[None]
    screenshot: Optional[str] = None

    def to_dict(self) -> dict[str, Any]:
        data = {}
        data['tabs'] = [tab.model_dump() for tab in self.tabs]
        data['screenshot'] = self.screenshot
        data['interacted_element'] = [
            el.to_dict() if el else None for el in self.interacted_element
        ]
        data['url'] = self.url
        data['title'] = self.title
        return data


class BrowserError(Exception):
    """Base class for all browser errors"""
