from dataclasses import dataclass, field
from typing import List, Dict, Union
import pytest

# ==============================================================================
# 1. Core Data Structures & Logic
# ==============================================================================


@dataclass
class Element:
    """Represents an element in the HTML stack."""

    tag_name: str
    attributes: Dict[str, str] = field(default_factory=dict)
    child_index: int = 0

    def __repr__(self) -> str:
        attrs = " ".join([f'{k}="{v}"' for k, v in self.attributes.items()])
        idx_str = f" (child:{self.child_index})"
        return f"<{self.tag_name}{idx_str} {attrs}>".replace(" >", ">")


class SimpleSelector:
    """
    Represents a simple CSS selector component (e.g., 'div.main[data-id]:nth-child(2)').
    """

    def __init__(
        self,
        tag_name: str = None,
        class_list: List[str] = None,
        attribute_presence: List[str] = None,
        attribute_values: Dict[str, str] = None,
        nth_child: int = None,
    ):
        self.tag_name = tag_name
        self.class_list = class_list or []
        self.attribute_presence = attribute_presence or []
        self.attribute_values = attribute_values or {}
        self.nth_child = nth_child

    def __repr__(self) -> str:
        parts = []
        if self.tag_name:
            parts.append(self.tag_name)
        if self.class_list:
            parts.extend([f".{c}" for c in self.class_list])
        if self.attribute_presence:
            parts.extend([f"[{a}]" for a in self.attribute_presence])
        if self.attribute_values:
            parts.extend(
                [f'[{k}="{v}"]' for k, v in self.attribute_values.items()])
        if self.nth_child is not None:
            parts.append(f":nth-child({self.nth_child})")
        return "".join(parts) or "*"

    def matches(self, element: Element) -> bool:
        """Checks if this simple selector matches a given element."""
        # Tag name check
        if self.tag_name and self.tag_name != element.tag_name:
            return False
        # Class list check
        if self.class_list:
            element_classes_str = element.attributes.get("class")
            if not element_classes_str:
                return False
            element_classes = set(element_classes_str.split())
            if not set(self.class_list).issubset(element_classes):
                return False
        # Attribute presence check (e.g., [href])
        for attr_key in self.attribute_presence:
            if attr_key not in element.attributes:
                return False
        # Attribute value check (e.g., [lang="en"])
        for attr_key, expected_value in self.attribute_values.items():
            if element.attributes.get(attr_key) != expected_value:
                return False
        # :nth-child check
        if self.nth_child is not None:
            if (self.nth_child - 1) != element.child_index:
                return False
        return True


class Combinator:
    """Abstract base class for combinators like ' ' (descendant) or '>' (child)."""

    def find_ancestor(
        self, simple_selector: SimpleSelector, stack: List[Element], start_idx: int
    ) -> int:
        raise NotImplementedError


class DescendantCombinator(Combinator):
    """Handles the descendant combinator (' ')."""

    def __repr__(self):
        return "' '"

    def find_ancestor(
        self, simple_selector: SimpleSelector, stack: List[Element], start_idx: int
    ) -> int:
        for i in range(start_idx, -1, -1):
            if simple_selector.matches(stack[i]):
                return i
        return -1


class ChildCombinator(Combinator):
    """Handles the direct child combinator ('>')."""

    def __repr__(self):
        return "'>'"

    def find_ancestor(
        self, simple_selector: SimpleSelector, stack: List[Element], start_idx: int
    ) -> int:
        if start_idx >= 0 and simple_selector.matches(stack[start_idx]):
            return start_idx
        return -1


def matches(
    stack: List[Element], selectors: list[Union[SimpleSelector, Combinator]]
) -> bool:
    """Checks if a full selector (AST) matches the current element stack."""
    if not selectors or not stack:
        return False

    stack_idx = len(stack) - 1
    selector_idx = len(selectors) - 1

    # The last selector must match the last element on the stack.
    last_selector = selectors[selector_idx]
    if not isinstance(last_selector, SimpleSelector) or not last_selector.matches(
        stack[stack_idx]
    ):
        return False

    selector_idx -= 1

    # Walk backwards through selectors and the stack.
    while selector_idx > 0:
        combinator = selectors[selector_idx]
        simple_selector = selectors[selector_idx - 1]
        ancestor_search_start_idx = stack_idx - 1

        match_idx = combinator.find_ancestor(
            simple_selector, stack, ancestor_search_start_idx
        )
        if match_idx == -1:
            return False

        stack_idx = match_idx
        selector_idx -= 2

    return True


# ==============================================================================
# 2. Pytest Tests
# ==============================================================================

# --- Test Fixtures ---


@pytest.fixture
def attribute_selector_stack():
    """
    Represents the parser state for attribute selector tests:
    <body>
        <div class="main" data-id="123">
            <p class="content" lang="en"></p>
        </div>
    </body>
    """
    return [
        Element("body", child_index=0),
        Element("div", attributes={"class": "main",
                "data-id": "123"}, child_index=1),
        Element("p", attributes={"class": "content",
                "lang": "en"}, child_index=1),
    ]


@pytest.fixture
def nth_child_selector_stack():
    """
    Represents the parser state for :nth-child tests:
    <body>
        <h1>...</h1>
        <div class="main">
            <p>...</p>
            <p class="content">...</p>
        </div>
    </body>
    """
    return [
        Element("body", child_index=1),
        Element("div", attributes={"class": "main"}, child_index=2),
        Element("p", attributes={"class": "content"},
                child_index=2),  # Current element
    ]


# --- SimpleSelector Unit Tests ---


@pytest.mark.parametrize(
    "selector, expected",
    [
        (SimpleSelector(attribute_presence=["lang"]), True),
        (SimpleSelector(attribute_presence=["lang", "class"]), True),
        (SimpleSelector(attribute_presence=["id"]), False),
        (SimpleSelector(attribute_values={"lang": "en"}), True),
        (SimpleSelector(attribute_values={"lang": "fr"}), False),
        (SimpleSelector(attribute_values={"class": "content"}), True),
        (
            SimpleSelector(
                "p", attribute_presence=["lang"], attribute_values={"class": "content"}
            ),
            True,
        ),
        (SimpleSelector(".content", attribute_values={"lang": "fr"}), False),
    ],
    ids=[
        "[lang]",
        "[lang][class]",
        "[id]",
        "[lang=en]",
        "[lang=fr]",
        "[class=content]",
        "p[lang][class=content]",
        ".content[lang=fr]",
    ],
)
def test_simple_selector_attributes(attribute_selector_stack, selector, expected):
    """Tests attribute selectors in SimpleSelector.matches()."""
    current_element = attribute_selector_stack[-1]
    assert selector.matches(current_element) == expected


@pytest.mark.parametrize(
    "selector, expected",
    [
        (
            SimpleSelector("p", nth_child=2),
            False,
        ),  # child_index is 0-based, nth-child is 1-based
        (SimpleSelector("p", nth_child=3), True),
        (SimpleSelector(class_list=["content"], nth_child=3), True),
        (SimpleSelector(nth_child=3), True),
        (SimpleSelector(nth_child=1), False),
    ],
    ids=[
        "p:nth-child(2)",
        "p:nth-child(3)",
        ".content:nth-child(3)",
        ":nth-child(3)",
        ":nth-child(1)",
    ],
)
def test_simple_selector_nth_child(nth_child_selector_stack, selector, expected):
    """Tests the :nth-child logic in SimpleSelector.matches()."""
    current_element = nth_child_selector_stack[-1]
    assert selector.matches(current_element) == expected


# --- Integration Tests for the main matches() function ---

# Define common selector components for reuse
BODY = SimpleSelector("body")
DESCENDANT = DescendantCombinator()
CHILD = ChildCombinator()


@pytest.mark.parametrize(
    "test_id, selector_ast, expected",
    [
        (
            "match_p_with_attr_val",
            [
                SimpleSelector(
                    "p", class_list=["content"], attribute_values={"lang": "en"}
                )
            ],
            True,
        ),
        (
            "fail_p_with_wrong_attr_val",
            [SimpleSelector("p", attribute_values={"lang": "fr"})],
            False,
        ),
        (
            "match_div_with_attr_presence_and_descendant_p",
            [
                SimpleSelector("div", attribute_presence=["data-id"]),
                DESCENDANT,
                SimpleSelector("p", attribute_values={"lang": "en"}),
            ],
            True,
        ),
        (
            "match_div_and_child_p_with_attrs",
            [
                SimpleSelector("div", attribute_values={"data-id": "123"}),
                CHILD,
                SimpleSelector("p", attribute_values={"lang": "en"}),
            ],
            True,
        ),
        (
            "fail_div_with_wrong_attr_val",
            [
                SimpleSelector("div", attribute_values={"data-id": "999"}),
                CHILD,
                SimpleSelector("p", attribute_values={"lang": "en"}),
            ],
            False,
        ),
    ],
)
def test_main_matches_with_attributes(
    test_id, selector_ast, expected, attribute_selector_stack
):
    """Tests the main matches() function with attribute selectors."""
    assert matches(attribute_selector_stack, selector_ast) == expected


@pytest.mark.parametrize(
    "test_id, selector_ast, expected",
    [
        ("match_p_as_3rd_child", [SimpleSelector("p", nth_child=3)], True),
        ("fail_p_as_1st_child", [SimpleSelector("p", nth_child=1)], False),
        (
            "match_p_descendant_of_div_as_3rd_child",
            [SimpleSelector("div"), DESCENDANT,
             SimpleSelector("p", nth_child=3)],
            True,
        ),
        (
            "match_p_child_of_div_as_3rd_child",
            [SimpleSelector("div"), CHILD, SimpleSelector("p", nth_child=3)],
            True,
        ),
        (
            "match_div_as_3rd_child_and_its_p_child",
            [
                SimpleSelector("div", nth_child=3),
                CHILD,
                SimpleSelector("p", class_list=["content"]),
            ],
            True,
        ),
        (
            "fail_div_as_1st_child",
            [SimpleSelector("div", nth_child=1), CHILD, SimpleSelector("p")],
            False,
        ),
        (
            "fail_p_is_not_child_of_body",
            [BODY, CHILD, SimpleSelector("p", nth_child=3)],
            False,
        ),
        (
            "match_body_has_descendant_p_3rd_child",
            [BODY, DESCENDANT, SimpleSelector("p", nth_child=3)],
            True,
        ),
    ],
)
def test_main_matches_with_nth_child(
    test_id, selector_ast, expected, nth_child_selector_stack
):
    """Tests the main matches() function with selectors including :nth-child."""
    assert matches(nth_child_selector_stack, selector_ast) == expected
