import datetime as dt
from base64 import b64encode
from collections.abc import Sequence
from dataclasses import field

from loguru import logger
from PIL import Image
from pydantic import BaseModel, Field

from notte_core.actions import InteractionAction
from notte_core.browser.dom_tree import A11yTree, DomNode, InteractionDomNode
from notte_core.utils.url import clean_url


class TabsData(BaseModel):
    tab_id: int
    title: str
    url: str


class ViewportData(BaseModel):
    scroll_x: int
    scroll_y: int
    viewport_width: int
    viewport_height: int
    total_width: int
    total_height: int

    @property
    def pixels_above(self) -> int:
        return self.scroll_y

    @property
    def pixels_below(self) -> int:
        return self.total_height - self.scroll_y - self.viewport_height


class SnapshotMetadata(BaseModel):
    title: str
    url: str
    viewport: ViewportData
    tabs: list[TabsData]
    timestamp: dt.datetime = field(default_factory=lambda: dt.datetime.now())


class BrowserSnapshot(BaseModel):
    metadata: SnapshotMetadata
    html_content: str
    a11y_tree: A11yTree | None
    dom_node: DomNode
    screenshot: bytes = Field(repr=False)

    model_config = {  # type: ignore[reportUnknownMemberType]
        "json_encoders": {
            bytes: lambda v: b64encode(v).decode("utf-8") if v else None,
        }
    }

    def display_screenshot(self) -> "Image.Image | None":
        from notte_core.utils.image import image_from_bytes

        return image_from_bytes(self.screenshot)

    @property
    def clean_url(self) -> str:
        return clean_url(self.metadata.url)

    def compare_with(self, other: "BrowserSnapshot") -> bool:
        inodes = {node.id for node in self.dom_node.interaction_nodes()}
        new_inodes = {node.id for node in other.dom_node.interaction_nodes()}
        identical = inodes == new_inodes
        if not identical:
            logger.trace(f"Interactive nodes changed: {new_inodes.difference(inodes)}")
        return identical

    def interaction_nodes(self) -> Sequence[InteractionDomNode]:
        return self.dom_node.interaction_nodes()

    def with_dom_node(self, dom_node: DomNode) -> "BrowserSnapshot":
        return BrowserSnapshot(
            metadata=self.metadata,
            html_content=self.html_content,
            a11y_tree=self.a11y_tree,
            dom_node=dom_node,
            screenshot=self.screenshot,
        )

    def subgraph_without(
        self, actions: Sequence[InteractionAction], roles: set[str] | None = None
    ) -> "BrowserSnapshot | None":
        if len(actions) == 0 and roles is not None:
            subgraph = self.dom_node.subtree_without(roles)
            return self.with_dom_node(subgraph)
        id_existing_actions = set([action.id for action in actions])
        failed_actions = {node.id for node in self.interaction_nodes() if node.id not in id_existing_actions}

        def only_failed_actions(node: DomNode) -> bool:
            return len(set(node.subtree_ids).intersection(failed_actions)) > 0

        filtered_graph = self.dom_node.subtree_filter(only_failed_actions)
        if filtered_graph is None:
            return None

        return self.with_dom_node(filtered_graph)
