import re
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Literal

from packaging.specifiers import InvalidSpecifier, SpecifierSet
from packaging.version import Version
from typing_extensions import assert_never, override

Environment = dict[str, list[str | Version | re.Pattern[str] | bool]]
Comparator = Literal["==", "!=", ">", "<", ">=", "<="]


class Node(ABC):
    """Base class for all nodes in the marker expression tree."""

    @abstractmethod
    def evaluate(self, environment: Environment) -> "Node":
        """Partially or fully evaluates the node based on the environment"""
        pass

    @override
    @abstractmethod
    def __str__(self) -> str:
        """Return a string representation of this node."""
        pass

    @property
    def left(self) -> "Node | None":
        return None

    @property
    def right(self) -> "Node | None":
        return None

    def contains(self, key: str) -> bool:
        if self.left is not None:
            if self.left.contains(key):
                return True
        if self.right is not None:
            if self.right.contains(key):
                return True
        return False


@dataclass(frozen=True)
class BooleanNode(Node):
    """A node representing a boolean literal value."""

    state: bool

    @override
    def __str__(self) -> str:
        return str(self.state)

    @override
    def evaluate(self, environment: Environment) -> "Node":
        return self  # No need to create new BooleanNode since they're immutable


TRUE = BooleanNode(True)
FALSE = BooleanNode(False)


@dataclass(frozen=True)
class ExpressionNode(Node):
    """A node representing a comparison expression (e.g., python_version > '3.7')."""

    lhs: str
    comparator: Comparator
    rhs: str

    @override
    def __str__(self) -> str:
        return f'{self.lhs} {self.comparator} "{self.rhs}"'

    @override
    def contains(self, key: str) -> bool:
        return self.lhs == key

    @override
    def evaluate(self, environment: Environment) -> "Node":
        if not self.lhs in environment:
            return self
        values = environment[self.lhs]
        result: bool | None = None
        for value in values:
            if isinstance(value, str):
                eval = self._evaluate_string(value)
                result = result if eval is None else result or eval
            elif isinstance(value, re.Pattern):
                eval = self._evaluate_pattern(value)
                result = result if eval is None else result or eval
            elif isinstance(value, Version):
                eval = self._evaluate_version(value)
                result = result if eval is None else result or eval
            elif isinstance(value, bool):
                result = value
                break
            else:
                assert_never(value)
        return self if result is None else BooleanNode(result)

    def _evaluate_string(self, value: str) -> "bool | None":
        if self.comparator == "==":
            return value == self.rhs
        elif self.comparator == "!=":
            return value != self.rhs
        else:
            return None

    def _evaluate_pattern(self, value: re.Pattern[str]) -> "bool | None":
        if self.comparator == "==":
            return value.match(self.rhs) is not None
        elif self.comparator == "!=":
            return not value.match(self.rhs)
        else:
            return None

    def _evaluate_version(self, value: Version) -> "bool | None":
        try:
            specifier = SpecifierSet(f"{self.comparator} {self.rhs}")
        except InvalidSpecifier:
            return None
        return specifier.contains(value)


@dataclass(frozen=True)
class OperatorNode(Node):
    """A node representing a boolean operation (and/or) between two child nodes."""

    operator: Literal["and", "or"]
    _left: Node
    _right: Node

    @property
    @override
    def left(self) -> "Node | None":
        return self._left

    @property
    @override
    def right(self) -> "Node | None":
        return self._right

    @override
    def __str__(self) -> str:
        return f"({self._left} {self.operator} {self._right})"

    @override
    def evaluate(self, environment: Environment) -> "Node":
        left = self._left.evaluate(environment)
        right = self._right.evaluate(environment)

        # If neither child changed, return self
        if left is self._left and right is self._right:
            return self

        if self.operator == "or":
            if isinstance(left, BooleanNode):
                return TRUE if left.state else right
            if isinstance(right, BooleanNode):
                return TRUE if right.state else left
            return OperatorNode(self.operator, left, right)
        elif self.operator == "and":
            if isinstance(left, BooleanNode):
                return right if left.state else FALSE
            if isinstance(right, BooleanNode):
                return left if right.state else FALSE
            return OperatorNode(self.operator, left, right)
        else:
            assert_never(self.operator)
